# E-OS v5 — Training (on-the-fly, chunked/resumable).
# Aufruf: python3 os5_train.py [--secs 520] [--chunks N]
# On-the-fly: Prozedurale Glas-Paare, kein Dataset auf Disk.
# Checkpoint: /home/z/my-project/eos-v5/train/v5_ckpt.pt
import argparse
import math
import os
import sys
import time

import numpy as np
import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from os5_model import FrameGeneratorV5, grad_loss  # noqa: E402

CKPT = "/home/z/my-project/eos-v5/train/v5_ckpt.pt"
PREV = "/home/z/my-project/eos-v5/previews"
BATCH = 4
TH, TW = 272, 128          # Train-Hint (H x W) -> Out 544 x 256
VAL_SEED0 = 900000


def make_batch(seeds, hh=None, hw=None):
    import os5_glass as G
    hh = hh or TH; hw = hw or TW
    hs = np.empty((len(seeds), 3, hh, hw), np.float32)
    vs = np.empty((len(seeds), 32), np.float32)
    ts = np.empty((len(seeds), hh * 2, hw * 2, 3), np.uint8)
    for i, sd in enumerate(seeds):
        h, v, t = G.sample_pair(int(sd))
        if hh != TH:   # App-Res: Hint hochskalieren (wie App) + Target = App-Res
            h = np.repeat(np.repeat(h, 2, 1), 2, 2)   # (3,272,128)->(3,544,256)
            t = np.repeat(np.repeat(t, 2, 0), 2, 1)
        hs[i], vs[i], ts[i] = h, v, t
    x = torch.from_numpy(hs)
    s = torch.from_numpy(vs)
    y = torch.from_numpy(ts).permute(0, 3, 1, 2).float() / 255.0
    return x, s, y


def evaluate(model, n=48):
    model.eval()
    psnrs = []
    with torch.no_grad():
        for i in range(0, n, 8):
            x, s, y = make_batch([VAL_SEED0 + j for j in range(i, i + 8)])
            p = model(x, s).clamp(0, 1)
            mse = ((p - y) ** 2).mean(dim=(1, 2, 3))
            psnrs += (10 * torch.log10(1.0 / mse)).tolist()
    model.train()
    return float(np.mean(psnrs))


def save_preview(model, step):
    import cv2
    os.makedirs(PREV, exist_ok=True)
    model.eval()
    rows = []
    with torch.no_grad():
        seeds = [200 + i for i in range(3)]
        x, s, y = make_batch(seeds)
        p = model(x, s).clamp(0, 1)
    for j in range(3):
        hint = x[j].permute(1, 2, 0).numpy()
        hint = np.repeat(np.repeat(hint, 2, 0), 2, 1)
        row = np.concatenate([hint, y[j].permute(1, 2, 0).numpy(),
                              p[j].permute(1, 2, 0).numpy()], axis=1)
        rows.append(row)
    grid = (np.clip(np.concatenate(rows, 0), 0, 1) * 255).astype(np.uint8)
    cv2.imwrite(os.path.join(PREV, "train_%06d.png" % step), grid)
    model.train()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--secs", type=int, default=500)
    ap.add_argument("--max-steps", type=int, default=200000)
    args = ap.parse_args()

    os.makedirs(os.path.dirname(CKPT), exist_ok=True)
    torch.set_num_threads(2)
    model = FrameGeneratorV5()
    opt = torch.optim.Adam(model.parameters(), lr=2.5e-4)
    step = 0
    hist = []
    if os.path.exists(CKPT):
        ck = torch.load(CKPT, map_location="cpu", weights_only=False)
        model.load_state_dict(ck["model"])
        opt.load_state_dict(ck["opt"])
        step = ck["step"]
        hist = ck["hist"]
        print("resume at step", step, flush=True)

    t0 = time.time()
    fbudget = args.secs
    next_eval = 400 if step == 0 else (step // 400 + 1) * 400
    lr_sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=12000, eta_min=5e-5)
    for _ in range(step):  # fast-forward scheduler position
        lr_sched.step()

    model.train()
    seeds = np.random.randint(0, 800000, size=(BATCH,))
    while step < args.max_steps:
        if time.time() - t0 > fbudget:
            break
        it0 = time.time()
        app_step = (step % 2 == 1)          # Mixed-Res: jede 2. Step in App-Aufloesung
        if app_step:
            x, s, y = make_batch(seeds[:1], 544, 256)
        else:
            x, s, y = make_batch(seeds)
        p = model(x, s)
        loss = (p - y).abs().mean() + 0.5 * grad_loss(p, y)
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
        opt.step()
        lr_sched.step()
        step += 1
        if step % 25 == 0:
            el = time.time() - t0
            print("step %d  loss %.4f  %.2f s/it  eta %d s" %
                  (step, loss.item(), el / max(1, step % 400 if step % 400 else 1) if False else (time.time() - it0),
                   max(0, fbudget - int(el))), flush=True)
        if step >= next_eval:
            ps = evaluate(model)
            hist.append((step, ps))
            save_preview(model, step)
            print("EVAL step %d  val PSNR %.2f dB" % (step, ps), flush=True)
            next_eval += 400
        seeds = np.random.randint(0, 800000, size=(BATCH,))

    ps = evaluate(model)
    hist.append((step, ps))
    torch.save({"model": model.state_dict(), "opt": opt.state_dict(),
                "step": step, "hist": hist}, CKPT)
    print("CHUNK DONE step %d  val PSNR %.2f dB  saved" % (step, ps), flush=True)


if __name__ == "__main__":
    main()