"""Fine-tune on Synth5 (real JL1 + synthetic). KIND siam: from siam_v2; KIND unet: from soup_all with L2-SP.
usage: python train12.py KIND MINUTES SNAP P_REAL OUTP INIT
"""
import sys, time
import torch, torch.nn.functional as F
from torch.utils.data import DataLoader
KIND, MINUTES, SNAP, P_REAL, OUTP, INIT = sys.argv[1], float(sys.argv[2]), float(sys.argv[3]), float(sys.argv[4]), sys.argv[5], sys.argv[6]
sys.argv = sys.argv[:1] + ["60", "32"]
from jl_synth import Synth5
from train3 import dice
init = torch.load(INIT, map_location="cpu", weights_only=True)["state_dict"]
if KIND == "siam":
    from model_v2 import SiamCD
    model = SiamCD(); model.load_state_dict(init); model = model.cuda()
    enc = [p for n, p in model.named_parameters() if n.startswith("enc.")]; rest = [p for n, p in model.named_parameters() if not n.startswith("enc.")]
    opt = torch.optim.AdamW([{"params": enc, "base": 2e-5}, {"params": rest, "base": 1e-4}], lr=1e-4, weight_decay=1e-4)
else:
    import segmentation_models_pytorch as smp
    model = smp.Unet(encoder_name="resnet18", encoder_weights=None, in_channels=6, classes=3); model.load_state_dict(init); model = model.cuda()
    anchor = {n: p.detach().clone() for n, p in model.named_parameters()}
    opt = torch.optim.AdamW([{"params": list(model.parameters()), "base": 1e-4}], lr=1e-4, weight_decay=0)
    MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).cuda(); STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).cuda()
dl = DataLoader(Synth5(10**7, P_REAL), batch_size=32 if KIND == "siam" else 64, num_workers=18, pin_memory=True, persistent_workers=True, prefetch_factor=4)
pw = torch.tensor([2.0, 2.0]).cuda().view(1, 2, 1, 1); w3 = torch.tensor([1.0, 2.0, 2.0]).cuda()
t0 = time.time(); step = 0; nsnap = 1
for pre, post, lab, pres, sp, sq in dl:
    pre = pre.cuda(non_blocking=True).float() / 255; post = post.cuda(non_blocking=True).float() / 255
    lab, pres, sp, sq = lab.cuda(non_blocking=True), pres.cuda(non_blocking=True), sp.cuda(non_blocking=True), sq.cuda(non_blocking=True)
    for g in opt.param_groups: g["lr"] = g["base"] * min(1, (step + 1) / 100)
    if KIND == "siam":
        with torch.autocast("cuda", dtype=torch.bfloat16):
            ch, pr, sem_p, sem_q = model(pre, post)
        ch, pr = ch.float(), pr.float()
        loss = F.binary_cross_entropy_with_logits(ch, lab, pos_weight=pw) + dice(ch, lab) + 0.3 * F.binary_cross_entropy_with_logits(pr, pres) \
            + 0.5 * (F.cross_entropy(sem_p.float(), sp, ignore_index=255) + F.cross_entropy(sem_q.float(), sq, ignore_index=255))
        loss = torch.nan_to_num(loss)
    else:
        y = torch.zeros_like(sp); y[lab[:, 1] > 0] = 2; y[lab[:, 0] > 0] = 1
        x = torch.cat([(pre - MEAN) / STD, (post - MEAN) / STD], 1)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            out = model(x)
        out = out.float(); p = out.softmax(1)
        dl_ = sum(1 - (2 * (p[:, c] * (y == c)).sum() + 1) / (p[:, c].sum() + (y == c).sum() + 1) for c in (1, 2)) / 2
        l2 = sum(((q - anchor[n]) ** 2).sum() for n, q in model.named_parameters())
        loss = F.cross_entropy(out, y, weight=w3) + dl_ + 1e-3 * l2
    opt.zero_grad(set_to_none=True); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step(); step += 1
    if step % 200 == 0: print(f"step {step} {time.time() - t0:.0f}s loss {loss.item():.4f}", flush=True)
    el = time.time() - t0
    if el > nsnap * SNAP * 60:
        torch.save({"state_dict": model.state_dict()}, f"{OUTP}_m{int(nsnap * SNAP)}.pt"); nsnap += 1; print("snapshot", int((nsnap - 1) * SNAP), flush=True)
    if el > MINUTES * 60: break
print("TRAIN_DONE", step, flush=True)