"""
phi.ast_parse — Layer 3: Parse tree analysis → τ and δ distributions

Converts equation strings to Python ASTs (after preprocessing math
notation), then computes:

  τ(E) — node-type frequency histogram over Δ_{k-1}
  δ(E) — child-ordering frequency histogram over Δ_{2k²-1}

Dependencies: Python ast, re (stdlib). Independent of phi.charclass.
"""

from __future__ import annotations

import ast
import re
from typing import Dict, List, Optional, Tuple

# ── Node types recognized by the parser ──────────────────────────────────

NODE_TYPES = [
    "bin_add", "bin_sub", "bin_mul", "bin_div",
    "bin_pow", "bin_eq", "un_neg", "un_not",
    "var", "const", "call", "subscript", "tuple",
    "list", "dict", "set", "lambda", "if_exp",
]


# ── Math notation preprocessor ───────────────────────────────────────────

def preprocess_math(text: str) -> str:
    """Preprocess math notation into Python-parseable form.

    Handles:
      - Single = → == (but preserves <=, >=, !=, ==)
      - |x| → abs(x)
      - Parenthesized subscripts: p_(i+j) → p[i+j]
      - Implicit multiplication: (x)(y) → (x)*(y)
      - TeX markers: removed
    """
    s = text.strip()
    if s.endswith(":"):
        return ""

    s = re.sub(r'\$|\$\$|\\\[|\\\]|\\\(|\\\)', '', s)
    s = re.sub(r'(?<![<>=!])=(?!=)', '==', s)
    s = re.sub(r'\|([^|=]+)\|', r'abs(\1)', s)
    s = re.sub(r'([a-zA-Z])_\(([^)]+)\)', r'\1[\2]', s)
    s = re.sub(r'\)\(', r')*(', s)
    return s


def parse_to_ast(equation: str) -> Optional[ast.AST]:
    """Parse an equation string into a Python AST.

    Tries direct parse first, then wraps in parens for expressions
    where math notation drops outer parentheses.
    """
    text = preprocess_math(equation)
    if not text:
        return None
    try:
        return ast.parse(text, mode="eval")
    except SyntaxError:
        pass
    try:
        return ast.parse(f"({text})", mode="eval")
    except SyntaxError:
        return None


# ── AST node classification ──────────────────────────────────────────────

def _classify_node(node: ast.AST) -> str:
    """Map an AST node to one of the NODE_TYPES."""
    if isinstance(node, ast.BinOp):
        if isinstance(node.op, ast.Add):
            return "bin_add"
        if isinstance(node.op, ast.Sub):
            return "bin_sub"
        if isinstance(node.op, ast.Mult):
            return "bin_mul"
        if isinstance(node.op, ast.Div):
            return "bin_div"
        if isinstance(node.op, ast.Pow):
            return "bin_pow"
        return "bin_add"
    if isinstance(node, ast.UnaryOp):
        if isinstance(node.op, ast.USub):
            return "un_neg"
        if isinstance(node.op, ast.Not):
            return "un_not"
        return "un_neg"
    if isinstance(node, ast.Name):
        return "var"
    if isinstance(node, ast.Constant):
        return "const"
    if isinstance(node, ast.Call):
        return "call"
    if isinstance(node, ast.Subscript):
        return "subscript"
    if isinstance(node, ast.Tuple):
        return "tuple"
    return "const"


# ── τ(E): node-type frequencies → Δ_{k-1} ───────────────────────────────

def compute_tau(equation: str) -> Optional[List[float]]:
    """Compute τ(E) — normalized node-type histogram.

    Returns 18 floats (one per NODE_TYPE) summing to 1.0,
    or None if the equation cannot be parsed.
    """
    tree = parse_to_ast(equation)
    if tree is None:
        return None

    counts = {t: 0 for t in NODE_TYPES}
    for node in ast.walk(tree):
        t = _classify_node(node)
        counts[t] = counts.get(t, 0) + 1

    total = sum(counts.values())
    if total == 0:
        return [1.0 / len(NODE_TYPES)] * len(NODE_TYPES)
    return [counts[t] / total for t in NODE_TYPES]


# ── δ(E): child-ordering frequencies → Δ_{2k²-1} ────────────────────────

# Which (parent, child) type combinations are tracked in δ
DELTA_PARENT_TYPES = [
    "bin_add", "bin_sub", "bin_mul", "bin_div", "bin_pow",
    "un_neg", "un_not", "call", "subscript",
]
DELTA_CHILD_TYPES = [
    "bin_add", "bin_sub", "bin_mul", "bin_div", "bin_pow",
    "un_neg", "var", "const", "call",
]
MAX_CHILDREN = 3


def compute_delta(equation: str) -> Optional[List[float]]:
    """Compute δ(E) — child-ordering frequency histogram.

    For each (parent_type, child_type, child_index) triple, returns
    the fraction of all parent-child edges.  Captures syntactic
    topology beyond raw node counts.

    Returns a flat list of length len(DELTA_PARENT_TYPES) ×
    len(DELTA_CHILD_TYPES) × MAX_CHILDREN, or None if parse fails.
    """
    tree = parse_to_ast(equation)
    if tree is None:
        return None

    edges: Dict[Tuple[str, str, int], int] = {}
    total_edges = 0

    def walk(node: ast.AST, _child_idx: int = 0):
        nonlocal total_edges
        parent_t = _classify_node(node)
        for i, child in enumerate(ast.iter_child_nodes(node)):
            child_t = _classify_node(child)
            key = (parent_t, child_t, i)
            edges[key] = edges.get(key, 0) + 1
            total_edges += 1
            walk(child, i)

    walk(tree)

    if total_edges == 0:
        return None

    result = []
    for pt in DELTA_PARENT_TYPES:
        for ct in DELTA_CHILD_TYPES:
            for i in range(MAX_CHILDREN):
                result.append(edges.get((pt, ct, i), 0) / total_edges)

    return result