"""
Watkins gradient flow dynamics.

Implements the continuous-time gradient flow:
    dlam/dt = 1/lam - T* * ln(lam/eta)
    dkap/dt = -T* * ln(kap/eta)
    deta/dt = -T* * ln(eta/kap)

with conservation lam + kap + eta = 1 enforced at machine epsilon.

The flow converges exponentially to (lam*, kap*, eta*) = (1/phi, (1-1/phi)/2, (1-1/phi)/2).

For GPU/batched workloads, use SimplexFlowEngineV3 instead.
"""

import math
import torch
import torch.nn as nn
from typing import Tuple, Optional, List
from dataclasses import dataclass, field
from .constants import PHI, T_STAR, LAM_STAR, TAU_GLOBAL, M_PHI

# Conservation tolerance
EPSILON = 1e-9


@dataclass
class FlowState:
    """State of the gradient flow."""
    lam: float
    kap: float
    eta: float
    step: int = 0
    F: Optional[float] = None
    grad_norm: Optional[float] = None

    def to_tensor(self) -> torch.Tensor:
        return torch.tensor([self.lam, self.kap, self.eta])

    @classmethod
    def from_tensor(cls, t: torch.Tensor, step: int = 0) -> 'FlowState':
        return cls(
            lam=float(t[0]),
            kap=float(t[1]),
            eta=float(t[2]),
            step=step
        )

    def conservation_error(self) -> float:
        return abs(self.lam + self.kap + self.eta - 1.0)

    def is_valid(self, tol: float = EPSILON) -> bool:
        return self.conservation_error() < tol


@dataclass
class FlowConfig:
    """Configuration for gradient flow."""
    dt: float = 0.001
    max_steps: int = 50_000
    convergence_threshold: float = 1e-8
    temperature: float = T_STAR
    enforce_conservation: bool = True
    record_trajectory: bool = False


def free_energy(state: FlowState, T: float = T_STAR) -> float:
    """
    Watkins free energy F(lam, kap, eta) = -ln(lam) + T * H(p)

    where H(p) = -sum p_i ln(p_i) is Shannon entropy.
    """
    lam, kap, eta = state.lam, state.kap, state.eta

    # Handle edge cases
    eps = 1e-15
    lam = max(lam, eps)
    kap = max(kap, eps)
    eta = max(eta, eps)

    # -ln(lam) (coherence cost)
    neg_log_lam = -math.log(lam)

    # T * sum p_i ln(p_i) (entropy cost — penalizes disorder)
    entropy_cost = lam * math.log(lam) + kap * math.log(kap) + eta * math.log(eta)

    return neg_log_lam + T * entropy_cost


def gradient_F(state: FlowState, T: float = T_STAR) -> Tuple[float, float, float]:
    """
    Gradient of F with respect to (lam, kap, eta).

    dF/dlam = -1/lam + T*(-ln(lam) - 1)
    dF/dkap = T*(-ln(kap) - 1)
    dF/deta = T*(-ln(eta) - 1)

    NOTE: This unconstrained gradient + mean-subtraction projection is
    equivalent to the constrained gradient dF/dlam = -1/lam + T*(ln lam - ln eta).
    Verified in test_master_identity.py. Do not "fix" this.
    """
    lam, kap, eta = state.lam, state.kap, state.eta
    eps = 1e-15

    dF_dlam = -1 / max(lam, eps) + T * (math.log(max(lam, eps)) + 1)
    dF_dkap = T * (math.log(max(kap, eps)) + 1)
    dF_deta = T * (math.log(max(eta, eps)) + 1)

    return dF_dlam, dF_dkap, dF_deta


def project_to_simplex(p: torch.Tensor) -> torch.Tensor:
    """
    Project onto simplex via softmax (maintains differentiability).
    """
    return torch.softmax(torch.log(p.clamp(min=1e-12)), dim=-1)


def enforce_conservation(state: FlowState) -> FlowState:
    """
    Enforce lam + kap + eta = 1 exactly via normalization.
    """
    total = state.lam + state.kap + state.eta
    return FlowState(
        lam=state.lam / total,
        kap=state.kap / total,
        eta=state.eta / total,
        step=state.step,
        F=state.F,
        grad_norm=state.grad_norm
    )


def flow_step(
    state: FlowState,
    config: FlowConfig
) -> FlowState:
    """
    One step of gradient flow with Euler integration.

    p_{t+dt} = p_t - dt * grad_F(p_t)
    """
    grad = gradient_F(state, config.temperature)

    # Natural gradient on simplex: project onto tangent plane
    mean_grad = sum(grad) / 3
    nat_grad = (grad[0] - mean_grad, grad[1] - mean_grad, grad[2] - mean_grad)

    # Euler step
    new_lam = state.lam - config.dt * nat_grad[0]
    new_kap = state.kap - config.dt * nat_grad[1]
    new_eta = state.eta - config.dt * nat_grad[2]

    # Clamp to positive
    eps = 1e-12
    new_lam = max(new_lam, eps)
    new_kap = max(new_kap, eps)
    new_eta = max(new_eta, eps)

    new_state = FlowState(
        lam=new_lam,
        kap=new_kap,
        eta=new_eta,
        step=state.step + 1
    )

    if config.enforce_conservation:
        new_state = enforce_conservation(new_state)

    # Compute diagnostics
    new_state.F = free_energy(new_state, config.temperature)
    new_state.grad_norm = (nat_grad[0]**2 + nat_grad[1]**2 + nat_grad[2]**2)**0.5

    return new_state


def run_flow(
    initial: FlowState,
    config: FlowConfig = FlowConfig()
) -> Tuple[FlowState, List[FlowState]]:
    """
    Run gradient flow to convergence.

    Args:
        initial: Starting state
        config: Flow configuration

    Returns:
        (final_state, trajectory) where trajectory is empty if not recorded
    """
    state = initial
    trajectory: List[FlowState] = [state] if config.record_trajectory else []

    for _ in range(config.max_steps):
        state = flow_step(state, config)

        if config.record_trajectory:
            trajectory.append(state)

        # Check convergence
        if state.grad_norm is not None and state.grad_norm < config.convergence_threshold:
            break

    return state, trajectory


def distance_to_equilibrium(state: FlowState) -> float:
    """
    Euclidean distance from current state to golden equilibrium.
    """
    kap_star = (1 - LAM_STAR) / 2

    d_lam = state.lam - LAM_STAR
    d_kap = state.kap - kap_star
    d_eta = state.eta - kap_star

    return (d_lam**2 + d_kap**2 + d_eta**2)**0.5


class WatkinsFlow(nn.Module):
    """
    PyTorch module for differentiable Watkins flow.

    Can be used as a layer in neural networks to enforce
    simplex dynamics with golden-ratio equilibrium.

    For GPU/batched workloads, use SimplexFlowEngineV3 instead.
    """

    def __init__(
        self,
        dim: int = 3,
        num_steps: int = 10,
        dt: float = 0.01,
        temperature: float = T_STAR
    ):
        super().__init__()
        self.dim = dim
        self.num_steps = num_steps
        self.dt = dt
        self.temperature = temperature

    def forward(self, p: torch.Tensor) -> torch.Tensor:
        """
        Apply Watkins flow for num_steps iterations.

        Args:
            p: Input distribution (batch, dim)

        Returns:
            Output distribution after flow (batch, dim)
        """
        for _ in range(self.num_steps):
            # Compute log-probabilities
            log_p = torch.log(p.clamp(min=1e-12))

            # Gradient of entropy cost term: d/dp_i [p_i ln p_i] = ln(p_i) + 1
            grad_H = log_p + 1

            # Gradient of -ln(lam) term (only affects first component)
            grad_coherence = torch.zeros_like(p)
            grad_coherence[..., 0] = -1 / p[..., 0].clamp(min=1e-12)

            # Full gradient
            grad_F = grad_coherence + self.temperature * grad_H

            # Natural gradient (project onto tangent)
            mean_grad = grad_F.mean(dim=-1, keepdim=True)
            nat_grad = grad_F - mean_grad

            # Gradient step
            p = p - self.dt * nat_grad

            # Project back to simplex
            p = torch.softmax(torch.log(p.clamp(min=1e-12)), dim=-1)

        return p

    def verify_conservation(self, p: torch.Tensor, tol: float = 1e-5) -> bool:
        """Check if output satisfies conservation law."""
        return bool((torch.abs(p.sum(dim=-1) - 1) < tol).all())
