"""Canonical runners for non-backprop baselines outside the SDIL family.

The methods intentionally do not all share one architecture: Forward-Forward
has goodness layers instead of a classifier readout, while Equilibrium
Propagation is an energy-based recurrent network with symmetric connections.
Forcing either into SDILNet would make the comparison look uniform while
silently changing the algorithm.  Every JSON result records the protocol.
"""
import argparse
import json
import os
import subprocess
import sys
import time

import torch
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sdil.data import get_dataset, onehot
from sdil.local_baselines import FANet, PEPITANet, FFNet, EPNet
from sdil import probes


METHOD_SOURCES = {
    "fa": {
        "paper": "https://www.nature.com/articles/ncomms13276",
        "protocol": "fixed random sequential feedback matrices",
    },
    "pepita": {
        "paper": "https://proceedings.mlr.press/v162/dellaferrera22a.html",
        "code": "https://github.com/GiorgiaD/PEPITA",
        "protocol": "ERIN two-forward-pass rule; He-uniform; ReLU; F scale 0.05",
    },
    "ff": {
        "paper": "https://www.cs.toronto.edu/~hinton/absps/FFXfinal.pdf",
        "reference": "https://github.com/mpezeshki/pytorch_forward_forward",
        "protocol": "greedy layerwise goodness; input length normalization; Adam",
    },
    "ep": {
        "paper": "https://www.frontiersin.org/journals/computational-neuroscience/articles/10.3389/fncom.2017.00024/full",
        "code": "https://github.com/bscellier/Towards-a-Biologically-Plausible-Backprop",
        "code_revision": "cb73f76d997924b6198355e22b50ed7e97cec684",
        "protocol": "two-phase energy-gradient dynamics; persistent free particles; random beta sign",
    },
}

_MNIST_STATS = {"mnist": (0.1307, 0.3081), "fmnist": (0.2860, 0.3530)}
_CIFAR_MEAN = torch.tensor([0.4914, 0.4822, 0.4465])
_CIFAR_STD = torch.tensor([0.2470, 0.2435, 0.2616])


def code_provenance():
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    try:
        commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=root, check=True,
                                capture_output=True, text=True).stdout.strip()
        dirty = bool(subprocess.run(
            ["git", "status", "--porcelain", "--untracked-files=no"], cwd=root,
            check=True, capture_output=True, text=True).stdout.strip())
        return {"git_commit": commit, "git_dirty": dirty}
    except (OSError, subprocess.CalledProcessError):
        return {"git_commit": None, "git_dirty": None}


def canonical_input(x, dataset, method, enabled=True):
    """Restore [0,1] pixels where the authors' protocol requires them."""
    if not enabled or method not in ("pepita", "ep"):
        return x
    if dataset in _MNIST_STATS:
        mean, std = _MNIST_STATS[dataset]
        return (x * std + mean).clamp(0, 1)
    if dataset == "cifar10":
        mean = _CIFAR_MEAN.to(x.device, x.dtype).view(1, 3, 1)
        std = _CIFAR_STD.to(x.device, x.dtype).view(1, 3, 1)
        return (x.view(-1, 3, 1024) * std + mean).clamp(0, 1).view_as(x)
    return x


@torch.no_grad()
def evaluate_method(net, loader, method, dataset, canonical):
    correct = total = 0
    loss_sum = 0.0
    for x, y in loader:
        x = canonical_input(x, dataset, method, canonical)
        if method == "ff":
            pred = net.predict(x)
        elif method == "ep":
            state = net._settle(x, beta=0.0, T=net.T_free)
            logits = state[-1]
            pred = logits.argmax(1)
            loss_sum += F.mse_loss(logits, onehot(y, 10), reduction="sum").item()
        else:
            logits = net.logits(x)
            pred = logits.argmax(1)
            loss_sum += F.cross_entropy(logits, y, reduction="sum").item()
        correct += (pred == y).sum().item()
        total += y.shape[0]
    return correct / total, loss_sum / total if method != "ff" else float("nan")


def ep_learning_rates(depth, base_eta=None):
    """Layerwise rates from the authors' one/two/three-hidden-layer runs."""
    canonical = {
        1: [0.1, 0.05],
        2: [0.4, 0.1, 0.01],
        3: [0.128, 0.032, 0.008, 0.002],
    }
    if base_eta is not None:
        return [base_eta / (4 ** l) for l in range(depth + 1)]
    if depth in canonical:
        return canonical[depth]
    raise ValueError("non-canonical EP depth requires an explicit --eta")


def ep_dynamics(depth, beta=None, free_steps=None, nudge_steps=None):
    """Original-model dynamics for the published 1/2/3-hidden-layer runs."""
    canonical = {
        1: {"beta": 0.5, "free_steps": 20, "nudge_steps": 4},
        2: {"beta": 1.0, "free_steps": 150, "nudge_steps": 6},
        3: {"beta": 1.0, "free_steps": 500, "nudge_steps": 8},
    }
    if depth not in canonical and (beta is None or free_steps is None or nudge_steps is None):
        raise ValueError(
            "non-canonical EP depth requires --ep_beta, --ep_free_steps, and --ep_nudge_steps")
    defaults = canonical.get(depth, {})
    return {
        "beta": defaults["beta"] if beta is None else beta,
        "free_steps": defaults["free_steps"] if free_steps is None else free_steps,
        "nudge_steps": defaults["nudge_steps"] if nudge_steps is None else nudge_steps,
        "canonical_depth": depth in canonical,
    }


def pepita_protocol(args):
    """Resolve only settings that are actually supported by a cited protocol.

    The original PEPITA paper specifies one hidden layer.  Deep PEPITA is very
    sensitive to architecture-specific learning rates, so silently reusing the
    shallow 0.1 default is worse than requiring an explicit value.
    """
    eta = args.eta
    if eta is None:
        if args.depth != 1:
            raise ValueError("deep PEPITA requires an explicit --eta")
        if args.dataset == "mnist":
            eta = 0.1
        elif args.dataset == "cifar10":
            eta = 0.01
        else:
            raise ValueError("PEPITA on this dataset requires an explicit --eta")
    if args.pepita_decay_epochs is not None:
        decay_epochs = [int(v) for v in args.pepita_decay_epochs.split(",") if v.strip()]
    elif args.depth == 1 and args.dataset == "mnist":
        decay_epochs = [60]
    elif args.depth == 1 and args.dataset == "cifar10":
        decay_epochs = [60, 90]
    else:
        decay_epochs = []
    return {"eta": eta, "decay_epochs": decay_epochs}


def build(args, n_in, device):
    sizes = [n_in] + [args.width] * args.depth + [10]
    if args.method == "fa":
        return FANet(sizes, act=args.act, device=device, seed=args.seed,
                     residual=bool(args.residual), b_scale=args.feedback_scale)
    if args.method == "pepita":
        return PEPITANet(sizes, act="relu", device=device, seed=args.seed,
                         f_scale=args.pepita_f_scale, keep_prob=args.keep_prob)
    if args.method == "ff":
        # FF has only goodness layers; a 10-unit SDIL-style output is not part
        # of the original supervised algorithm.
        return FFNet([n_in] + [args.width] * args.depth, device=device,
                     seed=args.seed, threshold=args.ff_threshold)
    if args.method == "ep":
        dyn = ep_dynamics(args.depth, args.ep_beta, args.ep_free_steps,
                          args.ep_nudge_steps)
        return EPNet(sizes, device=device, seed=args.seed, beta=dyn["beta"],
                     dt=args.ep_dt, T_free=dyn["free_steps"],
                     T_nudge=dyn["nudge_steps"], random_beta_sign=True)
    raise ValueError(args.method)


def train(args):
    torch.manual_seed(args.seed)
    device = args.device
    pepita_cfg = pepita_protocol(args) if args.method == "pepita" else None
    ep_dyn = (ep_dynamics(args.depth, args.ep_beta, args.ep_free_steps,
                          args.ep_nudge_steps) if args.method == "ep" else None)
    use_persistent_ep = args.method == "ep" and bool(args.ep_persistent)
    train_loader, test_loader, n_in, n_out = get_dataset(
        args.dataset, args.batch_size, device=device,
        shuffle_train=not use_persistent_ep,
        train_limit=args.train_examples or None)
    net = build(args, n_in, device)
    canonical = bool(args.canonical_preprocess)
    resolved_protocol = {}
    if pepita_cfg is not None:
        resolved_protocol.update(pepita_cfg)
    if ep_dyn is not None:
        resolved_protocol.update(ep_dyn)
        resolved_protocol["learning_rates"] = ep_learning_rates(args.depth, args.eta)
        resolved_protocol["persistent_particles"] = use_persistent_ep
        resolved_protocol["train_examples"] = train_loader.n
    log = {
        "args": vars(args),
        "resolved_protocol": resolved_protocol,
        "method_source": METHOD_SOURCES[args.method],
        "provenance": code_provenance(),
        "steps": [],
        "final": {},
    }
    if args.method == "fa":
        px, py = next(iter(test_loader))
        px, py = px[:args.probe_bs], py[:args.probe_bs]
        pyoh = onehot(py, n_out, device=device)
    t0 = time.time()

    if args.method == "ff":
        eta = 0.03 if args.eta is None else args.eta
        resolved_protocol["eta"] = eta
        resolved_protocol["epochs_per_layer"] = args.epochs
        for layer in range(args.depth):
            for epoch in range(args.epochs):
                batches = 0
                for x, y in train_loader:
                    loss = net.train_layer(layer, x, y, eta)
                    batches += 1
                    if args.max_batches and batches >= args.max_batches:
                        break
                print(f"[{args.tag}] layer {layer} epoch {epoch} local_loss {loss:.4f}",
                      flush=True)
            acc, test_loss = evaluate_method(net, test_loader, args.method,
                                             args.dataset, canonical)
            log["steps"].append({"layer_end": layer, "test_acc": acc,
                                 "local_loss": loss})
            print(f"[{args.tag}] layer {layer} test_acc {acc:.4f}", flush=True)
    else:
        if args.method == "fa":
            eta = 0.05 if args.eta is None else args.eta
            resolved_protocol["eta"] = eta
            resolved_protocol["feedback_scale"] = args.feedback_scale
        elif args.method == "pepita":
            eta = pepita_cfg["eta"]
        else:
            eta = args.eta
        ep_etas = ep_learning_rates(args.depth, eta) if args.method == "ep" else None
        ep_states = [None] * len(train_loader) if use_persistent_ep else None
        for epoch in range(args.epochs):
            if args.method == "pepita" and epoch in pepita_cfg["decay_epochs"]:
                eta *= 0.1
            batches = 0
            for batch_index, (x, y) in enumerate(train_loader):
                x = canonical_input(x, args.dataset, args.method, canonical)
                yoh = onehot(y, n_out, device=device)
                if args.method == "fa":
                    loss = net.fa_step(x, y, yoh, eta, args.momentum)
                elif args.method == "pepita":
                    loss = net.pepita_step(x, y, yoh, eta, args.momentum)
                else:
                    if use_persistent_ep:
                        loss, ep_states[batch_index] = net.train_step(
                            x, y, yoh, ep_etas, free_state=ep_states[batch_index],
                            return_free_state=True)
                    else:
                        loss = net.train_step(x, y, yoh, ep_etas)
                batches += 1
                if args.max_batches and batches >= args.max_batches:
                    break
            acc, test_loss = evaluate_method(net, test_loader, args.method,
                                             args.dataset, canonical)
            rec = {"epoch_end": epoch, "test_acc": acc,
                   "test_loss": test_loss, "train_loss": loss}
            msg = f"[{args.tag}] epoch {epoch} loss {loss:.4f} test_acc {acc:.4f}"
            if args.method == "fa":
                al = probes.fa_alignment_report(net, px, py, pyoh)
                rec["cos_fa_negg"] = al["cos_fa_negg"]
                msg += f" mean_cos(fa,-g) {sum(al['cos_fa_negg']) / len(al['cos_fa_negg']):+.3f}"
            log["steps"].append(rec)
            print(msg, flush=True)

    acc, test_loss = evaluate_method(net, test_loader, args.method,
                                     args.dataset, canonical)
    log["final"] = {"test_acc": acc, "test_loss": test_loss,
                    "wall_s": time.time() - t0}
    if args.method == "fa":
        log["final"].update(probes.fa_alignment_report(net, px, py, pyoh))
    os.makedirs(args.outdir, exist_ok=True)
    path = os.path.join(args.outdir, f"{args.tag}.json")
    with open(path, "w") as f:
        json.dump(log, f)
    print(f"[{args.tag}] DONE test_acc={acc:.4f} -> {path}", flush=True)
    return log


def get_args():
    p = argparse.ArgumentParser()
    p.add_argument("--method", required=True, choices=["fa", "pepita", "ff", "ep"])
    p.add_argument("--dataset", default="mnist", choices=["mnist", "fmnist", "cifar10"])
    p.add_argument("--depth", type=int, default=2)
    p.add_argument("--width", type=int, default=500)
    p.add_argument("--epochs", type=int, default=10,
                   help="FF: epochs per greedy layer; other methods: global epochs")
    p.add_argument("--batch_size", type=int, default=64)
    p.add_argument("--train_examples", type=int, default=0,
                   help="0 uses the full training split")
    p.add_argument("--max_batches", type=int, default=0)
    p.add_argument("--eta", type=float, default=None)
    p.add_argument("--momentum", type=float, default=0.9)
    p.add_argument("--act", default="tanh", choices=["tanh", "gelu", "silu", "relu"])
    p.add_argument("--residual", type=int, default=0)
    p.add_argument("--feedback_scale", type=float, default=1.0)
    p.add_argument("--pepita_f_scale", type=float, default=0.05)
    p.add_argument("--pepita_decay_epochs", default=None,
                   help="comma-separated; deep PEPITA defaults to no decay")
    p.add_argument("--keep_prob", type=float, default=0.9)
    p.add_argument("--ff_threshold", type=float, default=2.0)
    p.add_argument("--ep_beta", type=float, default=None)
    p.add_argument("--ep_dt", type=float, default=0.5)
    p.add_argument("--ep_free_steps", type=int, default=None)
    p.add_argument("--ep_nudge_steps", type=int, default=None)
    p.add_argument("--ep_persistent", type=int, default=1)
    p.add_argument("--canonical_preprocess", type=int, default=1)
    p.add_argument("--probe_bs", type=int, default=512)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--outdir", default="results")
    p.add_argument("--tag", default="baseline")
    return p.parse_args()


if __name__ == "__main__":
    train(get_args())