"""Defender training (paper Sec. IV-E): F_c frozen; train F_s + ArcFace head on
the server-side canonicalized feature z_can with V=2 protected views and the
multi-view consistency loss.  L = L_id + lambda_cons * L_cons.

Outputs checkpoints/defender.pt with the trained F_s and ArcFace head.
"""
import copy
import argparse

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

from models import load_ir50, ClientEncoder, ServerBackbone
from data import MS1MDataset
from decoyface import DecoyProtector
from common import ADAFACE_CKPT, CKPT, CFG, DEVICE

V = 2
LAMBDA_CONS = 0.1
BATCH = 256
EPOCHS = 15
LR = 0.1
DECAY_EPOCHS = [4, 8, 12]


class ArcFace(nn.Module):
    def __init__(self, in_dim, n_cls, s=64.0, m=0.5):
        super().__init__()
        self.W = nn.Parameter(torch.empty(n_cls, in_dim))
        nn.init.xavier_normal_(self.W)
        self.s, self.m = s, m

    def forward(self, emb, y):
        emb = F.normalize(emb)
        w = F.normalize(self.W)
        cos = emb @ w.T
        theta = torch.acos(cos.clamp(-1 + 1e-6, 1 - 1e-6))
        target = torch.cos(theta + self.m)
        onehot = F.one_hot(y, cos.shape[1]).float()
        logits = self.s * (onehot * target + (1 - onehot) * cos)
        return F.cross_entropy(logits, y)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ms1m_root", required=True)
    ap.add_argument("--num_ids", type=int, default=30000)
    ap.add_argument("--epochs", type=int, default=EPOCHS)
    args = ap.parse_args()
    dev = DEVICE

    ir50 = load_ir50(ADAFACE_CKPT, dev)
    Fc = ClientEncoder(ir50).to(dev).eval()
    for p in Fc.parameters():
        p.requires_grad_(False)
    Fs = ServerBackbone(copy.deepcopy(ir50)).to(dev).train()
    rssd = torch.load(CKPT / "rssd.pt", map_location=dev)
    bank = torch.load(CKPT / "bank.pt", map_location=dev)
    protector = DecoyProtector(rssd["P_rec"], rssd["P_N"], bank["mu_rec"],
                               bank["aux_bank"].float(), bank["perm"], CFG, dev)

    ds = MS1MDataset(args.ms1m_root, args.num_ids)
    head = ArcFace(512, ds.num_classes).to(dev)
    loader = DataLoader(ds, batch_size=BATCH, shuffle=True, num_workers=12,
                        drop_last=True, pin_memory=True, persistent_workers=True)

    params = list(Fs.parameters()) + list(head.parameters())
    opt = torch.optim.SGD(params, lr=LR, momentum=0.9, weight_decay=5e-4)
    wandb.init(project="decoyface-repro", name="defender", mode="online",
               config={**CFG, "num_ids": args.num_ids, "epochs": args.epochs})

    step = 0
    for epoch in range(args.epochs):
        lr = LR * (0.1 ** sum(epoch >= d for d in DECAY_EPOCHS))
        for g in opt.param_groups:
            g["lr"] = lr
        for x, y in loader:
            x, y = x.to(dev, non_blocking=True), y.to(dev, non_blocking=True)
            with torch.no_grad():
                z = Fc(x)
            embs = []
            id_loss = 0.0
            for _ in range(V):
                z_B = z.roll(torch.randint(1, x.shape[0], (1,)).item(), 0)
                z_can = protector.server_recover(protector.protect(z, z_B))
                e = Fs(z_can)
                id_loss = id_loss + head(e, y)
                embs.append(F.normalize(e))
            id_loss = id_loss / V
            cons = ((embs[0] - embs[1]) ** 2).sum(1).mean()
            loss = id_loss + LAMBDA_CONS * cons
            opt.zero_grad(set_to_none=True)
            loss.backward()
            opt.step()
            if step % 100 == 0:
                wandb.log({"def/loss": loss.item(), "def/id": id_loss.item(),
                           "def/cons": cons.item(), "def/lr": lr}, step=step)
            if step % 1000 == 0:
                print(f"ep{epoch} step{step} loss{loss.item():.3f} "
                      f"id{id_loss.item():.3f} cons{cons.item():.4f}", flush=True)
            step += 1
        torch.save({"Fs": Fs.state_dict(), "head": head.state_dict()},
                   CKPT / "defender.pt")
    wandb.finish()


if __name__ == "__main__":
    main()