"""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


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",
        "protocol": "two-phase energy-gradient dynamics; 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]
    # The original implementation already required rapidly shrinking rates as
    # depth grew. Continue its depth-3 geometric schedule for an explicit,
    # logged extrapolation rather than pretending there is a canonical one.
    return [0.128 / (4 ** l) for l in range(depth + 1)]


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":
        return EPNet(sizes, device=device, seed=args.seed, beta=args.ep_beta,
                     dt=args.ep_dt, T_free=args.ep_free_steps,
                     T_nudge=args.ep_nudge_steps, random_beta_sign=True)
    raise ValueError(args.method)


def train(args):
    torch.manual_seed(args.seed)
    device = args.device
    train_loader, test_loader, n_in, n_out = get_dataset(
        args.dataset, args.batch_size, device=device)
    net = build(args, n_in, device)
    canonical = bool(args.canonical_preprocess)
    log = {
        "args": vars(args),
        "method_source": METHOD_SOURCES[args.method],
        "provenance": code_provenance(),
        "steps": [],
        "final": {},
    }
    t0 = time.time()

    if args.method == "ff":
        eta = 0.03 if args.eta is None else args.eta
        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:
        eta = args.eta
        if eta is None:
            eta = 0.05 if args.method == "fa" else (0.1 if args.method == "pepita" else None)
        ep_etas = ep_learning_rates(args.depth, eta) if args.method == "ep" else None
        for epoch in range(args.epochs):
            if args.method == "pepita" and epoch in (60, 90):
                eta *= 0.1
            batches = 0
            for x, y in 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:
                    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)
            log["steps"].append({"epoch_end": epoch, "test_acc": acc,
                                 "test_loss": test_loss, "train_loss": loss})
            print(f"[{args.tag}] epoch {epoch} loss {loss:.4f} test_acc {acc:.4f}",
                  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}
    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("--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("--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=0.5)
    p.add_argument("--ep_dt", type=float, default=0.5)
    p.add_argument("--ep_free_steps", type=int, default=20)
    p.add_argument("--ep_nudge_steps", type=int, default=4)
    p.add_argument("--canonical_preprocess", type=int, default=1)
    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())