"""Satlas Siamese (SiamCD) from Satlas Aerial init on Synth4 (EarthView negatives), cosine LR.
usage: python train11.py MINUTES P_EV OUT
"""
import math, sys, time
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
MINUTES, P_EV, OUT = float(sys.argv[1]), float(sys.argv[2]), sys.argv[3]
sys.argv = sys.argv[:1] + ["60", "32"]
from ev_synth import Synth4
from train3 import dice
from model_v2 import SiamCD
model = SiamCD("/workspace/satlas/aerial_swinb_si.pth").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": 1e-4}, {"params": rest, "base": 6e-4}], lr=1e-4, weight_decay=1e-4)
dl = DataLoader(Synth4(10**7, P_EV), batch_size=32, num_workers=17, pin_memory=True, persistent_workers=True, prefetch_factor=4)
pw = torch.tensor([2.0, 2.0]).cuda().view(1, 2, 1, 1)
t0 = time.time(); step = 0; snapped = False
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)
    frac = min((time.time() - t0) / (MINUTES * 60), 1.0)
    for g in opt.param_groups:
        g["lr"] = g["base"] * 0.5 * (1 + math.cos(math.pi * frac)) * min(1, (step + 1) / 300)
    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.5 * (F.cross_entropy(sem_p.float(), sp, ignore_index=255) + F.cross_entropy(sem_q.float(), sq, ignore_index=255)) \
        + 0.3 * F.binary_cross_entropy_with_logits(pr, pres)
    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} ips {step * 32 / (time.time() - t0):.0f}", flush=True)
    if frac >= 0.75 and not snapped:
        torch.save({"state_dict": model.state_dict()}, OUT + "_m75.pt"); snapped = True
    if frac >= 1:
        break
torch.save({"state_dict": model.state_dict()}, OUT + ".pt")
print("TRAIN_DONE", step, flush=True)