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

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

  τ(E) — Layer 2: node-type frequency histogram over Δ_{k-1} (first 8 map to bases 8-15)
  δ(E) — Layer 3: child-ordering frequency histogram over Δ_{2k²-1} (first 8 map to bases 16-23)

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

    Examples:
        >>> preprocess_math("x = 5")
        'x == 5'
        >>> preprocess_math("|x|")
        'abs(x)'
        >>> preprocess_math("p_(i+j)")
        'p[i+j]'
        >>> preprocess_math("(x)(y)")
        '(x)*(y)'
        >>> preprocess_math("$x$")
        'x'
    """
    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'\)\s*\(', 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.

    Examples:
        >>> parse_to_ast("x + 1") is not None
        True
        >>> parse_to_ast("") is None
        True
        >>> parse_to_ast("invalid syntax @@@") is None
        True
    """
    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 a Python AST node to one of the 18 NODE_TYPES strings.

    Args:
        node: A Python AST node to classify.

    Returns:
        One of the following 18 type tags:

        - ``bin_add``, ``bin_sub``, ``bin_mul``, ``bin_div``, ``bin_pow``
          for :class:`ast.BinOp` nodes (unknown binary operators fall back to
          ``bin_add``)
        - ``un_neg``, ``un_not`` for :class:`ast.UnaryOp` nodes (unknown
          unary operators fall back to ``un_neg``)
        - ``var`` for :class:`ast.Name` nodes
        - ``const`` for :class:`ast.Constant` nodes
        - ``call`` for :class:`ast.Call` nodes
        - ``subscript`` for :class:`ast.Subscript` nodes
        - ``tuple`` for :class:`ast.Tuple` nodes

        Any unrecognised node type falls back to ``const``.
    """
    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) — weighted node-type histogram.

    Returns 18 floats (one per NODE_TYPE) summing to 1.0,
    weighted by the depth of each node in the parse tree. The first 8 values
    map directly to Layer 2 of the Φ DNA embedding (bases 8-15).

    Examples:
        >>> result = compute_tau("x + 1")
        >>> len(result)
        18
        >>> round(sum(result), 10)
        1.0
    """
    tree = parse_to_ast(equation)
    if tree is None:
        return None

    counts = {t: 0.0 for t in NODE_TYPES}
    total_weight = 0.0

    def walk(node: ast.AST, depth: int):
        nonlocal total_weight
        t = _classify_node(node)
        # Weight is proportional to (depth + 1)
        counts[t] += (depth + 1)
        total_weight += (depth + 1)
        for child in ast.iter_child_nodes(node):
            walk(child, depth + 1)

    walk(tree, 0)

    if total_weight == 0:
        return [1.0 / len(NODE_TYPES)] * len(NODE_TYPES)
    return [counts[t] / total_weight 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 = 8


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

    For each (parent_type, child_type, child_index) triple, returns
    the fraction of all parent-child edges, weighted by the depth 
    of the parent node. The first 8 values map directly to Layer 3
    of the Φ DNA embedding (bases 16-23).

    Examples:
        >>> result = compute_delta("x + 1")
        >>> len(result)
        648
        >>> all(v >= 0.0 for v in result)
        True
    """
    tree = parse_to_ast(equation)
    if tree is None:
        return None

    edges: Dict[Tuple[str, str, int], float] = {}
    total_weight = 0.0

    def walk(node: ast.AST, depth: int):
        nonlocal total_weight
        parent_t = _classify_node(node)
        # Weight is proportional to (depth + 1)
        weight = (depth + 1)
        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.0) + weight
            total_weight += weight
            walk(child, depth + 1)

    walk(tree, 0)

    if total_weight == 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.0) / total_weight)

    return result


# ── Depth coefficient λ(E) and recurrence vector R(E) ─────────────────────


def compute_lambda_and_r(equation: str) -> Tuple[Optional[float], Optional[List[float]]]:
    """Compute depth coefficient λ(E) and recurrence vector R(E).

    λ(E) — the maximum parse-tree depth (normalised by log₂(n)+1),
    a measure of structural nesting.  Returns None for unparseable input.

    R(E) — a normalised histogram of node-type recurrence: how often
    each node type appears in the AST, as a fraction of total nodes.

    Args:
        equation: The equation string.

    Returns:
        A tuple ``(lambda_val, r_vec)`` where:
        - ``lambda_val`` is a float in (0, 1] representing max depth
          (or None if unparseable).
        - ``r_vec`` is a list of 18 floats summing to 1.0 (or None).

    Examples:
        >>> lam, r = compute_lambda_and_r("x + 1")
        >>> lam is not None
        True
        >>> r is not None
        True
        >>> round(sum(r), 10)
        1.0
        >>> lam, r = compute_lambda_and_r("@#$%")
        >>> lam is None
        True
        >>> r is None
        True
    """
    tree = parse_to_ast(equation)
    if tree is None:
        return None, None

    max_depth = 0
    counts = {nt: 0 for nt in NODE_TYPES}
    total_nodes = 0

    def walk(node: ast.AST, depth: int) -> None:
        nonlocal max_depth, total_nodes
        max_depth = max(max_depth, depth)
        ntype = _classify_node(node)
        counts[ntype] = counts.get(ntype, 0) + 1
        total_nodes += 1
        for child in ast.iter_child_nodes(node):
            walk(child, depth + 1)

    walk(tree, 0)

    # λ: normalised depth; log₂(n)+1 is the depth of a balanced tree
    lambda_val = 1.0 - 1.0 / (max_depth + 2.0) if max_depth > 0 else 0.5

    if total_nodes == 0:
        return lambda_val, [1.0 / len(NODE_TYPES)] * len(NODE_TYPES)

    r_vec = [counts[nt] / total_nodes for nt in NODE_TYPES]
    return lambda_val, r_vec


if __name__ == "__main__":
    # ── 1.  preprocess_math conversion rules ───────────────────────────────
    assert preprocess_math("x = 5") == "x == 5"
    assert preprocess_math("<= x") == "<= x"         # <= preserved
    assert preprocess_math("|x|") == "abs(x)"
    assert preprocess_math("p_(i+j)") == "p[i+j]"
    assert preprocess_math("(x)(y)") == "(x)*(y)"
    assert preprocess_math("$x$") == "x"             # TeX markers removed

    # ── 2.  parse_to_ast valid / invalid ───────────────────────────────────
    assert parse_to_ast("x + 1") is not None
    assert parse_to_ast("") is None
    assert parse_to_ast("invalid syntax @@@") is None

    # ── 3.  compute_tau: len=18, sums to 1.0 ──────────────────────────────
    tau = compute_tau("x + 1")
    assert tau is not None
    assert len(tau) == 18
    assert abs(sum(tau) - 1.0) < 1e-10

    # also works for other parseable equations
    tau2 = compute_tau("E = m * c**2")
    assert tau2 is not None
    assert len(tau2) == 18
    assert abs(sum(tau2) - 1.0) < 1e-10

    # ── 4.  compute_delta: len=648, sum captures tracked edges ─────────────
    # NOTE: sum is < 1.0 because the Expression wrapper node (parent type
    # "const") and Name→Load edges fall outside the DELTA_* type lists.
    delta = compute_delta("x + 1")
    assert delta is not None
    assert len(delta) == 648, f"expected 648, got {len(delta)}"
    assert all(v >= 0.0 for v in delta)
    assert 0 < sum(delta) <= 1.0

    delta2 = compute_delta("a + b * c")
    assert delta2 is not None
    assert len(delta2) == 648
    assert all(v >= 0.0 for v in delta2)
    assert 0 < sum(delta2) <= 1.0

    # ── 5.  Edge case: empty string → None ────────────────────────────────
    assert parse_to_ast("") is None
    assert compute_tau("") is None
    assert compute_delta("") is None

    print("ast_parse.py: all verifications PASSED")