"""Utilities to convert a `dysts` dynamical system object's rhs to SymPy.

This module inspects the source of an object's RHS method (by default
named ``rhs``), parses the function using ``ast``, and converts the
returned expression(s) into SymPy expressions.

The conversion is intentionally conservative and aims to handle common
patterns used in simple rhs implementations, e.g. returning a tuple/list
of arithmetic expressions, using indexing into a state vector (``x[0]``),
and calls to common ``numpy``/``math`` functions (``np.sin``, ``math.exp``, ...).

Limitations:
- It does not execute arbitrary code from the inspected function.
- Complex control flow, loops, or non-trivial Python constructs may not
  be fully supported.

Example
-------
from dysts.flows import Lorenz
from inspect_to_sympy import object_to_sympy_rhs

lor = Lorenz()
symbols, exprs, lambda_rhs = object_to_sympy_rhs(lor)
# `symbols` is a list of SymPy symbols for the state vector
# `exprs` is a list of SymPy expressions for the RHS
# `lambda_rhs` is a SymPy Lambda mapping state symbols -> rhs expressions
"""

from __future__ import annotations

import ast
import inspect
import textwrap
from typing import Any, Callable, Dict, List, Tuple

import numpy as np
import sympy as sp
from dysts.base import BaseDyn


def _is_name(node: ast.AST, name: str) -> bool:
    return isinstance(node, ast.Name) and node.id == name


class _ASTToSympy(ast.NodeVisitor):
    def __init__(
        self,
        state_name: str,
        state_symbols: List[sp.Symbol],
        locals_map: Dict[str, Any],
    ):
        self.state_name = state_name
        self.state_symbols = state_symbols
        self.locals = dict(locals_map)

    def generic_visit(self, node):
        raise NotImplementedError(f"AST node not supported: {node!r}")

    def visit_Constant(self, node: ast.Constant):
        return sp.sympify(node.value)

    def visit_Num(self, node: ast.Num):
        return sp.sympify(node.n)

    def visit_Name(self, node: ast.Name):
        if node.id in self.locals:
            return self.locals[node.id]
        return sp.Symbol(node.id)

    def visit_Tuple(self, node: ast.Tuple):
        elems = []
        for elt in node.elts:
            val = self.visit(elt)
            if isinstance(val, (list, tuple)):
                elems.extend(list(val))
            else:
                elems.append(val)
        return tuple(elems)

    def visit_List(self, node: ast.List):
        elems = []
        for elt in node.elts:
            val = self.visit(elt)
            if isinstance(val, (list, tuple)):
                elems.extend(list(val))
            else:
                elems.append(val)
        return elems

    def visit_Starred(self, node: ast.Starred):
        # Handle starred expressions like `*x` in list/tuple literals.
        # If the starred value is the state vector name, expand to state symbols.
        if isinstance(node.value, ast.Name) and node.value.id == self.state_name:
            return tuple(self.state_symbols)
        # Otherwise, evaluate the value and if it is a sequence, return its items
        val = self.visit(node.value)
        if isinstance(val, (list, tuple)):
            return tuple(val)
        raise NotImplementedError(
            "Unsupported starred expression; cannot expand non-iterable"
        )

    def visit_BinOp(self, node: ast.BinOp):
        left = self.visit(node.left)
        right = self.visit(node.right)
        if isinstance(node.op, ast.Add):
            return left + right
        if isinstance(node.op, ast.Sub):
            return left - right
        if isinstance(node.op, ast.Mult):
            return left * right
        if isinstance(node.op, ast.Div):
            return left / right
        if isinstance(node.op, ast.Pow):
            return left**right
        if isinstance(node.op, ast.Mod):
            return left % right
        raise NotImplementedError(f"Binary op not supported: {node.op!r}")

    def visit_UnaryOp(self, node: ast.UnaryOp):
        operand = self.visit(node.operand)
        if isinstance(node.op, ast.USub):
            return -operand
        if isinstance(node.op, ast.UAdd):
            return +operand
        raise NotImplementedError(f"Unary op not supported: {node.op!r}")

    def visit_Call(self, node: ast.Call):
        # Determine function name
        func = node.func
        func_name = None
        mod_name = None

        if isinstance(func, ast.Name):
            func_name = func.id
        elif isinstance(func, ast.Attribute):
            # e.g. np.sin or math.exp
            if isinstance(func.value, ast.Name):
                mod_name = func.value.id
            func_name = func.attr
        else:
            raise NotImplementedError(
                f"Call to unsupported func node: {ast.dump(func)}"
            )

        # Map common numpy/math functions to sympy
        func_map = {
            "sin": sp.sin,
            "cos": sp.cos,
            "tan": sp.tan,
            "exp": sp.exp,
            "log": sp.log,
            "sqrt": sp.sqrt,
            "abs": sp.Abs,
            "atan": sp.atan,
            "asin": sp.asin,
            "acos": sp.acos,
        }

        args = [self.visit(a) for a in node.args]

        # Special-case array constructors: return underlying list/tuple
        if func_name in ("array", "asarray") and mod_name in ("np", "numpy"):
            # expect a single positional arg that's a list/tuple
            if len(args) == 1:
                return args[0]

        if func_name in func_map:
            return func_map[func_name](*args)

        # Unknown function: create a Sympy Function
        symf = sp.Function(func_name)
        return symf(*args)

    def visit_Subscript(self, node: ast.Subscript):
        # Support patterns like x[0] where x is the state vector name
        value = node.value
        # handle simple constant index
        if _is_name(value, self.state_name):
            # Python >=3.9: slice is directly the node.slice
            idx_node = node.slice
            if isinstance(idx_node, ast.Constant):
                idx = idx_node.value
            else:
                raise NotImplementedError(
                    "Only constant indices into state vector supported"
                )
            return self.state_symbols[idx]

        # If it's something else, try to evaluate generically
        base = self.visit(value)
        # slice may be constant
        if isinstance(node.slice, ast.Constant):
            key = node.slice.value
            return base[key]
        raise NotImplementedError("Unsupported subscript pattern")


def _numeric_consistency_check(
    dysts_flow: BaseDyn,
    rhsfunc: Callable,
    arg_names: List[str],
    state_names: List[str],
    vector_mode: bool,
    sys_dim: int,
    lambda_rhs: sp.Lambda,
) -> None:
    """Compare the original dysts rhs function to the SymPy-derived lambda.

    Raises a RuntimeError if they disagree.
    """
    # default to nonnegative support (e.g. Lotka volterra)
    random_state = np.random.standard_exponential(size=sys_dim)

    # Construct call arguments for the original function (bound method).
    call_args = []
    for name in arg_names:
        if name == "self":
            continue
        if name in state_names and not vector_mode:
            idx = state_names.index(name)
            call_args.append(random_state[idx])
        elif name in state_names and vector_mode:
            call_args.append(np.asarray(random_state, dtype=float))
        elif name == "t":
            call_args.append(float(np.random.standard_normal(size=())))
        else:
            call_args.append(dysts_flow.params[name])

    dysts_val = rhsfunc(*call_args)
    orig_arr = np.asarray(dysts_val, dtype=float).ravel()

    sym_val = lambda_rhs(*tuple(random_state))
    sym_arr = np.asarray(sym_val, dtype=float).ravel()

    if orig_arr.shape != sym_arr.shape:
        raise RuntimeError(
            f"_rhs shape {orig_arr.shape} != sympy shape {sym_arr.shape}"
        )

    if not np.allclose(orig_arr, sym_arr, rtol=1e-6, atol=1e-9):
        raise RuntimeError("Numeric mismatch between original and sympy conversion.")


def dynsys_to_sympy(
    obj: Any, func_name: str = "_rhs"
) -> Tuple[List[sp.Symbol], List[sp.Expr], sp.Lambda]:
    """Inspect ``obj`` for a method named ``func_name`` and return a SymPy
    representation of its RHS.

    Returns:
        a tuple ``(state_symbols, exprs, lambda_rhs)`` where ``state_symbols``
    is a list of SymPy symbols for the state vector, ``exprs`` is a list of
    SymPy expressions for the RHS components, and ``lambda_rhs`` is a SymPy
    Lambda mapping the state symbols to the RHS vector.

    Example:

    >>> from dysts.flows import Lorenz
    >>> from inspect_to_sympy import dynsys_to_sympy
    >>> lor = Lorenz()
    >>> symbols, exprs, lambda_rhs = dynsys_to_sympy(lor)
    >>> print(lor._rhs(1, 2, 3, t=0.0, **lor.params))
    (10, 23, -6.0009999999999994)

    >>> print(tuple(lambda_rhs(1, 2, 3)))
    (10, 23, -6.00100000000000)

    """

    if not hasattr(obj, func_name):
        raise AttributeError(f"Object has no attribute {func_name!r}")

    func = getattr(obj, func_name)
    src = inspect.getsource(func)
    src = textwrap.dedent(src)

    parsed = ast.parse(src)

    # Find first FunctionDef
    fndef = None
    for node in parsed.body:
        if isinstance(node, ast.FunctionDef):
            fndef = node
            break
    if fndef is None:
        raise RuntimeError("No function definition found in source")

    # Determine state argument names. Common dysts signature:
    # (self, *states, t, *parameters). Prefer obj.dimension when available.
    arg_names = [a.arg for a in fndef.args.args]
    if len(arg_names) == 0:
        raise RuntimeError("Function has no arguments")

    start_idx = 0
    if arg_names[0] == "self":
        start_idx = 1

    vector_mode = False
    state_args: List[str]
    t_idx = None
    if "t" in arg_names:
        t_idx = arg_names.index("t")

    if hasattr(obj, "dimension") and isinstance(getattr(obj, "dimension"), int):
        n_state = int(getattr(obj, "dimension"))
        if t_idx is not None:
            potential = arg_names[start_idx:t_idx]
            if len(potential) >= n_state:
                state_args = potential[:n_state]
            else:
                state_args = [arg_names[start_idx]]
                vector_mode = True
        else:
            potential = arg_names[start_idx:]
            if len(potential) >= n_state:
                state_args = potential[:n_state]
            else:
                state_args = [arg_names[start_idx]]
                vector_mode = True
    else:
        if t_idx is not None:
            state_args = arg_names[start_idx:t_idx]
            if len(state_args) == 0:
                state_args = [arg_names[start_idx]]
                vector_mode = True
            elif len(state_args) == 1:
                # single name could be vector or scalar; assume vector-mode
                vector_mode = True
        else:
            state_args = [arg_names[start_idx]]
            vector_mode = True

    # If vector_mode, inspect AST for subscript/index usage or tuple unpacking
    if vector_mode:
        state_name = state_args[0]
        max_index = -1
        unpack_size = None
        for node in ast.walk(fndef):
            if (
                isinstance(node, ast.Subscript)
                and isinstance(node.value, ast.Name)
                and node.value.id == state_name
            ):
                sl = node.slice
                if isinstance(sl, ast.Constant) and isinstance(sl.value, int):
                    if sl.value > max_index:
                        max_index = sl.value
            if isinstance(node, ast.Assign):
                if isinstance(node.value, ast.Name) and node.value.id == state_name:
                    targets = node.targets
                    if len(targets) == 1 and isinstance(
                        targets[0], (ast.Tuple, ast.List)
                    ):
                        unpack_size = len(targets[0].elts)

        if unpack_size is not None:
            n_state = unpack_size
        elif max_index >= 0:
            n_state = max_index + 1
        else:
            n_state = int(getattr(obj, "dimension", 3))

        state_symbols = [sp.Symbol(f"x{i}") for i in range(n_state)]
        primary_state_name = state_name
    else:
        # individual state args -> use their arg names as symbol names
        state_symbols = [sp.Symbol(n) for n in state_args]
        primary_state_name = state_args[0] if len(state_args) > 0 else "x"

    # Build locals mapping from known state arg names and parameters
    locals_map: Dict[str, Any] = {}
    for i, name in enumerate(state_args):
        if i < len(state_symbols):
            locals_map[name] = state_symbols[i]

    # map parameters (if present) to numeric values or symbols
    if hasattr(obj, "parameters") and isinstance(getattr(obj, "parameters"), dict):
        params = getattr(obj, "parameters")
        if t_idx is not None:
            param_arg_names = arg_names[t_idx + 1 :]
        else:
            param_arg_names = []
        for pname in param_arg_names:
            if pname in params:
                locals_map[pname] = sp.sympify(params[pname])
            else:
                locals_map[pname] = sp.Symbol(pname)

    converter = _ASTToSympy(primary_state_name, state_symbols, locals_map)

    return_expr = None
    # Walk through function body statements, handle Assign and Return
    for stmt in fndef.body:
        if isinstance(stmt, ast.Assign):
            # only simple single-target assignments supported
            if len(stmt.targets) != 1:
                raise ValueError("Only single-target assignments supported")
            target = stmt.targets[0]
            if isinstance(target, ast.Name):
                value_expr = converter.visit(stmt.value)
                locals_map[target.id] = value_expr
            elif (
                isinstance(target, (ast.Tuple, ast.List))
                and isinstance(stmt.value, ast.Name)
                and stmt.value.id == state_name
            ):
                # unpacking like a,b,c = x -> map names to state symbols
                for i, elt in enumerate(target.elts):
                    if isinstance(elt, ast.Name):
                        locals_map[elt.id] = state_symbols[i]
        elif isinstance(stmt, ast.Return):
            return_expr = stmt.value

    if return_expr is None:
        # maybe last statement is an Expr with list construction;
        # try to find a Return node deep
        for node in ast.walk(fndef):
            if isinstance(node, ast.Return):
                return_expr = node.value
                break

    if return_expr is None:
        raise RuntimeError("No return expression found in function body")

    # Refresh converter with updated locals
    converter = _ASTToSympy(primary_state_name, state_symbols, locals_map)
    rhs_val = converter.visit(return_expr)

    if isinstance(rhs_val, (list, tuple)):
        exprs = list(rhs_val)
    else:
        # single expression: treat as 1-dim RHS
        exprs = [rhs_val]

    lambda_rhs = sp.Lambda(tuple(state_symbols), sp.Matrix(exprs))

    # Run numeric consistency guard (raises on mismatch)
    _numeric_consistency_check(
        obj,
        func,
        arg_names,
        state_args,
        vector_mode,
        len(state_symbols),
        lambda_rhs,
    )

    return state_symbols, exprs, lambda_rhs


__all__ = ["dynsys_to_sympy"]
