#!/usr/bin/env python3
"""
ADP Theory v2.0 — Asymptotic Derivative Plateau
=================================================
MNIST MLP training with stochastic photosphere detection,
EWMA-based statistics, and Type I/II regime classification.

Key v2.0 Enhancements:
- EWMA (Exponential Moving Average) smoothing
- Continuous opacity field τ (no binary photosphere)
- Signal-to-Noise Ratio (SNR) tracking
- Type I (noisy attractor) vs Type II (manifold drift) detection
- 2D diffusion radius monitoring
- Curvature index for regime identification

Usage:
    python train_mnist_adp_v2.py
"""

import os
import json
import time
import math
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from collections import deque


# ─── Configuration ────────────────────────────────────────────────────
LR = 0.01
BATCH_SIZE = 128
EPOCHS = 8
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
SEED = 42
LOG_DIR = "./adp_logs_v2"
LOG_INTERVAL = 4
EWMA_ALPHA = 0.1  # smoothing factor for EWMA
DIFFUSION_WINDOW = 200


# ─── ADP State Manager (EWMA-based) ──────────────────────────────────
class ADPState:
    """
    Maintains exponentially-weighted moving averages of loss derivatives.
    Replaces fixed-window statistics with adaptive smoothing.
    """
    def __init__(self, alpha=EWMA_ALPHA):
        self.alpha = alpha
        self.drift_mean = 0.0      # EWMA of dL/dt
        self.drift_var = 0.0       # EWMA of (dL/dt - mean)^2
        self.ema_loss = None       # Exponentially smoothed loss
        self.loss_history = deque(maxlen=500)  # Recent losses for diffusion tracking
        self.step_count = 0
        
    def update(self, loss):
        """Update state with new loss observation."""
        self.step_count += 1
        self.loss_history.append(loss)
        
        if self.ema_loss is None:
            self.ema_loss = loss
            
        delta = loss - self.ema_loss
        self.ema_loss += self.alpha * delta
        
        # Online EWMA for mean and variance
        self.drift_mean += self.alpha * (delta - self.drift_mean)
        self.drift_var += self.alpha * ((delta - self.drift_mean)**2 - self.drift_var)
        
    def get_snr(self):
        """Signal-to-Noise Ratio of the loss derivative."""
        std = math.sqrt(max(self.drift_var, 1e-12))
        return abs(self.drift_mean) / std
    
    def get_opacity(self):
        """
        Continuous opacity field τ.
        τ = SNR² where high τ = opaque (core), low τ = transparent (photosphere/corona).
        """
        snr = self.get_snr()
        return snr ** 2
    
    def get_transparency(self):
        """Transparency score T = exp(-τ)."""
        tau = self.get_opacity()
        return math.exp(-tau)
    
    def is_in_equilibrium(self):
        """Check if system is in stochastic equilibrium (SNR < 1)."""
        return self.get_snr() < 1.0
    
    def is_in_photosphere(self):
        """Check if system is in transition layer (0.5 < SNR <= 2.0)."""
        snr = self.get_snr()
        return 0.5 < snr <= 2.0


# ─── 2D Observer with Diffusion Tracking ─────────────────────────────
class TwoDObserver:
    """
    Projects full parameter space onto 2D subspace.
    Tracks trajectory, diffusion radius, and curvature index.
    """
    def __init__(self, model, rank=12345):
        self.model = model
        self.init_params = None
        self.proj_a = None
        self.proj_b = None
        self.path = []  # List of (theta1, theta2, loss, timestamp)
        self.init_coords = None
        
        # Diffusion tracking
        self.diffusion_window = deque(maxlen=DIFFUSION_WINDOW)
        self.diffusion_radius_sq = 0.0
        self.prev_diffusion_radius_sq = 0.0
        
        self._init_projection()
    
    def _init_projection(self):
        """Random orthogonal projection into 2D."""
        all_params = self._flat_params()
        rng = np.random.RandomState(12345)
        self.proj_a = rng.randn(len(all_params))
        self.proj_b = rng.randn(len(all_params))
        
        # Gram-Schmidt orthogonalization
        self.proj_b -= np.dot(self.proj_b, self.proj_a) / np.dot(self.proj_a, self.proj_a) * self.proj_a
        self.proj_a /= np.linalg.norm(self.proj_a)
        self.proj_b /= np.linalg.norm(self.proj_b)
        
        self.init_params = all_params.copy()
        self.init_coords = (
            np.dot(all_params, self.proj_a),
            np.dot(all_params, self.proj_b),
        )
    
    def _flat_params(self):
        return torch.cat([p.view(-1) for p in self.model.parameters()]).detach().cpu().numpy()
    
    def project(self, loss=None, timestamp=None):
        """Project current params to 2D and record trajectory point."""
        p = self._flat_params()
        x = float(np.dot(p, self.proj_a))
        y = float(np.dot(p, self.proj_b))
        self.path.append((x, y, loss, timestamp))
        
        # Update diffusion window
        self.diffusion_window.append((x, y))
        
        # Compute diffusion radius
        if len(self.diffusion_window) > 10:
            points = list(self.diffusion_window)
            cx = sum(p[0] for p in points) / len(points)
            cy = sum(p[1] for p in points) / len(points)
            self.diffusion_radius_sq = sum((p[0]-cx)**2 + (p[1]-cy)**2 for p in points) / len(points)
        
        return (x, y)
    
    def get_diffusion_growth_rate(self):
        """
        Rate of change of diffusion radius.
        Positive = unbounded drift (Type II), negative/zero = bounded (Type I).
        """
        growth = self.diffusion_radius_sq - self.prev_diffusion_radius_sq
        self.prev_diffusion_radius_sq = self.diffusion_radius_sq
        return growth
    
    def get_curvature_index(self):
        """
        Curvature index κ = ||Δv̂|| / ||v||
        High κ = jittery random walk (diffusion), Low κ = smooth trajectory.
        """
        if len(self.path) < 3:
            return 1.0
        
        # Recent velocity vectors
        v1 = (self.path[-2][0] - self.path[-3][0], self.path[-2][1] - self.path[-3][1])
        v2 = (self.path[-1][0] - self.path[-2][0], self.path[-1][1] - self.path[-2][1])
        
        # Unit vectors
        n1 = math.hypot(v1[0], v1[1])
        n2 = math.hypot(v2[0], v2[1])
        
        if n1 < 1e-12 or n2 < 1e-12:
            return 1.0
        
        u1 = (v1[0]/n1, v1[1]/n1)
        u2 = (v2[0]/n2, v2[1]/n2)
        
        # Change in direction
        delta_u = math.hypot(u2[0] - u1[0], u2[1] - u1[1])
        avg_speed = (n1 + n2) / 2.0
        
        return delta_u / (avg_speed + 1e-12)
    
    def get_heading(self):
        """Current heading angle in radians."""
        if len(self.path) < 2:
            return 0.0
        x2, y2 = self.path[-1][0], self.path[-1][1]
        x1, y1 = self.path[-2][0], self.path[-2][1]
        return float(np.arctan2(y2 - y1, x2 - x1))
    
    def get_speed(self):
        """Current speed (step magnitude)."""
        if len(self.path) < 2:
            return 0.0
        x2, y2 = self.path[-1][0], self.path[-1][1]
        x1, y1 = self.path[-2][0], self.path[-2][1]
        return float(np.hypot(x2 - x1, y2 - y1))
    
    def get_direction_from_init(self):
        """Unit vector from initial position to current position."""
        if len(self.path) == 0:
            return (1.0, 0.0)
        x, y = self.path[-1][0], self.path[-1][1]
        dx = x - self.init_coords[0]
        dy = y - self.init_coords[1]
        norm = math.hypot(dx, dy)
        if norm < 1e-12:
            return (1.0, 0.0)
        return (dx / norm, dy / norm)


# ─── Regime Classifier ────────────────────────────────────────────────
def classify_regime(snr, diffusion_growth, curvature_index):
    """
    Classify current optimization regime based on v2.0 theory.
    
    Returns:
        str: Regime identifier
    """
    if snr > 2.0:
        return "OPAQUE_CORE"
    elif 0.5 < snr <= 2.0:
        return "PHOTOSPHERE"
    elif snr <= 0.5:
        if diffusion_growth > 0.01:
            return "TYPE_II_DRIFT"
        else:
            return "TYPE_I_EQUILIBRIUM"
    return "UNKNOWN"


# ─── Solar Analogy Label ─────────────────────────────────────────────
def get_solar_layer(opacity_tau):
    """
    Map opacity to solar structure analogy.
    """
    if opacity_tau > 10:
        return "Radiative Core"
    elif opacity_tau > 1:
        return "Tachocline"
    elif opacity_tau > 0.1:
        return "Photosphere"
    else:
        return "Chromosphere/Corona"


# ─── MLP Model ────────────────────────────────────────────────────────
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 10),
        )
    
    def forward(self, x):
        return self.net(x)


# ─── Training Loop ────────────────────────────────────────────────────
def train():
    device = DEVICE
    print(f"[ADP-v2.0] Device: {device}")
    print(f"[ADP-v2.0] LR={LR}, BS={BATCH_SIZE}, Epochs={EPOCHS}")
    print(f"[ADP-v2.0] EWMA alpha={EWMA_ALPHA}")
    print(f"[ADP-v2.0] Log dir: {LOG_DIR}")
    print("=" * 80)
    
    os.makedirs(LOG_DIR, exist_ok=True)
    
    # Data
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])
    train_ds = datasets.MNIST("./data", train=True, download=True, transform=transform)
    test_ds = datasets.MNIST("./data", train=False, transform=transform)
    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)
    test_loader = DataLoader(test_ds, batch_size=256, shuffle=False)
    
    # Model
    model = MLP().to(device)
    optimizer = optim.SGD(model.parameters(), lr=LR, momentum=0.9)
    
    # ADP observers
    adp_state = ADPState(alpha=EWMA_ALPHA)
    observer = TwoDObserver(model)
    
    print(f"[ADP-v2.0] 2D projection initialized")
    print(f"  Initial (θ₁, θ₂) = ({observer.init_coords[0]:+.6f}, {observer.init_coords[1]:+.6f})")
    print()
    
    # Logging accumulators
    log_rows = []
    overall_start = time.time()
    last_diffusion_growth = 0.0
    
    for epoch in range(EPOCHS):
        model.train()
        epoch_loss_sum = 0.0
        epoch_batches = 0
        
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            
            optimizer.zero_grad()
            output = model(data)
            loss = nn.functional.cross_entropy(output, target)
            loss.backward()
            optimizer.step()
            
            steps_since_init = len(adp_state.loss_history)
            step_time = time.time()
            epoch_loss_sum += loss.item()
            epoch_batches += 1
            
            # Update ADP state
            adp_state.update(loss.item())
            coords = observer.project(loss=loss.item(), timestamp=step_time)
            
            # Compute metrics
            snr = adp_state.get_snr()
            opacity_tau = adp_state.get_opacity()
            transparency = adp_state.get_transparency()
            regime = classify_regime(snr, last_diffusion_growth, observer.get_curvature_index())
            solar_layer = get_solar_layer(opacity_tau)
            diffusion_growth = observer.get_diffusion_growth_rate()
            curvature = observer.get_curvature_index()
            heading = observer.get_heading()
            speed = observer.get_speed()
            direction = observer.get_direction_from_init()
            
            log_entry = {
                "step": steps_since_init,
                "epoch": epoch + 1,
                "batch_in_epoch": batch_idx,
                "loss": round(loss.item(), 6),
                "ema_loss": round(adp_state.ema_loss, 6) if adp_state.ema_loss else None,
                "theta1": round(coords[0], 6),
                "theta2": round(coords[1], 6),
                "snr": round(snr, 4),
                "opacity_tau": round(opacity_tau, 4),
                "transparency": round(transparency, 4),
                "regime": regime,
                "solar_layer": solar_layer,
                "diffusion_growth": round(diffusion_growth, 8),
                "diffusion_radius_sq": round(observer.diffusion_radius_sq, 8),
                "curvature_index": round(curvature, 4),
                "heading_rad": round(heading, 4),
                "speed": round(speed, 6),
                "eta2d_x": round(direction[0], 6),
                "eta2d_y": round(direction[1], 6),
            }
            log_rows.append(log_entry)
            
            # Runtime print (every 50 steps or epoch end)
            if (batch_idx % 50 == 0) or (batch_idx == len(train_loader) - 1):
                pbar = (epoch * len(train_loader) + batch_idx) / (EPOCHS * len(train_loader))
                regime_emoji = {
                    "OPAQUE_CORE": "🔥",
                    "PHOTOSPHERE": "🌅",
                    "TYPE_I_EQUILIBRIUM": "⚖️",
                    "TYPE_II_DRIFT": "↗️",
                }.get(regime, "❓")
                
                print(
                    f"  Epoch {epoch+1:2d} | Batch {batch_idx:4d}/{len(train_loader):4d} "
                    f"| Loss {loss.item():.4f} | SNR {snr:.3f} | τ {opacity_tau:.2f} "
                    f"| {solar_layer:15s} {regime_emoji:4s} | Speed {speed:.2e} | Curv {curvature:.2f}"
                )
            
            last_diffusion_growth = diffusion_growth
        
        # End of epoch summary
        epoch_avg = epoch_loss_sum / max(epoch_batches, 1)
        final_snr = adp_state.get_snr()
        final_tau = adp_state.get_opacity()
        final_transparency = adp_state.get_transparency()
        final_regime = classify_regime(final_snr, 0, 0)
        final_solar = get_solar_layer(final_tau)
        
        print(
            f"\n  Epoch {epoch+1} Summary: avg_loss={epoch_avg:.4f} | "
            f"SNR={final_snr:.3f} | τ={final_tau:.2f} | "
            f"T={final_transparency:.3f} | Regime: {final_regime} | "
            f"Layer: {final_solar}"
        )
        
        # Detection messages
        if final_snr < 1.0 and epoch >= 2:
            print(f"  [✓] STOCHASTIC EQUILIBRIUM: SNR < 1, training plateau reached")
            if final_regime == "TYPE_I_EQUILIBRIUM":
                print(f"  [✓] TYPE I: Bounded noisy attractor, model has converged")
            elif final_regime == "TYPE_II_DRIFT":
                print(f"  [⚠] TYPE II: Unbounded manifold drift, check for implicit bias")
    
    total_time = time.time() - overall_start
    print(f"\n[ADP-v2.0] Training complete in {total_time:.1f}s")
    print(f"[ADP-v2.0] Total steps: {len(log_rows)}")
    
    # Save logs
    _save_logs(log_rows, LOG_DIR)
    
    return model, log_rows


def _save_logs(log_rows, log_dir):
    """Save training logs to JSON and CSV."""
    # JSON
    json_file = os.path.join(log_dir, "training_logs_v2.json")
    with open(json_file, "w") as f:
        json.dump(log_rows, f, indent=2, allow_nan=True)
    print(f"[ADP-v2.0] Structured logs → {json_file}")
    
    # CSV
    csv_file = os.path.join(log_dir, "training_v2.csv")
    if log_rows:
        header = ",".join(log_rows[0].keys())
        with open(csv_file, "w") as f:
            f.write(header + "\n")
            for row in log_rows:
                vals = [str(row[k]) for k in row]
                f.write(",".join(vals) + "\n")
    print(f"[ADP-v2.0] CSV logs → {csv_file}")


# ─── Test Accuracy ────────────────────────────────────────────────────
def evaluate(model):
    model.eval()
    correct = 0
    total = 0
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])
    test_ds = datasets.MNIST("./data", train=False, transform=transform)
    loader = DataLoader(test_ds, batch_size=256, shuffle=False)
    with torch.no_grad():
        for data, target in loader:
            data, target = data.to(DEVICE), target.to(DEVICE)
            output = model(data)
            pred = output.argmax(dim=1)
            correct += (pred == target).sum().item()
            total += target.size(0)
    return correct / total


# ─── Main ─────────────────────────────────────────────────────────────
if __name__ == "__main__":
    model, logs = train()
    acc = evaluate(model)
    print(f"\n[ADP-v2.0] Test Accuracy: {acc:.4f} ({acc*100:.1f}%)")
    print(f"[ADP-v2.0] All logs in: {LOG_DIR}/")