"""
Field math utilities (PDE/field-space helpers).

Provides a grid-based PDE solver and sampling utilities for projecting
field PDEs into particle motion without hardcoded dynamics.
"""

from __future__ import annotations

from typing import Dict, List, Tuple

import re
import numpy as np

from .symbols import normalize_symbols, normalize_lhs


def normalize_pde(pde: str) -> str:
    return (pde or "").strip()


def init_grid(width: int, height: int, dx: float) -> Tuple[np.ndarray, float]:
    nx = max(2, int(width / dx))
    nz = max(2, int(height / dx))
    grid = np.zeros((nz, nx), dtype=float)
    return grid, dx


def _preprocess_expr(expr: str) -> str:
    expr = (expr or "").strip()
    # Drop common annotation tokens (LLM outputs sometimes include these).
    expr = re.sub(r"\(\s*approx\s*\)", "", expr, flags=re.IGNORECASE)
    expr = re.sub(r"\bapprox\b", "", expr, flags=re.IGNORECASE)
    # If the RHS contains multiple '=' (e.g. an annotated expansion), prefer the
    # rightmost expression, which is typically the most explicit.
    if "=" in expr:
        expr = expr.split("=")[-1]
    expr = normalize_symbols(expr)
    return expr.strip()


def parse_pde(pde: str) -> Tuple[str, int, str, str]:
    """
    Returns (var, order, lhs_coeff_expr, rhs_expr).
    Supports forms:
      ∂u/∂t = RHS
      ∂²u/∂t² = RHS
      ∂^n u / ∂t^n = RHS
      rho(∂²u/∂t²) = RHS
    """
    pde = normalize_pde(pde)
    if "=" not in pde:
        return "u", 1, "1", pde

    lhs, rhs = pde.split("=", 1)
    lhs = normalize_lhs(lhs.strip())
    rhs = rhs.strip()

    def _extract_coeff(lhs_expr: str, deriv_expr: str) -> str:
        coeff = lhs_expr.replace(deriv_expr, "").strip()
        if not coeff:
            return "1"
        # Remove leading or trailing multiplication symbols
        coeff = coeff.strip("*")
        return _preprocess_expr(coeff) or "1"

    # N-th order time derivative: ∂^n v / ∂t^n
    # Matches patterns like: ∂^2 u / ∂t^2, ∂3u/∂t3, ∂ u / ∂t
    # We normalized unicode `²` to `^2` in normalize_lhs.
    
    # Regex for derivative term
    # breakdown:
    #   ∂               literal partial
    #   \s*             whitespace
    #   (?:\^?(\d+))?   optional order: ^2, 2, or nothing (order 1)
    #   \s*             whitespace
    #   ([a-zA-Z_]\w*)  variable name (group 2)
    #   \s*             whitespace
    #   /               division
    #   \s*             whitespace
    #   ∂t              ∂t
    #   (?:\^?(\d+))?   optional order at bottom: ^2, 2, or nothing
    
    pattern = r"∂\s*(?:\^?(\d+))?\s*([a-zA-Z_]\w*)\s*/\s*∂t\s*(?:\^?(\d+))?"
    m = re.search(pattern, lhs)
    
    if m:
        # group 1: order top, group 2: var, group 3: order bottom
        # if inconsistent orders, we trust top or bottom if they match?
        # Usually they match. If missing, assume 1.
        ord1 = m.group(1)
        var = m.group(2)
        ord2 = m.group(3)
        
        order = 1
        if ord1:
            order = int(ord1)
        elif ord2:
            order = int(ord2)
            
        coeff = _extract_coeff(lhs, m.group(0))
        return var, order, coeff, _preprocess_expr(rhs)

    return "u", 1, "1", _preprocess_expr(rhs)


def evolve_pde(grid: np.ndarray, dt: float, constants: dict) -> np.ndarray:
    """
    Generic PDE evolution step. Requires constants to define dynamics.
    If no constants are provided, returns the grid unchanged.

    Expected constants:
    - alpha: diffusion coefficient (optional)
    """
    return grid


def sample_gradient(grid: np.ndarray, x: float, z: float, dx: float) -> Tuple[float, float]:
    if grid is None:
        return 0.0, 0.0
    nz, nx = grid.shape
    ix = int(np.clip(x / dx, 1, nx - 2))
    iz = int(np.clip(z / dx, 1, nz - 2))
    du_dx = (grid[iz, ix + 1] - grid[iz, ix - 1]) / (2 * dx)
    du_dz = (grid[iz + 1, ix] - grid[iz - 1, ix]) / (2 * dx)
    return du_dx, du_dz


def _build_eval_env(vars_grid: Dict[str, np.ndarray], constants: Dict[str, float], grid_dx: float) -> Dict[str, object]:
    def lap(u):
        return dxx(u) + dzz(u)

    def dx(u):
        return (np.roll(u, -1, axis=1) - np.roll(u, 1, axis=1)) / (2 * grid_dx)

    def dz(u):
        return (np.roll(u, -1, axis=0) - np.roll(u, 1, axis=0)) / (2 * grid_dx)

    def dxx(u):
        return (np.roll(u, -1, axis=1) - 2 * u + np.roll(u, 1, axis=1)) / (grid_dx ** 2)

    def dzz(u):
        return (np.roll(u, -1, axis=0) - 2 * u + np.roll(u, 1, axis=0)) / (grid_dx ** 2)

    def gradmag(u):
        return dx(u) ** 2 + dz(u) ** 2

    def gradl1(u):
        return np.abs(dx(u)) + np.abs(dz(u))

    def grad(u):
        # Return a 2x(HxW) array so scalar multiplication broadcasts naturally
        # and div(...) can consume it.
        return np.stack((dx(u), dz(u)), axis=0)

    def div(v):
        # Accept either a tuple/list (vx, vz) or a stacked array with shape (2, ...).
        if isinstance(v, (tuple, list)) and len(v) == 2:
            return dx(v[0]) + dz(v[1])
        if isinstance(v, np.ndarray) and v.ndim >= 3 and v.shape[0] == 2:
            return dx(v[0]) + dz(v[1])
        return np.zeros_like(next(iter(vars_grid.values())))

    def sech(u):
        return 1.0 / np.cosh(u)

    def sign(u):
        return np.sign(u)

    def pos(u):
        return np.maximum(u, 0.0)

    env = {
        "np": np,
        "sin": np.sin,
        "cos": np.cos,
        "tan": np.tan,
        "sinh": np.sinh,
        "cosh": np.cosh,
        "tanh": np.tanh,
        "arcsin": np.arcsin,
        "arccos": np.arccos,
        "arctan": np.arctan,
        "log": np.log,
        "log10": np.log10,
        "log2": np.log2,
        "exp": np.exp,
        "sqrt": np.sqrt,
        "abs": np.abs,
        "pi": np.pi,
        "inf": np.inf,
        "lap": lap,
        "dx": dx,
        "dz": dz,
        "dxx": dxx,
        "dzz": dzz,
        "gradmag": gradmag,
        "gradl1": gradl1,
        "grad": grad,
        "div": div,
        "pos": pos,
        "sech": sech,
        "sign": sign,
    }

    # Add constants
    for k, v in (constants or {}).items():
        env[k] = v

    if "E" not in env:
        env["E"] = np.e
    if "e" not in env:
        env["e"] = np.e

    # Normalize common constant aliases used in test corpora / LLM outputs.
    # - Users often provide 'lambda' and 'eps' in JSON to avoid Unicode / reserved keywords.
    if "lambda" in env and "lam" not in env:
        env["lam"] = env["lambda"]
    if "eps" in env and "epsilon" not in env:
        env["epsilon"] = env["eps"]

    # Common aliasing for constants
    if "w" in env and "omega" not in env:
        env["omega"] = env["w"]
    if "w2" in env and "omega2" not in env:
        env["omega2"] = env["w2"]
    if "ax" in env and "alphax" not in env:
        env["alphax"] = env["ax"]
    if "az" in env and "alphaz" not in env:
        env["alphaz"] = env["az"]
    if "th" in env and "theta" not in env:
        env["theta"] = env["th"]
    if "damp" in env and "zeta" not in env:
        env["zeta"] = env["damp"]

    # Provide time symbol if present
    if "t" not in env:
        env["t"] = float((constants or {}).get("t", 0.0))

    # Add variable grids
    for k, v in vars_grid.items():
        env[k] = v

    return env


def evaluate_rhs(rhs_expr: str, vars_grid: Dict[str, np.ndarray], constants: Dict[str, float], grid_dx: float) -> np.ndarray:
    """
    Evaluate RHS expression on grid using numpy operations.
    """
    rhs_expr = _preprocess_expr(rhs_expr)
    env = _build_eval_env(vars_grid, constants, grid_dx)

    # Provide default zero grids for missing symbols referenced in RHS
    if vars_grid:
        zero_grid = np.zeros_like(next(iter(vars_grid.values())))
        identifiers = set(re.findall(r"\b([A-Za-z_][A-Za-z0-9_]*)\b", rhs_expr))
        reserved = {
            "np",
            "sin", "cos", "tan", "sinh", "cosh", "tanh",
            "arcsin", "arccos", "arctan",
            "exp", "sqrt", "abs", "log", "log10", "log2",
            "sech", "sign",
            "pi", "inf", "E", "e",
            "lap", "dx", "dz", "dxx", "dzz", "grad", "div", "gradmag", "gradl1", "pos", "t",
        }
        for name in identifiers:
            if name in reserved:
                continue
            if name not in env:
                env[name] = zero_grid

    if not vars_grid:
        return np.zeros((2, 2), dtype=float)

    try:
        return eval(rhs_expr, {}, env)
    except Exception as exc:
        raise RuntimeError(f"PDE RHS eval failed for '{rhs_expr}': {exc}")


def evaluate_rhs_compiled(rhs_expr: str, vars_grid: Dict[str, np.ndarray], constants: Dict[str, float], grid_dx: float) -> np.ndarray:
    """Evaluate already-compiled RHS expression (no preprocessing)."""
    env = _build_eval_env(vars_grid, constants, grid_dx)

    if vars_grid:
        zero_grid = np.zeros_like(next(iter(vars_grid.values())))
        identifiers = set(re.findall(r"\b([A-Za-z_][A-Za-z0-9_]*)\b", rhs_expr))
        reserved = {
            "np",
            "sin", "cos", "tan", "sinh", "cosh", "tanh",
            "arcsin", "arccos", "arctan",
            "exp", "sqrt", "abs", "log", "log10", "log2",
            "sech", "sign",
            "pi", "inf", "E", "e",
            "lap", "dx", "dz", "dxx", "dzz", "grad", "div", "gradmag", "gradl1", "pos", "t",
        }
        for name in identifiers:
            if name in reserved:
                continue
            if name not in env:
                env[name] = zero_grid

    if not vars_grid:
        return np.zeros((2, 2), dtype=float)

    try:
        return eval(rhs_expr, {}, env)
    except Exception as exc:
        raise RuntimeError(f"PDE RHS eval failed for '{rhs_expr}': {exc}")


def evaluate_scalar(expr: str, vars_grid: Dict[str, np.ndarray], constants: Dict[str, float], grid_dx: float) -> float:
    """Evaluate a scalar coefficient expression with the same env as RHS."""
    if expr in ("", "1"):
        return 1.0
    rhs = evaluate_rhs(expr, vars_grid, constants, grid_dx)
    if np.isscalar(rhs):
        return float(rhs)
    try:
        return float(np.mean(rhs))
    except Exception as exc:
        raise RuntimeError(f"PDE coeff eval failed for '{expr}': {exc}")


def evaluate_scalar_compiled(expr: str, vars_grid: Dict[str, np.ndarray], constants: Dict[str, float], grid_dx: float) -> float:
    """Evaluate scalar coefficient expression without preprocessing."""
    if expr in ("", "1"):
        return 1.0
    rhs = evaluate_rhs_compiled(expr, vars_grid, constants, grid_dx)
    if np.isscalar(rhs):
        return float(rhs)
    try:
        return float(np.mean(rhs))
    except Exception as exc:
        raise RuntimeError(f"PDE coeff eval failed for '{expr}': {exc}")


def step_pdes(
    field_pdes: List[str],
    grids: Dict[str, np.ndarray],
    constants: Dict[str, float],
    grid_dx: float,
    dt: float,
    external_grids: Dict[str, np.ndarray] | None = None,
) -> List[str]:
    """Advance PDEs by one time step. Returns list of error strings."""
    errors: List[str] = []
    if not field_pdes or not grids:
        return errors

    vars_grid = dict(grids)
    if external_grids:
        vars_grid.update(external_grids)

    for pde in field_pdes:
        try:
            var, order, coeff_expr, rhs_expr = parse_pde(pde)
            if var not in grids:
                grids[var] = np.zeros_like(next(iter(grids.values())))

            rhs = evaluate_rhs(rhs_expr, vars_grid, constants, grid_dx)
            coeff = evaluate_scalar(coeff_expr, vars_grid, constants, grid_dx)
            rhs = rhs / float(coeff or 1.0)

            if order == 2:
                vname = f"{var}_t"
                if vname not in grids:
                    grids[vname] = np.zeros_like(grids[var])
                grids[vname] = grids[vname] + dt * rhs
                grids[var] = grids[var] + dt * grids[vname]
            else:
                grids[var] = grids[var] + dt * rhs
        except Exception as exc:
            errors.append(f"PDE eval failed for '{pde}': {exc}")

    return errors


def step_compiled_pdes(
    compiled_pdes: List[Dict[str, object]],
    grids: Dict[str, np.ndarray],
    constants: Dict[str, float],
    grid_dx: float,
    dt: float,
    external_grids: Dict[str, np.ndarray] | None = None,
) -> List[str]:
    """Advance already-compiled PDEs. Returns list of error strings."""
    errors: List[str] = []
    if not compiled_pdes or not grids:
        return errors

    vars_grid = dict(grids)
    if external_grids:
        vars_grid.update(external_grids)

    for entry in compiled_pdes:
        try:
            var = entry.get("var", "u")
            order = int(entry.get("order", 1))
            coeff_expr = str(entry.get("coeff_expr", "1"))
            rhs_expr = str(entry.get("rhs_expr", "0"))

            if var not in grids:
                grids[var] = np.zeros_like(next(iter(grids.values())))

            rhs = evaluate_rhs_compiled(rhs_expr, vars_grid, constants, grid_dx)
            coeff = evaluate_scalar_compiled(coeff_expr, vars_grid, constants, grid_dx)
            rhs = rhs / float(coeff or 1.0)

            if order == 2:
                vname = f"{var}_t"
                if vname not in grids:
                    grids[vname] = np.zeros_like(grids[var])
                grids[vname] = grids[vname] + dt * rhs
                grids[var] = grids[var] + dt * grids[vname]
            else:
                grids[var] = grids[var] + dt * rhs
        except Exception as exc:
            errors.append(f"PDE eval failed for '{entry}': {exc}")

    return errors
