#!/usr/bin/env python3
"""
expr_tree.py — Expression Tree Data Structure for SilverSight Symbolic Regression

Expression trees are the core data structure for representing mathematical
expressions. Each node is either:
  - A leaf: variable 'x' or a constant float
  - An internal node: unary operator (sin, cos, log, exp, sqrt, neg)
                   or binary operator (+, -, *, /, pow)

Operations:
  - evaluate(xs): compute expression value for input array
  - to_string(): human-readable expression
  - tree_size(): node count
  - tree_depth(): max depth
  - copy(): deep copy
"""

from __future__ import annotations

import math
import numpy as np
from dataclasses import dataclass, field
from typing import Optional, List


# ── Operator definitions ────────────────────────────────────────────────

UNARY_OPS = {
    "sin":  (np.sin,  "sin"),
    "cos":  (np.cos,  "cos"),
    "log":  (np.log,  "log"),    # safe log: |x|
    "exp":  (np.exp,  "exp"),
    "sqrt": (np.sqrt, "sqrt"),   # safe sqrt: |x|
    "neg":  (np.negative, "neg"),
    "abs":  (np.abs,  "abs"),
}

BINARY_OPS = {
    "+":   (np.add,      "+"),
    "-":   (np.subtract, "-"),
    "*":   (np.multiply, "*"),
    "/":   (np.divide,   "/"),    # safe div: x/y with floor
    "pow": (np.power,    "pow"),  # safe pow: |x|^y
}

# Cost weights for BIC complexity
OP_COST = {
    "x": 1, "const": 3,
    "sin": 2, "cos": 2, "log": 2, "exp": 2, "sqrt": 2, "neg": 1, "abs": 1,
    "+": 1, "-": 1, "*": 1, "/": 2, "pow": 3,
}


# ── Expression Node ─────────────────────────────────────────────────────

@dataclass
class ExprNode:
    """A node in an expression tree."""
    op: str                                   # operator name
    left: Optional['ExprNode'] = None         # left child (or unary child)
    right: Optional['ExprNode'] = None        # right child (None for unary/leaf)
    value: Optional[float] = None             # constant value (for const leaves)
    depth: int = 0                            # max depth from this node
    size: int = 0                             # total node count from this node

    def evaluate(self, xs: np.ndarray) -> np.ndarray:
        """Evaluate expression on input array."""
        if self.op == "x":
            return xs.copy()
        elif self.op == "const":
            return np.full_like(xs, self.value, dtype=np.float64)
        elif self.op in UNARY_OPS:
            child_val = self.left.evaluate(xs)
            if self.op == "log":
                return np.log(np.abs(child_val) + 1e-15)
            elif self.op == "sqrt":
                return np.sqrt(np.abs(child_val))
            elif self.op == "neg":
                return -child_val
            elif self.op == "abs":
                return np.abs(child_val)
            else:
                return UNARY_OPS[self.op][0](child_val)
        elif self.op in BINARY_OPS:
            left_val = self.left.evaluate(xs)
            right_val = self.right.evaluate(xs)
            if self.op == "/":
                return left_val / np.where(np.abs(right_val) < 1e-15, 1.0, right_val)
            elif self.op == "pow":
                return np.power(np.abs(left_val), np.clip(right_val, -10, 10))
            else:
                return BINARY_OPS[self.op][0](left_val, right_val)
        else:
            raise ValueError(f"Unknown op: {self.op}")

    def to_string(self) -> str:
        """Human-readable expression string."""
        if self.op == "x":
            return "x"
        elif self.op == "const":
            if self.value == int(self.value) and abs(self.value) < 100:
                return str(int(self.value))
            return f"{self.value:.4g}"
        elif self.op in UNARY_OPS:
            child_str = self.left.to_string()
            if self.op == "neg":
                return f"(-{child_str})"
            return f"{self.op}({child_str})"
        elif self.op in BINARY_OPS:
            left_str = self.left.to_string()
            right_str = self.right.to_string()
            return f"({left_str} {self.op} {right_str})"
        return f"?{self.op}?"

    def update_stats(self) -> None:
        """Recompute depth and size from this node down."""
        if self.left is None and self.right is None:
            self.depth = 0
            self.size = 1
        elif self.right is None:
            self.left.update_stats()
            self.depth = self.left.depth + 1
            self.size = self.left.size + 1
        else:
            self.left.update_stats()
            self.right.update_stats()
            self.depth = max(self.left.depth, self.right.depth) + 1
            self.size = self.left.size + self.right.size + 1

    def copy(self) -> ExprNode:
        """Deep copy of this subtree."""
        node = ExprNode(op=self.op, value=self.value)
        if self.left is not None:
            node.left = self.left.copy()
        if self.right is not None:
            node.right = self.right.copy()
        node.depth = self.depth
        node.size = self.size
        return node

    def complexity(self) -> int:
        """Weighted complexity for BIC (variables=1, constants=3, ops=cost)."""
        cost = OP_COST.get(self.op, 1)
        if self.left is not None:
            cost += self.left.complexity()
        if self.right is not None:
            cost += self.right.complexity()
        return cost


# ── Factory functions ────────────────────────────────────────────────────

def var_x() -> ExprNode:
    """Variable leaf: x"""
    return ExprNode(op="x", depth=0, size=1)

def const(v: float) -> ExprNode:
    """Constant leaf."""
    return ExprNode(op="const", value=v, depth=0, size=1)

def unary(op: str, child: ExprNode) -> ExprNode:
    """Unary operator node."""
    node = ExprNode(op=op, left=child)
    node.update_stats()
    return node

def binary(op: str, left: ExprNode, right: ExprNode) -> ExprNode:
    """Binary operator node."""
    node = ExprNode(op=op, left=left, right=right)
    node.update_stats()
    return node


# ── Common expressions ──────────────────────────────────────────────────

def x_squared() -> ExprNode:
    """x²"""
    return binary("*", var_x(), var_x())

def x_cubed() -> ExprNode:
    """x³"""
    return binary("*", x_squared(), var_x())

def sqrt_x() -> ExprNode:
    """√x"""
    return unary("sqrt", var_x())

def x_sqrt_x() -> ExprNode:
    """x·√x = x^1.5"""
    return binary("*", var_x(), sqrt_x())

def sin_x() -> ExprNode:
    """sin(x)"""
    return unary("sin", var_x())

def cos_x() -> ExprNode:
    """cos(x)"""
    return unary("cos", var_x())

def exp_neg_x2() -> ExprNode:
    """exp(-x²)"""
    return unary("exp", binary("*", const(-1.0), x_squared()))