"""
Hamiltonian Mechanics Implementation

Provides tools for deriving Hamilton's equations of motion.
"""

from typing import Dict, List, Optional

import sympy as sp

from ..base import PhysicsDomain


class HamiltonianMechanics(PhysicsDomain):
    """
    Hamiltonian mechanics domain using Hamilton's equations.

    dq/dt = ∂H/∂p
    dp/dt = -∂H/∂q
    """

    def __init__(self, name: str = "hamiltonian_system"):
        super().__init__(name)
        self._hamiltonian: Optional[sp.Expr] = None
        self._symbol_cache: Dict[str, sp.Symbol] = {}

    def get_symbol(self, name: str, **assumptions) -> sp.Symbol:
        """Get or create a symbol with caching."""
        if name not in self._symbol_cache:
            default_assumptions = {"real": True}
            default_assumptions.update(assumptions)
            self._symbol_cache[name] = sp.Symbol(name, **default_assumptions)
        return self._symbol_cache[name]

    def set_hamiltonian(self, expr: sp.Expr) -> None:
        """Set the Hamiltonian directly."""
        self._hamiltonian = expr

    def define_lagrangian(self) -> sp.Expr:
        """
        Convert Hamiltonian to Lagrangian (inverse Legendre transform).

        L = Σ(p_i * q̇_i) - H
        """
        if self._hamiltonian is None:
            raise ValueError("Hamiltonian not defined")

        L = sp.S.Zero
        for q in self.coordinates:
            q_dot = self.get_symbol(f"{q}_dot")
            p = self.get_symbol(f"p_{q}")
            L += p * q_dot

        L = L - self._hamiltonian
        return L

    def define_hamiltonian(self) -> sp.Expr:
        """Return the Hamiltonian expression."""
        if self._hamiltonian is None:
            raise ValueError("Hamiltonian not defined")
        return self._hamiltonian

    def derive_equations_of_motion(self) -> Dict[str, sp.Expr]:
        """
        Derive Hamilton's equations of motion.

        Returns:
            Dictionary with q_dot and p_dot expressions
        """
        if self._hamiltonian is None:
            raise ValueError("Hamiltonian not defined")

        H = self._hamiltonian
        equations = {}

        for q in self.coordinates:
            q_sym = self.get_symbol(q)
            p_sym = self.get_symbol(f"p_{q}")

            # dq/dt = ∂H/∂p
            q_dot_expr = sp.diff(H, p_sym)
            equations[f"{q}_dot"] = sp.simplify(q_dot_expr)

            # dp/dt = -∂H/∂q
            p_dot_expr = -sp.diff(H, q_sym)
            equations[f"p_{q}_dot"] = sp.simplify(p_dot_expr)

        self.equations = equations
        self._is_compiled = True
        return equations

    def get_state_variables(self) -> List[str]:
        """Get state variables (q, p pairs) for ODE system."""
        state = []
        for q in self.coordinates:
            state.extend([q, f"p_{q}"])
        return state

    def get_conserved_quantities(self) -> Dict[str, sp.Expr]:
        """Hamiltonian is conserved for time-independent systems."""
        if self._hamiltonian is not None:
            return {"hamiltonian": self._hamiltonian}
        return {}
