#!/usr/bin/env python3
"""Stage-2 benchmark for the Markov Spectral Equivariant Layer idea.
Uses the read-only bench vision track and bench.train_model for both systems.
"""
import json, random, sys
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
import bench
from bench.protocol import evaluate

SEEDS = tuple(range(8))
# Union of baseline and idea learning-rate grids; identical on both sides.
GRID = [{"lr": 0.001}, {"lr": 0.003}, {"lr": 0.01}]
EPOCHS = 5
NTRAIN, NTEST = 400, 400


def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


class FejerMarkov2d(nn.Module):
    """Positive unit-mass finite-rank analogue: outer-product Fejer kernel.
    The [1,2,1]/4 kernel is nonnegative, sums to one, and is a finite
    Fourier taper (the 2-D spectrum is the product of triangular tapers).
    """
    def __init__(self, channels, gamma=0.7):
        super().__init__()
        self.gamma = float(gamma)
        k1 = torch.tensor([1., 2., 1.]) / 4.
        kernel = torch.outer(k1, k1).view(1, 1, 3, 3)
        self.register_buffer("kernel", kernel.repeat(channels, 1, 1, 1))
        self.channels = channels

    def forward(self, x):
        smoothed = nn.functional.conv2d(x, self.kernel, padding=1,
                                        groups=self.channels)
        return self.gamma * smoothed + (1.0 - self.gamma) * x


class MarkovCNN(nn.Module):
    def __init__(self, out_dim=10, gamma=0.7):
        super().__init__()
        self.c1 = nn.Conv2d(3, 32, 3, padding=1)
        self.c2 = nn.Conv2d(32, 64, 3, padding=1)
        self.c3 = nn.Conv2d(64, 96, 3, padding=1)
        self.f1, self.f2, self.f3 = (FejerMarkov2d(c, gamma) for c in (32,64,96))
        self.relu = nn.ReLU(); self.pool = nn.MaxPool2d(2)
        self.head = nn.Sequential(nn.Flatten(), nn.Linear(96*4*4,128), nn.ReLU(), nn.Linear(128,out_dim))

    def forward(self, x):
        x = self.pool(self.relu(self.f1(self.c1(x))))
        x = self.pool(self.relu(self.f2(self.c2(x))))
        x = self.pool(self.relu(self.f3(self.c3(x))))
        return self.head(x)


def baseline_model():
    return bench.make_model("cnn_small", (3,16,16), 10)

def idea_model():
    return MarkovCNN(10, gamma=0.7)

def run_one(kind, seed, lr, return_model=False):
    seed_all(seed)
    ds = bench.get_dataset("vision", seed=seed, n_train=NTRAIN, n_test=NTEST)
    model = baseline_model() if kind == "baseline" else idea_model()
    trained, metric, history = bench.train_model(model, ds, epochs=EPOCHS, lr=lr,
                                                  batch=128, weight_decay=0.0,
                                                  log=lambda *_: None)
    if trained is None: raise RuntimeError("bench training failed")
    return (float(metric), trained, ds) if return_model else float(metric)


def sweep(kind):
    def make_fn(cfg):
        return lambda seed: run_one(kind, seed, cfg["lr"])
    return bench.sweep_baseline(make_fn, GRID, seeds=(0,1,2,3))


def activation_signature(kind, lr, seed=0):
    metric, model, ds = run_one(kind, seed, lr, return_model=True)
    vals = []
    hooks=[]
    # Signature probing is outside bench.train_model; force the robust CPU path.
    model = model.to("cpu").eval()
    # Compare each trained system's post-convolution feature sup norm to its
    # pre-filter/pre-ReLU convolution output on the same held-out batch.
    if kind == "idea":
        pairs = [(model.c1, model.f1), (model.c2, model.f2), (model.c3, model.f3)]
        for conv, filt in pairs:
            def hook(module, inp, out, filt=filt):
                with torch.no_grad():
                    z = filt(out)
                    vals.append(float(z.abs().amax() / (out.abs().amax()+1e-12)))
            hooks.append(conv.register_forward_hook(hook))
    else:
        for conv in [model.net[0], model.net[3], model.net[6]]:
            def hook(module, inp, out):
                with torch.no_grad():
                    vals.append(float(out.abs().amax() / (inp[0].abs().amax()+1e-12)))
            hooks.append(conv.register_forward_hook(hook))
    old_cudnn = torch.backends.cudnn.enabled
    torch.backends.cudnn.enabled = False
    try:
        with torch.no_grad(): model(ds["xte"][:64].to("cpu"))
    finally:
        torch.backends.cudnn.enabled = old_cudnn
    for h in hooks: h.remove()
    return {"test_metric": metric, "layer_ratios": vals,
            "max_ratio": max(vals), "prediction": "Markov filter ratio <= 1 relative to pre-filter activation"}


def main():
    base = sweep("baseline")
    idea_sweep = []
    for cfg in GRID:
        r = evaluate(lambda seed, cfg=cfg: run_one("idea", seed, cfg["lr"]), seeds=(0,1,2,3))
        idea_sweep.append({"cfg": cfg, "mean": r["mean"]})
    best_cfg = min(idea_sweep, key=lambda x:x["mean"])["cfg"]
    idea_full = evaluate(lambda seed: run_one("idea", seed, best_cfg["lr"]), seeds=SEEDS)
    base_sig = activation_signature("baseline", base["best_cfg"]["lr"])
    idea_sig = activation_signature("idea", best_cfg["lr"])
    report = bench.make_report("vision", "cnn_small", base, idea_full,
        extra={"prediction": "positive normalized filter is non-expansive in sup norm",
               "baseline_signature": base_sig, "idea_signature": idea_sig,
               "predicted_max_ratio": 1.0,
               "observed_max_ratio": idea_sig["max_ratio"],
               "confirmed": idea_sig["max_ratio"] <= 1.05})
    report["idea_sweep"] = idea_sweep
    report["protocol"] = {"paired_seeds": list(SEEDS), "epochs": EPOCHS,
                           "n_train": NTRAIN, "n_test": NTEST,
                           "filter_gamma": 0.7, "same_lr_grid": True}
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == "__main__": main()