#!/usr/bin/env python3
"""
ADP (Asymptotic Derivative Plateau) Theory — MNIST CNN Training
================================================================
Projects full parameter space onto a 2D subspace, logs trajectory,
loss derivative, drift direction, and photosphere crossing.

Usage:
    python train_mnist_adp.py

Logs directory: ./adp_logs/
"""

import os
import json
import time
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

# ─── 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"
LOG_INTERVAL = 4  # log every N iterations per epoch

# Set seeds
torch.manual_seed(SEED)
np.random.seed(SEED)

# ─── 2D Projection Observer ──────────────────────────────────────────
class TwoDObserver:
    """
    Projects the full parameter vector into 2D using random fixed directions.
    Tracks: trajectory, step vectors, speed, heading, curvature proxy.
    """

    def __init__(self, model, rank=12345):
        self.model = model
        self.init_params = None
        self.proj_a = None  # first 2D axis
        self.proj_b = None  # second 2D axis
        self.path = []  # list of (theta1, theta2, loss, time)
        self.dirs = []  # unit direction vectors between consecutive steps

        self._init_projection()

    def _init_projection(self):
        """Random projection of all parameters into 2D."""
        all_params = self._flat_params()
        rng = np.random.RandomState(12345)
        # Two orthogonal random directions
        self.proj_a = rng.randn(len(all_params))
        self.proj_b = rng.randn(len(all_params))
        # Gram-Schmidt orthogonalize
        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):
        """Return (x, y) in 2D space."""
        p = self._flat_params()
        return (np.dot(p, self.proj_a), np.dot(p, self.proj_b))

    def direction_from_init(self):
        """Unit vector from initial point to current 2D position."""
        x, y = self.project()
        dx, dy = x - self.init_coords[0], y - self.init_coords[1]
        norm = np.hypot(dx, dy)
        if norm < 1e-12:
            return (1.0, 0.0)
        return (dx / norm, dy / norm)

    def speed(self):
        """Speed of the most recent step."""
        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 np.hypot(x2 - x1, y2 - y1)

    def heading(self):
        """Heading angle in radians of the most recent step."""
        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 np.arctan2(y2 - y1, x2 - x1)


# ─── 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)


# ─── Compute Loss Derivative ─────────────────────────────────────────
def compute_loss_derivative(loss_history, window=5):
    """
    Return (mean_drift, derivative_variance) over last `window` steps.
    mean_drift ≈ dL/dt (constant if ADP regime).
    """
    if len(loss_history) < window:
        return 0.0, float("inf")
    recent = loss_history[-window:]
    diffs = [recent[i + 1] - recent[i] for i in range(len(recent) - 1)]
    mean = np.mean(diffs)
    var = np.var(diffs)
    return mean, var


# ─── Photosphere Crossing Detector ────────────────────────────────────
def check_photosphere(loss_history, window=10, threshold=0.0003):
    """
    Detect when trajectory exits the opaque interior.
    Returns True if derivative variance has collapsed (low opacity).
    """
    if len(loss_history) < window:
        return False, 0.0
    recent = loss_history[-window:]
    diffs = [recent[i + 1] - recent[i] for i in range(len(recent) - 1)]
    var = np.var(diffs)
    return var < threshold, var


# ─── Training ─────────────────────────────────────────────────────────
def train():
    device = DEVICE
    print(f"[ADP-Train] Device: {device}")
    print(f"[ADP-Train] LR={LR}, BS={BATCH_SIZE}, Epochs={EPOCHS}")
    print(f"[ADP-Train] Log dir: {LOG_DIR}")
    print("=" * 70)

    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)

    # Observer
    observer = TwoDObserver(model)
    print(f"[ADP-Observer] 2D projection initialized.")
    print(f"  Initial (θ₁, θ₂) = ({observer.init_coords[0]:+.6f}, {observer.init_coords[1]:+.6f})")

    # Logging accumulators
    loss_history = []
    steps_since_init = 0
    log_rows = []  # structured log for JSON
    photo_crossed = False

    overall_start = time.time()

    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 += 1
            step_time = time.time()
            epoch_loss_sum += loss.item()
            epoch_batches += 1

            # Record: projected 2D pos, loss, time, direction
            coords = observer.project()
            direction = observer.direction_from_init()
            speed = observer.speed()
            heading = observer.heading()

            loss_history.append(loss.item())

            if steps_since_init % LOG_INTERVAL == 0:
                drift, drift_var = compute_loss_derivative(loss_history)
                crossed, opacity_var = check_photosphere(loss_history)
                # Eta2D direction (direction of 2D step)
                eta2d = f"({direction[0]:+.4f}, {direction[1]:+.4f})"

                entry = {
                    "step": steps_since_init,
                    "epoch": epoch + 1,
                    "batch_in_epoch": batch_idx,
                    "theta1": round(coords[0], 6),
                    "theta2": round(coords[1], 6),
                    "loss": round(loss.item(), 6),
                    "eta2d_x": round(direction[0], 6),
                    "eta2d_y": round(direction[1], 6),
                    "speed": round(speed, 6),
                    "heading_rad": round(heading, 6),
                    "drift_mean": round(drift, 8),
                    "drift_var": round(drift_var, 10),
                    "opacity_var": round(opacity_var, 10),
                    "photosphere_crossed": crossed,
                }
                log_rows.append(entry)

                # Runtime print (every 50 logged steps or epoch boundary)
                if (batch_idx % 50 == 0) or (batch_idx == len(train_loader) - 1):
                    pbar = (epoch * len(train_loader) + batch_idx) / (EPOCHS * len(train_loader))
                    state = "✧ DRIFTING" if crossed else "↻ SCATTERING"
                    print(
                        f"  Epoch {epoch+1:2d} | Batch {batch_idx:4d}/{len(train_loader)} "
                        f"| Loss {loss.item():.4f} | Drift {drift:+.6f} Var {drift_var:.2e} | "
                        f"Dir {eta2d} | Speed {speed:.4f} | Heading {heading:.2f}rad | {state}"
                    )

        # End of epoch
        epoch_avg = epoch_loss_sum / max(epoch_batches, 1)
        drift_full, drift_var_full = compute_loss_derivative(loss_history)
        crossed, _ = check_photosphere(loss_history)
        if crossed and not photo_crossed:
            photo_crossed = True
            print(f"\n[✧ PHOTOsphere CROSSED at epoch {epoch+1}, step {steps_since_init}]")
            print(f"  Drift variance collapsed: {drift_var_full:.2e}")
            print(f"  Mean loss derivative: {drift_full:+.6f}")

        print(
            f"  Epoch {epoch+1}  avg_loss={epoch_avg:.4f}  "
            f"full_drift_mean={drift_full:+.6f}  drift_var={drift_var_full:.2e}  "
            f"{'⟶ ADP' if photo_crossed else '… scattering'}"
        )

    total_time = time.time() - overall_start
    print(f"\n[ADP-Train] Training complete in {total_time:.1f}s. {len(log_rows)} log rows.")
    print(f"[ADP-Train] Total steps: {steps_since_init}")
    print(f"[ADP-Train] Photosphere crossed: {photo_crossed}")

    # ─── Save Logs ──────────────────────────────────────────────────
    # 1. Full structured JSON
    log_file = os.path.join(LOG_DIR, "training_logs.json")
    with open(log_file, "w") as f:
        json.dump(log_rows, f, indent=2)
    print(f"[ADP-Train] Structured logs → {log_file}")

    # 2. CSV-style for plotting
    csv_file = os.path.join(LOG_DIR, "training.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-Train] CSV logs → {csv_file}")

    # 3. Summary JSON
    final_drift, final_drift_var = compute_loss_derivative(loss_history, window=20)
    summary = {
        "experiment": "mnist_adp",
        "epochs": EPOCHS,
        "lr": LR,
        "batch_size": BATCH_SIZE,
        "device": str(device),
        "total_steps": steps_since_init,
        "total_time_s": round(total_time, 2),
        "final_loss": round(loss_history[-1], 6) if loss_history else None,
        "final_drift_mean": round(final_drift, 8),
        "final_drift_var": round(final_drift_var, 10),
        "photosphere_crossed_epoch": epoch + 1 if photo_crossed else None,
        "log_rows": len(log_rows),
        "init_theta": [round(observer.init_coords[0], 6), round(observer.init_coords[1], 6)],
        "final_theta": [round(observer.project()[0], 6), round(observer.project()[1], 6)],
    }
    summary_file = os.path.join(LOG_DIR, "summary.json")
    with open(summary_file, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"[ADP-Train] Summary → {summary_file}")

    return model, log_rows, summary


# ─── 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, summary = train()
    acc = evaluate(model)
    summary["test_accuracy"] = round(acc, 4)

    # Append test accuracy
    with open(os.path.join(LOG_DIR, "summary.json"), "w") as f:
        json.dump(summary, f, indent=2)
    print(f"\n[ADP-Train] Test Accuracy: {acc:.4f} ({acc*100:.1f}%)")
    print(f"[ADP-Train] All logs in: {LOG_DIR}/")