import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import numpy as np
from torch.utils.data import DataLoader
from sklearn.metrics import confusion_matrix

torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

EPOCHS = 60          # OneCycle is tuned to this horizon; change both together
MAX_LR = 3e-3        # peak learning rate for the OneCycle schedule
LAMBDA_SPECTRAL = 0.01
LAMBDA_COMMUTATOR = 0.1   # strength of the inter-layer commutator penalty

# -------------------------------
# 1. 3x3 Algebraic Cell (vectorized spectral reg)
# -------------------------------
class AlgebraicCell3x3(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_channels, in_channels, 3, 3) * 0.1)
        self.bias = nn.Parameter(torch.zeros(out_channels))

    def forward(self, x):
        return F.conv2d(x, self.weight, bias=self.bias, padding=1)

    def mean_kernel(self):
        """Layer's representative 3x3 operator (mean over all channels)."""
        return self.weight.mean(dim=(0, 1))

    def spectral_regularization(self):
        mats = self.weight.mean(dim=1)
        sym = mats + mats.transpose(-1, -2)
        eigvals = torch.linalg.eigvalsh(sym)
        psd_pen = F.relu(-eigvals[:, 0]).sum() * 0.1
        gaps = eigvals[:, 1:] - eigvals[:, :-1]
        collision_pen = torch.exp(-gaps.var(dim=1, unbiased=True)).sum() * 0.01
        return psd_pen + collision_pen

# -------------------------------
# 2. Network — now with BatchNorm (the single biggest accuracy accelerator)
# -------------------------------
class InvariantCellNet(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = AlgebraicCell3x3(3, 64)
        self.bn1 = nn.BatchNorm2d(64)
        self.conv2 = AlgebraicCell3x3(64, 128)
        self.bn2 = nn.BatchNorm2d(128)
        self.conv3 = AlgebraicCell3x3(128, 256)
        self.bn3 = nn.BatchNorm2d(256)
        self.pool = nn.MaxPool2d(2)
        self.drop = nn.Dropout(0.3)
        self.fc1 = nn.Linear(256 * 4 * 4, 512)
        self.fc2 = nn.Linear(512, num_classes)

    def forward(self, x):
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        x = self.pool(F.relu(self.bn3(self.conv3(x))))
        x = x.flatten(1)
        x = F.relu(self.fc1(x))
        x = self.drop(x)
        return self.fc2(x)

    def spectral_reg_loss(self):
        return (self.conv1.spectral_regularization() +
                self.conv2.spectral_regularization() +
                self.conv3.spectral_regularization())

    def commutator_reg_loss(self):
        """Lie-flow orthogonality between successive cells.

        Penalizes ||[M1,M2]||_F^2 + ||[M2,M3]||_F^2 on the layers' mean
        3x3 kernels. Driving the commutator toward the zero MATRIX (a
        fine-grained zero, not a scalar zero) pushes successive feature
        extractors toward order-independence (zero phase-shift).
        """
        m1, m2, m3 = (self.conv1.mean_kernel(),
                      self.conv2.mean_kernel(),
                      self.conv3.mean_kernel())
        c12 = m1 @ m2 - m2 @ m1
        c23 = m2 @ m3 - m3 @ m2
        return c12.pow(2).sum() + c23.pow(2).sum()

    @torch.no_grad()
    def commutator_norms(self):
        """Diagnostic: Frobenius norms of the inter-layer commutators."""
        m1, m2, m3 = (self.conv1.mean_kernel(),
                      self.conv2.mean_kernel(),
                      self.conv3.mean_kernel())
        return (torch.linalg.matrix_norm(m1 @ m2 - m2 @ m1).item(),
                torch.linalg.matrix_norm(m2 @ m3 - m3 @ m2).item())

# -------------------------------
# 3. Class Equity Monitor (detection + reweighting, as before)
# -------------------------------
class ClassEquityMonitor:
    def __init__(self, num_classes, device, ema=0.99, eta=1.0, warn_z=1.5):
        self.num_classes = num_classes
        self.ema = ema
        self.eta = eta
        self.warn_z = warn_z
        self.loss_ema = torch.zeros(num_classes, device=device)
        self.seen = torch.zeros(num_classes, device=device)

    @torch.no_grad()
    def update(self, per_sample_loss, targets):
        sums = torch.zeros(self.num_classes, device=per_sample_loss.device)
        cnts = torch.zeros(self.num_classes, device=per_sample_loss.device)
        sums.scatter_add_(0, targets, per_sample_loss)
        cnts.scatter_add_(0, targets, torch.ones_like(per_sample_loss))
        present = cnts > 0
        batch_mean = sums[present] / cnts[present]
        first_time = self.seen[present] == 0
        new_vals = self.ema * self.loss_ema[present] + (1 - self.ema) * batch_mean
        self.loss_ema[present] = torch.where(first_time, batch_mean, new_vals)
        self.seen += cnts

    @torch.no_grad()
    def weights(self):
        return torch.softmax(self.eta * self.loss_ema, dim=0) * self.num_classes

    @torch.no_grad()
    def sacrificed_classes(self):
        if (self.seen > 0).sum() < self.num_classes:
            return torch.tensor([], dtype=torch.long)
        mu = self.loss_ema.mean()
        sd = self.loss_ema.std().clamp_min(1e-8)
        z = (self.loss_ema - mu) / sd
        return torch.nonzero(z > self.warn_z).flatten().cpu()

    def report(self, class_names=None):
        flagged = self.sacrificed_classes()
        ema = self.loss_ema.cpu().numpy()
        print(f"  Per-class EMA loss: {np.array2string(ema, precision=3)}")
        if len(flagged) > 0:
            names = [class_names[i] if class_names else str(int(i)) for i in flagged]
            print(f"  !! SACRIFICE WARNING: {names} — escalating to gradient surgery.")
        else:
            print("  No class sacrifice detected (reweighting only).")

# -------------------------------
# 4. Equity Gradient Surgery (interference-free circuitry, PCGrad-style)
# -------------------------------
def equity_gradient_surgery(model, scaler, loss_lag, loss_dom):
    """When sacrifice is flagged, escalate from reweighting to projection.

    Computes separate gradients for the lagging-class loss and the
    dominant-class loss. If they conflict (negative inner product across
    the full parameter space), the dominant gradient's conflicting
    component is projected off the lagging direction:

        g_dom <- g_dom - (<g_dom, g_lag> / ||g_lag||^2) g_lag   if <g_dom,g_lag> < 0

    Dominant classes may then only improve along directions ORTHOGONAL to
    the lagging classes' descent direction — the literal zero-inner-product
    condition: structural orthogonality enforced on gradient flow.
    Uniform AMP loss scaling cancels in the projection coefficient.
    """
    params = [p for p in model.parameters() if p.requires_grad]

    # gradient of the lagging-class loss
    scaler.scale(loss_lag).backward(retain_graph=True)
    g_lag = [p.grad.detach().clone() if p.grad is not None else None for p in params]
    for p in params:
        if p.grad is not None:
            p.grad = None

    # gradient of the dominant-class loss (+ any reg already folded in)
    scaler.scale(loss_dom).backward()
    g_dom = [p.grad.detach().clone() if p.grad is not None else None for p in params]

    # conflict test over the whole parameter vector
    dot = sum((gd * gl).sum() for gd, gl in zip(g_dom, g_lag)
              if gd is not None and gl is not None)
    if dot < 0:
        norm2 = sum((gl * gl).sum() for gl in g_lag if gl is not None).clamp_min(1e-12)
        coef = dot / norm2
        g_dom = [gd - coef * gl if (gd is not None and gl is not None) else gd
                 for gd, gl in zip(g_dom, g_lag)]

    # final gradient: protected lagging direction + de-conflicted dominant
    for p, gd, gl in zip(params, g_dom, g_lag):
        if gd is None and gl is None:
            continue
        p.grad = (gd if gd is not None else 0) + (gl if gl is not None else 0)

# -------------------------------
# 5. Accuracy Acceleration Meter
# -------------------------------
class AccelerationMeter:
    """Tracks test-accuracy velocity (1st difference) and acceleration
    (2nd difference) so 'training force' is a measured quantity."""
    def __init__(self):
        self.acc, self.vel = [], []

    def update(self, acc):
        self.acc.append(acc)
        v = self.acc[-1] - self.acc[-2] if len(self.acc) > 1 else 0.0
        self.vel.append(v)
        a = self.vel[-1] - self.vel[-2] if len(self.vel) > 1 else 0.0
        print(f"  [accel-meter] acc={acc:.2f}%  velocity={v:+.2f}%/ep  "
              f"acceleration={a:+.2f}%/ep²")
        return v, a

# -------------------------------
# 6. Diagnostics
# -------------------------------
def detect_eigenvalue_collision(model):
    model.eval()
    with torch.no_grad():
        net = getattr(model, '_orig_mod', model)
        w = net.conv1.weight[0].mean(dim=0).cpu().numpy()
        eigvals = np.linalg.eigvalsh(w + w.T)
        gap_var = np.var(np.diff(eigvals))
        print(f"Eigenvalue collision score: {gap_var:.4f} -> "
              f"{'COLLAPSED' if gap_var < 0.01 else 'OK'}")
        return gap_var < 0.01

def simpson_paradox_check(model, test_loader, device):
    model.eval()
    all_preds, all_labels = [], []
    with torch.no_grad():
        for images, labels in test_loader:
            images = images.to(device, non_blocking=True)
            with torch.autocast(device_type=device.type, enabled=(device.type == 'cuda')):
                outputs = model(images)
            all_preds.append(outputs.argmax(1).cpu())
            all_labels.append(labels)
    all_preds = torch.cat(all_preds).numpy()
    all_labels = torch.cat(all_labels).numpy()
    cm = confusion_matrix(all_labels, all_preds)
    per_class_acc = cm.diagonal() / cm.sum(axis=1)
    overall_acc = np.mean(per_class_acc)
    low = np.where(per_class_acc < overall_acc - 0.1)[0]
    if len(low) > 0:
        print(f"Simpson's paradox warning: classes {low} at "
              f"{per_class_acc[low]} vs overall {overall_acc:.3f}")
    return overall_acc, per_class_acc

# -------------------------------
# 7. Training loop
# -------------------------------
def train(model, device, train_loader, optimizer, scheduler, scaler, epoch, monitor,
          lambda_spectral=LAMBDA_SPECTRAL, lambda_comm=LAMBDA_COMMUTATOR,
          spectral_every=4, min_lag_samples=8):
    model.train()
    total_loss, surgeries = 0.0, 0
    correct = torch.zeros(1, device=device)
    use_amp = device.type == 'cuda'

    for batch_idx, (data, target) in enumerate(train_loader):
        data = data.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)
        optimizer.zero_grad(set_to_none=True)

        class_w = monitor.weights()
        flagged = monitor.sacrificed_classes().to(device)

        with torch.autocast(device_type=device.type, enabled=use_amp):
            output = model(data)
            per_sample = F.cross_entropy(output, target, reduction='none',
                                         label_smoothing=0.1)
        monitor.update(per_sample.detach().float(), target)

        # regularizers (fp32, on tiny matrices)
        reg = torch.zeros((), device=device)
        if batch_idx % spectral_every == 0:
            reg = reg + lambda_spectral * model.spectral_reg_loss() * spectral_every
        reg = reg + lambda_comm * model.commutator_reg_loss()

        lag_mask = torch.isin(target, flagged) if flagged.numel() > 0 else None
        if (lag_mask is not None and lag_mask.sum() >= min_lag_samples
                and (~lag_mask).sum() >= min_lag_samples):
            # ESCALATION: sacrifice in progress -> gradient surgery
            loss_lag = (class_w[target[lag_mask]] * per_sample[lag_mask]).mean()
            loss_dom = (class_w[target[~lag_mask]] * per_sample[~lag_mask]).mean() + reg
            equity_gradient_surgery(model, scaler, loss_lag, loss_dom)
            loss_val = (loss_lag + loss_dom).item()
            surgeries += 1
        else:
            # normal regime: equity reweighting only
            loss = (class_w[target] * per_sample).mean() + reg
            scaler.scale(loss).backward()
            loss_val = loss.item()

        scaler.step(optimizer)
        scaler.update()
        scheduler.step()

        total_loss += loss_val
        correct += (output.argmax(dim=1) == target).sum()

    avg_loss = total_loss / len(train_loader)
    acc = 100. * correct.item() / len(train_loader.dataset)
    extra = f", surgeries: {surgeries}/{len(train_loader)}" if surgeries else ""
    print(f'Train Epoch {epoch}: Loss: {avg_loss:.4f}, Accuracy: {acc:.2f}%, '
          f'LR: {scheduler.get_last_lr()[0]:.4f}{extra}')
    return avg_loss, acc

def test(model, device, test_loader):
    model.eval()
    test_loss = 0.0
    correct = torch.zeros(1, device=device)
    use_amp = device.type == 'cuda'
    with torch.no_grad():
        for data, target in test_loader:
            data = data.to(device, non_blocking=True)
            target = target.to(device, non_blocking=True)
            with torch.autocast(device_type=device.type, enabled=use_amp):
                output = model(data)
                test_loss += F.cross_entropy(output, target, reduction='sum').item()
            correct += (output.argmax(dim=1) == target).sum()
    test_loss /= len(test_loader.dataset)
    acc = 100. * correct.item() / len(test_loader.dataset)
    print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {acc:.2f}%')
    return acc

# -------------------------------
# 8. Main
# -------------------------------
def main():
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])

    trainset = torchvision.datasets.CIFAR10(root='../data', train=True, download=True,
                                            transform=transform_train)
    testset = torchvision.datasets.CIFAR10(root='../data', train=False, download=True,
                                           transform=transform_test)
    class_names = trainset.classes

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    pin = device.type == 'cuda'
    train_loader = DataLoader(trainset, batch_size=256, shuffle=True,
                              num_workers=4, pin_memory=pin, persistent_workers=True)
    test_loader = DataLoader(testset, batch_size=512, shuffle=False,
                             num_workers=4, pin_memory=pin, persistent_workers=True)

    model = InvariantCellNet(num_classes=10).to(device)
    try:
        model = torch.compile(model)
    except Exception:
        pass

    optimizer = optim.AdamW(model.parameters(), lr=MAX_LR / 10, weight_decay=5e-4)
    scheduler = optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=MAX_LR,
        steps_per_epoch=len(train_loader), epochs=EPOCHS,
        pct_start=0.15, div_factor=10, final_div_factor=100)
    scaler = torch.amp.GradScaler(enabled=(device.type == 'cuda'))
    monitor = ClassEquityMonitor(num_classes=10, device=device,
                                 ema=0.99, eta=1.0, warn_z=1.5)
    accel = AccelerationMeter()

    print("=== METAFOUNDRY-ML Pre-Run Diagnosis ===")
    detect_eigenvalue_collision(model)
    c12, c23 = getattr(model, '_orig_mod', model).commutator_norms()
    print(f"Initial commutator norms: ||[M1,M2]||={c12:.4f}, ||[M2,M3]||={c23:.4f}")

    best_acc = 0.0
    for epoch in range(1, EPOCHS + 1):
        train(model, device, train_loader, optimizer, scheduler, scaler, epoch, monitor)
        test_acc = test(model, device, test_loader)
        accel.update(test_acc)
        monitor.report(class_names)
        if test_acc > best_acc:
            best_acc = test_acc
            torch.save(model.state_dict(), 'best_invariant_cellnet.pth')

        if epoch % 10 == 0:
            print("\n--- Self-Diagnostic Report ---")
            detect_eigenvalue_collision(model)
            c12, c23 = getattr(model, '_orig_mod', model).commutator_norms()
            print(f"Commutator norms: ||[M1,M2]||={c12:.4f}, ||[M2,M3]||={c23:.4f} "
                  f"(falling = Lie-flow orthogonality emerging)")
            overall, per_class = simpson_paradox_check(model, test_loader, device)
            print(f"Overall test accuracy: {overall*100:.2f}%")
            print(f"Per-class accuracies: {per_class}")
            print("-------------------------------\n")

    model.load_state_dict(torch.load('best_invariant_cellnet.pth'))
    test(model, device, test_loader)
    print(f"\n=== Final Result ===")
    print(f"Best test accuracy: {best_acc:.2f}%")
    c12, c23 = getattr(model, '_orig_mod', model).commutator_norms()
    print(f"Final commutator norms: ||[M1,M2]||={c12:.4f}, ||[M2,M3]||={c23:.4f}")

if __name__ == "__main__":
    main()