"""RSSD: train the proxy inversion network Phi, then estimate the
reconstruction-sensitive subspace from the gradient covariance (paper Sec. IV-B).

Outputs checkpoints/rssd.pt with P_rec, P_N (C x C), U_rec, eigenvalues, and the
trained proxy checkpoint checkpoints/proxy_unet.pt.
"""
import argparse

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
import wandb

from models import load_ir50, ClientEncoder, UNetInverter
from data import MS1MDataset
from common import ADAFACE_CKPT, CKPT, CFG, DEVICE

PROXY_IDS = 5000        # identity folders used to train the proxy
PROXY_ITERS = 30000     # paper: 30k iterations
BATCH = 128
COV_SAMPLES = 20000     # samples used to estimate the gradient covariance


def infinite(loader):
    while True:
        for batch in loader:
            yield batch


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ms1m_root", required=True)
    args = ap.parse_args()

    CKPT.mkdir(exist_ok=True)
    dev = DEVICE
    wandb.init(project="decoyface-repro", name="rssd", mode="online", config=CFG)

    ir50 = load_ir50(ADAFACE_CKPT, dev)
    Fc = ClientEncoder(ir50).to(dev).eval()
    for p in Fc.parameters():
        p.requires_grad_(False)

    ds = MS1MDataset(args.ms1m_root, PROXY_IDS)
    loader = DataLoader(ds, batch_size=BATCH, shuffle=True, num_workers=8,
                        drop_last=True, pin_memory=True, persistent_workers=True)
    it = infinite(loader)

    proxy = UNetInverter().to(dev)
    opt = torch.optim.Adam(proxy.parameters(), lr=1e-3)

    proxy_path = CKPT / "proxy_unet.pt"
    resume_path = CKPT / "proxy_unet_resume.pt"
    start = 0
    if resume_path.exists():
        ck = torch.load(resume_path, map_location=dev)
        proxy.load_state_dict(ck["model"])
        opt.load_state_dict(ck["opt"])
        start = ck["step"]
        print(f"resuming proxy from step {start}", flush=True)
    for step in range(start, PROXY_ITERS):
        x, _ = next(it)
        x = x.to(dev, non_blocking=True)
        with torch.no_grad():
            z = Fc(x)
        rec = proxy(z)
        loss = F.l1_loss(rec, x)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step()
        if step % 50 == 0:
            wandb.log({"proxy/l1": loss.item()}, step=step)
        if step % 2000 == 0:
            print(f"proxy step {step} l1 {loss.item():.4f}", flush=True)
            torch.save({"step": step, "model": proxy.state_dict(),
                        "opt": opt.state_dict()}, resume_path)
    torch.save(proxy.state_dict(), proxy_path)

    # ---- gradient covariance over the channel-averaged reconstruction grad ----
    proxy.eval()
    gbar = []
    seen = 0
    for x, _ in DataLoader(ds, batch_size=64, shuffle=True, num_workers=8):
        x = x.to(dev)
        z = Fc(x).detach().clone().requires_grad_(True)
        loss = F.l1_loss(proxy(z), x)
        g = torch.autograd.grad(loss, z)[0]          # (B,C,H,W)
        gbar.append(g.mean(dim=(2, 3)).detach().cpu())  # (B,C)
        seen += x.shape[0]
        if seen >= COV_SAMPLES:
            break
    gbar = torch.cat(gbar, 0)                          # (N,C)
    mu = gbar.mean(0, keepdim=True)
    gc = gbar - mu
    sigma = (gc.T @ gc) / (gbar.shape[0] - 1)          # (C,C)
    evals, evecs = torch.linalg.eigh(sigma)            # ascending
    order = torch.argsort(evals, descending=True)
    evals, evecs = evals[order], evecs[:, order]
    U_rec = evecs[:, :CFG["r"]]                        # (C,r)
    P_rec = U_rec @ U_rec.T
    P_N = torch.eye(CFG["C"]) - P_rec

    torch.save({"P_rec": P_rec, "P_N": P_N, "U_rec": U_rec,
                "evals": evals, "sigma": sigma}, CKPT / "rssd.pt")
    ev = evals / evals.sum()
    print("top-8 eigenvalue energy:", ev[:8].tolist(), flush=True)
    print(f"RSS energy captured by top-{CFG['r']}: "
          f"{ev[:CFG['r']].sum().item():.4f}", flush=True)
    wandb.log({"rssd/rss_energy": ev[:CFG["r"]].sum().item()})
    wandb.finish()


if __name__ == "__main__":
    main()