# E-OS v4 — chunked/resumable trainer for FrameGenerator.
# Usage:  timeout 570 python3 train_os4.py            (repeat until "TRAINING DONE")
# Phases: 0 = half-res (hint 96x208 -> frame 192x416, batch 6)
#         1 = app-res  (hint 192x416 -> frame 384x832, batch 2)
# Checkpoint: /home/z/my-project/checkpoints4/gen_ckpt.pt every 30 steps (atomic)
import os, sys, time, math, json
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from PIL import Image
sys.path.insert(0, "/home/z/my-project/scripts/os4")
from os4_models import FrameGenerator

DATA = "/home/z/my-project/data/einkos4"
CKPT = "/home/z/my-project/checkpoints4/gen_ckpt.pt"
N_TOTAL = 3200
PHASE0_STEPS = 6000
PHASE1_STEPS = 900

torch.set_num_threads(2)


class Os4DS(Dataset):
    def __init__(self, phase, idxs):
        self.phase = phase
        self.idxs = idxs
        self.fm = np.memmap(os.path.join(DATA, "frames.mem"), dtype=np.uint8, mode="r",
                            shape=(N_TOTAL, 832, 384, 3))
        self.states = np.load(os.path.join(DATA, "states.npy"))

    def __len__(self):
        return len(self.idxs)

    def __getitem__(self, i):
        idx = self.idxs[i]
        fr = np.asarray(self.fm[idx])                      # 832x384x3 uint8
        hint = np.asarray(Image.open(os.path.join(DATA, "hint", "%06d.png" % idx)).convert("RGB"))
        if self.phase == 0:
            fr = fr.reshape(416, 2, 192, 2, 3).mean(axis=(1, 3))
            hint = hint.reshape(208, 2, 96, 2, 3).mean(axis=(1, 3))
        t = torch.from_numpy(np.ascontiguousarray(fr)).permute(2, 0, 1).float() / 255.0
        h = torch.from_numpy(np.ascontiguousarray(hint)).permute(2, 0, 1).float() / 255.0
        s = torch.from_numpy(self.states[idx].copy())
        return h, s, t


def sobel_mag(x):
    kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3)
    ky = kx.transpose(2, 3).contiguous()
    g = x.mean(1, keepdim=True)
    gx = F.conv2d(g, kx, padding=1)
    gy = F.conv2d(g, ky, padding=1)
    return torch.sqrt(gx * gx + gy * gy + 1e-6)


def loss_fn(pred, tgt):
    l1 = (pred - tgt).abs().mean()
    e = (sobel_mag(pred) - sobel_mag(tgt)).abs().mean()
    return l1 + 0.25 * e, l1.item(), e.item()


def make_val(idxs_all):
    return [i for i in idxs_all if i % 16 == 5][:24]


def main():
    dev = "cpu"
    model = FrameGenerator().to(dev)
    opt = None
    step = 0
    phase = 0
    best = 1e9
    if os.path.exists(CKPT):
        ck = torch.load(CKPT, map_location=dev, weights_only=False)
        model.load_state_dict(ck["m"])
        phase = ck["phase"]; step = ck["step"]; best = ck.get("best", 1e9)
        print("resume step %d phase %d best %.4f" % (step, phase, best), flush=True)
    lr = 2.5e-4 if phase == 0 else 8e-5
    lr_end = 1.2e-4 if phase == 0 else 2.5e-5
    extra = int(sys.argv[1]) if len(sys.argv) > 1 else 0
    p0 = step if phase == 0 else 0
    p1 = step - PHASE0_STEPS if phase == 1 else 0
    bs = 6 if phase == 0 else 2
    total = PHASE0_STEPS if phase == 0 else PHASE1_STEPS + extra
    done = p0 if phase == 0 else p1

    idxs = list(range(N_TOTAL))
    import random
    random.Random(7).shuffle(idxs)
    ds = Os4DS(phase, idxs)
    dl = DataLoader(ds, batch_size=bs, shuffle=True, num_workers=2,
                    drop_last=True, persistent_workers=False, prefetch_factor=2)
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5)
    if os.path.exists(CKPT) and phase == ck["phase"]:
        try:
            opt.load_state_dict(ck["o"])
        except Exception:
            pass

    t0 = time.time()
    ep = 0
    while True:
        for hb, sb, tb in dl:
            # cosine lr
            prog = min(1.0, done / float(total))
            cur = lr_end + 0.5 * (lr - lr_end) * (1 + math.cos(math.pi * prog))
            for g in opt.param_groups:
                g["lr"] = cur
            pred = model(hb, sb)
            loss, l1v, ev = loss_fn(pred, tb)
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            done += 1; step += 1
            if done % 30 == 0:
                tmp = CKPT + ".tmp"
                torch.save({"m": model.state_dict(), "o": opt.state_dict(),
                            "step": step, "phase": phase, "best": best}, tmp)
                os.replace(tmp, CKPT)
            if done % 150 == 0:
                el = time.time() - t0
                print("p%d step %d/%d loss %.4f (l1 %.4f e %.4f) lr %.1e %.2f s/it" %
                      (phase, done, total, loss.item(), l1v, ev, cur, el / 150.0), flush=True)
                t0 = time.time()
            if done % 600 == 0 or done == total:
                model.eval()
                ps = []
                with torch.no_grad():
                    for vi in make_val(idxs):
                        h, s, t = ds[vi]
                        p = model(h.unsqueeze(0), s.unsqueeze(0))
                        mse = ((p - t.unsqueeze(0)) ** 2).mean().item()
                        ps.append(10 * math.log10(1.0 / max(1e-8, mse)))
                v = float(np.mean(ps))
                print("  [val] phase %d step %d PSNR %.2f dB" % (phase, done, v), flush=True)
                if v < best:
                    best = v
                    torch.save({"m": model.state_dict(), "step": step, "phase": phase,
                                "best": best, "opt": None}, CKPT.replace(".pt", "_best.pt"))
                model.train()
            if done >= total:
                if phase == 0:
                    phase = 1
                    ck = torch.load(CKPT, map_location=dev, weights_only=False)
                    ck["phase"] = 1; ck["step"] = PHASE0_STEPS
                    tmp = CKPT + ".tmp"
                    torch.save(ck, tmp); os.replace(tmp, CKPT)
                    print("PHASE 0 COMPLETE -> phase 1 (app-res fine-tune)", flush=True)
                    return
                else:
                    torch.save({"m": model.state_dict(), "step": step, "phase": 2,
                                "best": best}, CKPT.replace(".pt", "_final.pt"))
                    print("TRAINING DONE step %d best-val %.2f dB" % (step, best), flush=True)
                    return
        ep += 1


if __name__ == "__main__":
    main()