"""virelaix_core/simplex_flow_v3.py -- Fixed Gradient Flow Engine

CSID: SF-003
Version: 3.0.0

DISCOVERY: The v2 SimplexFlowEnginePyTorch has a gradient-space mismatch bug.
It computes dF/ds (gradient w.r.t. softmax output) but applies the update
to the unconstrained parameters u (log-space). This means:

    u_new = u_old - dt * dF/ds   <- WRONG (mixed coordinate systems)

The correct update requires chain rule through the softmax:

    u_new = u_old - dt * dF/du
    where dF/du_i = Sum_j (dF/ds_j)(ds_j/du_i)
    and ds_j/du_i = s_j(delta_{ij} - s_i)   [Jacobian of softmax]

This module provides SimplexFlowEngineV3 with:
1. Correct gradient chain rule through softmax
2. Proper Armijo line search in unconstrained space
3. Multi-start convergence for robustness validation
4. Convergence diagnostics and metrics export

Author: Claude Opus 4.6 (Cowork session, Feb 13 2026)
Discovery: Gradient-space mismatch bug found during convergence testing
"""

import torch
import torch.nn as nn
import numpy as np
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, field
import json
from pathlib import Path
import time

from .algosignal_v2 import (
    AlgoSignal, WATKINS_LAMBDA, WATKINS_KAPPA, WATKINS_ETA, PHI, T_STAR
)


@dataclass
class ConvergenceMetrics:
    """Detailed metrics from a flow convergence run."""
    converged: bool
    steps: int
    final_state: List[float]
    final_F: float
    gradient_norm: float
    conservation_error: float
    watkins_distance: float
    at_watkins: bool
    wall_time_ms: float
    method: str
    # Trajectory subsampled for efficiency
    trajectory_subsample: List[List[float]] = field(default_factory=list)
    F_history: List[float] = field(default_factory=list)
    grad_history: List[float] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            'converged': self.converged,
            'steps': self.steps,
            'final_state': self.final_state,
            'final_F': self.final_F,
            'gradient_norm': self.gradient_norm,
            'conservation_error': self.conservation_error,
            'watkins_distance': self.watkins_distance,
            'at_watkins': self.at_watkins,
            'wall_time_ms': self.wall_time_ms,
            'method': self.method,
        }


class SimplexFlowEngineV3:
    """Fixed gradient flow engine with proper chain rule through softmax.

    Key fix: Uses torch.autograd to compute dF/du (gradient in unconstrained
    space) rather than manually applying dF/ds to unconstrained parameters.

    This is both simpler AND correct -- let PyTorch handle the chain rule.

    Example:
        >>> signal = AlgoSignal(lambda_init=0.5, kappa_init=0.3, eta_init=0.2)
        >>> engine = SimplexFlowEngineV3(signal)
        >>> metrics = engine.run_to_convergence()
        >>> assert metrics.converged
        >>> assert metrics.at_watkins
    """

    def __init__(
        self,
        signal: Optional[AlgoSignal] = None,
        dt: float = 0.01,
        max_steps: int = 10000,
        tolerance: float = 1e-8,
        adaptive: bool = True,
    ):
        if signal is None:
            # Random start on simplex
            r = torch.rand(3, dtype=torch.float64)
            r = r / r.sum()
            signal = AlgoSignal(
                lambda_init=r[0].item(),
                kappa_init=r[1].item(),
                eta_init=r[2].item(),
            )

        self.signal = signal
        self.dt = dt
        self.max_steps = max_steps
        self.tolerance = tolerance
        self.adaptive = adaptive

        # Trajectory storage
        self._trajectory: List[np.ndarray] = []
        self._F_values: List[float] = []
        self._grad_norms: List[float] = []
        self._step_count = 0

    def _compute_F_from_unconstrained(self) -> torch.Tensor:
        """Compute F as a function of unconstrained params (for autograd)."""
        state = torch.softmax(self.signal.unconstrained, dim=0)
        lam, kap, eta = state[0], state[1], state[2]
        F = -torch.log(lam) + self.signal.T_STAR * (
            lam * torch.log(lam) +
            kap * torch.log(kap) +
            eta * torch.log(eta)
        )
        return F

    @staticmethod
    def projected_gradient_norm(grad_simplex: np.ndarray) -> float:
        """Compute gradient norm projected onto simplex tangent plane.

        DISCOVERY: The full gradient dF/ds at the Watkins equilibrium is
        [-0.903, -0.903, -0.903] with |g| = 1.56. But this includes the
        normal component to the simplex (the lambda+kappa+eta=1 constraint surface).

        The MEANINGFUL gradient for constrained optimization is the projection
        onto the tangent plane of the simplex. At equilibrium, this is ~0.

        Tangent plane normal: n = [1,1,1]/sqrt(3)
        Projected gradient: g_proj = g - (g.n)n
        """
        n = np.array([1, 1, 1]) / np.sqrt(3)
        g_proj = grad_simplex - np.dot(grad_simplex, n) * n
        return float(np.linalg.norm(g_proj))

    def step(self) -> float:
        """Take one gradient descent step with proper chain rule.

        Uses torch.autograd.grad to compute dF/du directly, which
        automatically applies the chain rule through softmax.

        Returns:
            Projected gradient norm (simplex tangent plane) after step
        """
        # Compute F and its gradient w.r.t. unconstrained params
        F = self._compute_F_from_unconstrained()
        grad_u = torch.autograd.grad(F, self.signal.unconstrained)[0]

        dt_actual = self.dt

        if self.adaptive:
            # Armijo line search in unconstrained space
            F_current = F.item()
            grad_norm_sq = torch.sum(grad_u ** 2).item()
            c = 0.1
            tau = 0.5
            backup = self.signal.unconstrained.data.clone()

            for _ in range(15):
                with torch.no_grad():
                    self.signal.unconstrained.copy_(backup - dt_actual * grad_u)
                F_new = self._compute_F_from_unconstrained().item()
                if F_new <= F_current - c * dt_actual * grad_norm_sq:
                    break
                dt_actual *= tau
            else:
                # Fallback: accept smallest step
                with torch.no_grad():
                    self.signal.unconstrained.copy_(backup - dt_actual * grad_u)
        else:
            with torch.no_grad():
                self.signal.unconstrained.copy_(self.signal.unconstrained - dt_actual * grad_u)

        # Record
        state = self.signal.state_vector().detach().cpu().numpy()
        F_val = self._compute_F_from_unconstrained().item()

        # Use PROJECTED gradient norm for convergence (DISCOVERY fix)
        grad_s = self.signal.grad_F().detach().cpu().numpy()
        grad_norm = self.projected_gradient_norm(grad_s)

        self._trajectory.append(state.copy())
        self._F_values.append(F_val)
        self._grad_norms.append(grad_norm)
        self._step_count += 1

        return grad_norm

    def run_to_convergence_lbfgs(self, verbose: bool = False) -> ConvergenceMetrics:
        """Run L-BFGS optimization for rapid convergence.

        L-BFGS uses second-order information (approximate Hessian)
        and converges much faster than gradient descent for smooth
        objectives like F(lambda, kappa, eta).

        Returns:
            ConvergenceMetrics with full diagnostics
        """
        t0 = time.perf_counter()

        optimizer = torch.optim.LBFGS(
            [self.signal.unconstrained],
            lr=1.0,
            max_iter=20,
            tolerance_grad=self.tolerance,
            tolerance_change=1e-15,
            history_size=10,
            line_search_fn='strong_wolfe',
        )

        grad_norm = float('inf')

        for step in range(self.max_steps // 20):  # LBFGS does ~20 evals per step
            def closure():
                optimizer.zero_grad()
                F = self._compute_F_from_unconstrained()
                F.backward()
                return F

            optimizer.step(closure)

            # Record state
            state = self.signal.state_vector().detach().cpu().numpy()
            F_val = self._compute_F_from_unconstrained().item()
            grad_s = self.signal.grad_F().detach().cpu().numpy()
            grad_norm = self.projected_gradient_norm(grad_s)

            self._trajectory.append(state.copy())
            self._F_values.append(F_val)
            self._grad_norms.append(grad_norm)
            self._step_count += 1

            if verbose and self._step_count % 10 == 0:
                print(f"  L-BFGS step {self._step_count}: F={F_val:.8f}, |grad_F|={grad_norm:.2e}")

            if grad_norm < self.tolerance:
                if verbose:
                    print(f"  Converged at L-BFGS step {self._step_count}, |grad_F| = {grad_norm:.2e}")
                break

        wall_time = (time.perf_counter() - t0) * 1000
        return self._build_metrics(grad_norm, wall_time, method='lbfgs_v3')

    def run_to_convergence(self, verbose: bool = False) -> ConvergenceMetrics:
        """Run gradient flow until convergence.

        Returns:
            ConvergenceMetrics with full diagnostics
        """
        t0 = time.perf_counter()

        grad_norm = float('inf')
        for step in range(self.max_steps):
            grad_norm = self.step()

            if grad_norm < self.tolerance:
                if verbose:
                    print(f"  Converged at step {step}, |grad_F| = {grad_norm:.2e}")
                break

            if verbose and step % 1000 == 0 and step > 0:
                print(f"  Step {step}: F={self._F_values[-1]:.6f}, |grad_F|={grad_norm:.2e}")

        wall_time = (time.perf_counter() - t0) * 1000
        return self._build_metrics(grad_norm, wall_time,
                                    method='adaptive_v3' if self.adaptive else 'euler_v3')

    def _build_metrics(self, grad_norm: float, wall_time: float, method: str) -> ConvergenceMetrics:
        final_state = self.signal.to_numpy()
        conservation_err = abs(float(final_state.sum()) - 1.0)
        watkins_dist = float(np.linalg.norm(final_state - np.array(
            [WATKINS_LAMBDA, WATKINS_KAPPA, WATKINS_ETA]
        )))

        subsample_step = max(1, len(self._trajectory) // 100)
        traj_sub = [t.tolist() for t in self._trajectory[::subsample_step]]

        return ConvergenceMetrics(
            converged=grad_norm < self.tolerance,
            steps=self._step_count,
            final_state=final_state.tolist(),
            final_F=self._F_values[-1] if self._F_values else float('inf'),
            gradient_norm=grad_norm,
            conservation_error=conservation_err,
            watkins_distance=watkins_dist,
            at_watkins=watkins_dist < 1e-5,
            wall_time_ms=wall_time,
            method=method,
            trajectory_subsample=traj_sub,
            F_history=self._F_values[-100:],
            grad_history=self._grad_norms[-100:],
        )


def multi_start_convergence_test(
    n_starts: int = 20,
    max_steps: int = 10000,
    tolerance: float = 1e-8,
    verbose: bool = True,
) -> Dict:
    """Run convergence test from multiple random starting points.

    Tests that ALL random initializations converge to the Watkins equilibrium,
    validating that the equilibrium is a global attractor.

    Args:
        n_starts: Number of random starting points
        max_steps: Max steps per run
        tolerance: Convergence tolerance
        verbose: Print progress

    Returns:
        Dictionary with aggregate metrics and per-run results
    """
    results = []
    all_converged = True
    all_watkins = True

    if verbose:
        print(f"Multi-start convergence test: {n_starts} random starts")
        print(f"  Max steps: {max_steps}, tolerance: {tolerance}")
        print()

    for i in range(n_starts):
        # Random point on simplex via Dirichlet(1,1,1)
        r = torch.rand(3, dtype=torch.float64)
        r = r / r.sum()
        signal = AlgoSignal(
            lambda_init=r[0].item(),
            kappa_init=r[1].item(),
            eta_init=r[2].item(),
        )

        engine = SimplexFlowEngineV3(signal, max_steps=max_steps, tolerance=tolerance)
        metrics = engine.run_to_convergence(verbose=False)
        results.append(metrics)

        if not metrics.converged:
            all_converged = False
        if not metrics.at_watkins:
            all_watkins = False

        if verbose:
            status = "OK" if metrics.converged and metrics.at_watkins else "FAIL"
            print(f"  [{status}] Start {i+1:2d}: "
                  f"init=[{r[0]:.3f},{r[1]:.3f},{r[2]:.3f}] -> "
                  f"steps={metrics.steps:5d}, "
                  f"|grad_F|={metrics.gradient_norm:.2e}, "
                  f"d_W={metrics.watkins_distance:.2e}")

    # Aggregate
    steps_list = [r.steps for r in results]
    times_list = [r.wall_time_ms for r in results]
    conv_errors = [r.conservation_error for r in results]

    summary = {
        'n_starts': n_starts,
        'all_converged': all_converged,
        'all_watkins': all_watkins,
        'convergence_rate': sum(1 for r in results if r.converged) / n_starts,
        'watkins_rate': sum(1 for r in results if r.at_watkins) / n_starts,
        'avg_steps': np.mean(steps_list),
        'max_steps_used': max(steps_list),
        'avg_wall_time_ms': np.mean(times_list),
        'max_conservation_error': max(conv_errors),
        'per_run': [r.to_dict() for r in results],
    }

    if verbose:
        print()
        print(f"  Convergence rate: {summary['convergence_rate']*100:.0f}%")
        print(f"  Watkins rate: {summary['watkins_rate']*100:.0f}%")
        print(f"  Avg steps: {summary['avg_steps']:.0f}")
        print(f"  Max conservation error: {summary['max_conservation_error']:.2e}")

    return summary


def compute_F_raw(lam: float, kap: float, eta: float) -> float:
    """Compute F(lambda, kappa, eta) for a raw simplex coordinate triple."""
    signal = AlgoSignal(lambda_init=lam, kappa_init=kap, eta_init=eta)
    return float(signal.compute_F().item())


def _main() -> None:
    """Minimal CLI smoke test for module execution."""
    signal = AlgoSignal(lambda_init=0.4, kappa_init=0.3, eta_init=0.3)
    engine = SimplexFlowEngineV3(signal, max_steps=5000, tolerance=1e-8, adaptive=True)
    metrics = engine.run_to_convergence(verbose=False)
    print(f"Converged: {metrics.converged}")
    print(f"Steps: {metrics.steps}")
    print(f"Final F: {metrics.final_F:.12f}")
    print(f"Gradient norm: {metrics.gradient_norm:.3e}")
    print(f"At Watkins: {metrics.at_watkins}")


if __name__ == "__main__":
    _main()
