# E-OS v6 — chunked/resumable trainer for the Liquid-Glass FrameGenerator.
# Usage:  timeout 570 python3 train_os6.py [extra_phase1_steps]
# Phases: 0 = half-res (hint 128x272 -> frame 256x544, batch 6)
#         1 = app-res  (hint 256x544 -> frame 512x1088, batch 2)
# Checkpoint: /home/z/my-project/checkpoints6/gen_ckpt.pt every 30 steps.
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/os6")
from os6_models import FrameGenerator
from os6_draw import rasterize_hint, build_state

DATA = "/home/z/my-project/data/einkos6"
CKPT = "/home/z/my-project/checkpoints6/gen_ckpt.pt"
PHASE0_STEPS = 4200
PHASE1_STEPS = 700

torch.set_num_threads(2)
try:
    torch.multiprocessing.set_sharing_strategy("file_system")   # tiny /dev/shm
except Exception:
    pass


class GlassDS(Dataset):
    def __init__(self, phase, recs):
        self.phase = phase
        self.recs = recs

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

    def __getitem__(self, i):
        rec = self.recs[i]
        fr = np.asarray(Image.open(os.path.join(
            DATA, "frames", "%06d.jpg" % rec["i"])).convert("RGB"))
        if self.phase == 0:
            fr = fr.reshape(544, 2, 256, 2, 3).mean(axis=(1, 3))
            hw, hh = 128, 272
        else:
            hw, hh = 256, 544
        hint = rasterize_hint(rec["shapes"], rec["ripples"], hw, hh)
        st = build_state(rec["theme"], rec["shimmer"], rec["day"], rec["pressed"])
        t = torch.from_numpy(np.ascontiguousarray(fr)).permute(2, 0, 1).float() / 255.0
        h = torch.from_numpy(hint)
        s = torch.from_numpy(st)
        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.5 * e, l1.item(), e.item()


def main():
    recs = [json.loads(l) for l in open(os.path.join(DATA, "layouts.jsonl"))]
    dev = "cpu"
    model = FrameGenerator().to(dev)
    step = 0
    phase = 0
    best = 1e9
    ck = None
    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(len(recs)))
    import random
    random.Random(7).shuffle(idxs)
    ds = GlassDS(phase, [recs[i] for i in idxs])
    dl = DataLoader(ds, batch_size=bs, shuffle=True, num_workers=0,
                    drop_last=True)
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5)
    if ck is not None and phase == ck["phase"]:
        try:
            opt.load_state_dict(ck["o"])
        except Exception:
            pass

    def val_psnr(idxs_val):
        ps = []
        with torch.no_grad():
            for vi in idxs_val:
                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)))
        return float(np.mean(ps))

    t0 = time.time()
    ep = 0
    while True:
        for hb, sb, tb in dl:
            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()
                vidx = [i for i in range(len(ds)) if i % 16 == 5][:24]
                v = val_psnr(vidx)
                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:
                    ck2 = torch.load(CKPT, map_location=dev, weights_only=False)
                    ck2["phase"] = 1; ck2["step"] = PHASE0_STEPS
                    tmp = CKPT + ".tmp"
                    torch.save(ck2, 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()