"""
train.py - train PixelModel v1 on caption/image pairs, save into model.png.

Expects the npz produced by fetch_coco_subset.py:
  images   uint8 (N, 64, 64, 3)
  captions json list of N strings (stored alongside as captions.json)

Each step samples a batch of captions and a random subset of pixel
coordinates (the CPPN decoder makes per-pixel training natural), computes
MSE against the target pixels, and Adam-steps the weights. Weights are
clamped to [-WMAX, WMAX] so they always round-trip through the 16-bit
PNG codec. model.png is written every epoch.

Usage:
  python train.py --data ../pm-work/coco_train.npz
  python train.py --data ../pm-work/coco_train.npz --epochs 40 --lr 2e-3
"""

import argparse
import json
import os
import time

import numpy as np
import torch

from model import (
    NATIVE_RES, N_PARAMS, PARAM_SPECS, WMAX,
    coord_features, decode_pixels, encode_prompt,
    init_weights, load_model, prompts_to_embeddings, save_model,
)

MODEL_PATH = "model.png"


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--data",     required=True, help="npz with images (N,64,64,3) uint8")
    p.add_argument("--captions", default=None,  help="json list of captions (default: <data>.captions.json)")
    p.add_argument("--epochs",   type=int,   default=30)
    p.add_argument("--batch",    type=int,   default=128)
    p.add_argument("--pixels",   type=int,   default=768, help="pixel coords sampled per step")
    p.add_argument("--lr",       type=float, default=2e-3)
    p.add_argument("--seed",     type=int,   default=0)
    p.add_argument("--device",   default="auto", help="auto, cpu, cuda, or a PyTorch device string")
    p.add_argument("--resume",   action="store_true", help="continue from existing model.png")
    args = p.parse_args()

    torch.manual_seed(args.seed)
    rng = np.random.default_rng(args.seed)
    device = torch.device("cuda" if args.device == "auto" and torch.cuda.is_available()
                          else "cpu" if args.device == "auto" else args.device)
    print(f"device: {device}")

    data = np.load(args.data)
    images = data["images"]                                  # (N, 64, 64, 3) uint8
    cap_path = args.captions or args.data.replace(".npz", ".captions.json")
    with open(cap_path, encoding="utf-8") as f:
        captions = json.load(f)
    N = len(captions)
    assert images.shape[0] == N
    res = images.shape[1]
    print(f"dataset: {N} pairs @ {res}x{res}")

    print("precomputing prompt embeddings...")
    embs = prompts_to_embeddings(captions).to(device)        # (N, 64)

    targets = torch.from_numpy(images.reshape(N, res * res, 3).astype(np.float32) / 255.0).to(device)
    feats_all = coord_features(res).to(device)               # (res*res, 18)

    if args.resume and os.path.exists(MODEL_PATH):
        weights = load_model(MODEL_PATH)
        print(f"resumed from {MODEL_PATH}")
    else:
        weights = init_weights(args.seed)
    for w in weights.values():
        w.data = w.data.to(device)
        w.requires_grad_(True)
    params = [weights[n] for n, _ in PARAM_SPECS]
    opt = torch.optim.Adam(params, lr=args.lr)

    steps_per_epoch = N // args.batch
    print(f"training: {args.epochs} epochs x {steps_per_epoch} steps "
          f"(batch={args.batch}, pixels/step={args.pixels}, lr={args.lr}, "
          f"params={N_PARAMS})\n")

    t0 = time.time()
    for epoch in range(1, args.epochs + 1):
        order = rng.permutation(N)
        ep_loss, ep_steps = 0.0, 0
        for s in range(steps_per_epoch):
            idx = order[s * args.batch:(s + 1) * args.batch]
            pix = torch.from_numpy(rng.choice(res * res, size=args.pixels, replace=False)).to(device)

            emb = embs[idx]
            tgt = targets[idx][:, pix, :]                    # (B, P, 3)

            z = encode_prompt(weights, emb)
            pred = decode_pixels(weights, z, feats_all[pix])
            loss = torch.nn.functional.mse_loss(pred, tgt)

            opt.zero_grad()
            loss.backward()
            opt.step()
            with torch.no_grad():
                for w in params:
                    w.clamp_(-WMAX, WMAX)

            ep_loss += loss.item()
            ep_steps += 1

        save_model(weights, MODEL_PATH)
        elapsed = time.time() - t0
        print(f"epoch {epoch:>3}/{args.epochs}  loss={ep_loss / ep_steps:.5f}  "
              f"elapsed={elapsed:.0f}s  -> saved {MODEL_PATH}", flush=True)

    print(f"\nDone in {time.time() - t0:.0f}s. Final model saved to {MODEL_PATH}")


if __name__ == "__main__":
    main()