# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 Lucasgitguy-co
import json
import random
import time
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader

from .data import CachedDataset, image_tensor, load_manifest
from .model import Student, distillation_loss
from .splats import Splats, write_ply


def device_for(name="auto"):
    if name == "auto":
        name = "mps" if torch.backends.mps.is_available() else "cpu"
    if name == "mps" and not torch.backends.mps.is_available():
        raise ValueError("Metal GPU unavailable in this process. Run in a normal macOS Terminal, or use --device cpu.")
    if name not in ("mps", "cpu"):
        raise ValueError("This prototype supports cpu and mps")
    torch.set_num_threads(min(4, torch.get_num_threads()))
    return torch.device(name)


def split_rows(rows, fraction, seed):
    if not 0 <= fraction < 1:
        raise ValueError("Validation fraction must lie in [0,1)")
    # Keep repeated views/crops of the same object together. Supply object_id in cache metadata.
    groups = {}
    for row in rows:
        group = row.get("metadata", {}).get("object_id", row["id"])
        groups.setdefault(group, []).append(row)
    keys = sorted(groups)
    random.Random(seed).shuffle(keys)
    n = max(1, round(len(keys) * fraction)) if fraction else 0
    if n >= len(keys):
        raise ValueError("Need at least two different objects for validation, or explicitly use --validation-fraction 0 for an overfit test")
    val_keys = set(keys[:n])
    train = [r for k in keys if k not in val_keys for r in groups[k]]
    val = [r for k in keys if k in val_keys for r in groups[k]]
    if {r["image_sha256"] for r in train} & {r["image_sha256"] for r in val}:
        raise ValueError("Identical images occur across training/validation. Assign the same object-id to related samples.")
    return train, val


def on_device(target, device):
    return {k: v.to(device) for k, v in target.items()}


@torch.no_grad()
def evaluate(model, loader, device):
    model.eval()
    sums = {}
    for image, target in loader:
        target = on_device(target, device)
        pred = model(image.to(device))
        loss, terms = distillation_loss(pred, target, model.grid)
        occupied = pred["occ"].sigmoid() > .5
        truth = target["occ"] > .5
        terms = {**terms, "loss": loss, "occupancy_iou": (occupied & truth).sum() / (occupied | truth).sum().clamp_min(1)}
        for key, value in terms.items():
            sums[key] = sums.get(key, 0.) + float(value)
    return {k: v / len(loader) for k, v in sums.items()}


def save_checkpoint(path, model, optimizer, step, metadata):
    payload = {"format": 1, "config": model.config, "model": {k: v.detach().cpu() for k, v in model.state_dict().items()},
               "optimizer": optimizer.state_dict(), "step": step, "metadata": metadata}
    tmp = Path(path).with_suffix(".tmp")
    torch.save(payload, tmp)
    tmp.replace(path)


def train(args):
    if args.steps < 1 or args.batch_size < 1 or args.lr <= 0 or args.eval_every < 1:
        raise ValueError("Steps, batch size, learning rate, and evaluation interval must be positive")
    device = device_for(args.device)
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    rows = load_manifest(args.manifest)
    synthetic = any(r["source"] == "synthetic" for r in rows)
    if synthetic and not args.allow_synthetic:
        raise ValueError("Synthetic fixtures are code tests, not teacher data. Add --allow-synthetic for an explicit smoke test.")
    train_rows, val_rows = split_rows(rows, args.validation_fraction, args.seed)
    train_data = CachedDataset(args.manifest, train_rows)
    generator = torch.Generator().manual_seed(args.seed)
    loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=0, generator=generator)
    val_loader = DataLoader(CachedDataset(args.manifest, val_rows), batch_size=1) if val_rows else None
    model = Student(train_data.grid, args.width).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
    run = Path(args.out)
    if run.exists() and any(run.iterdir()):
        raise ValueError("Run directory must be empty; choose a new --out directory")
    run.mkdir(parents=True, exist_ok=True)
    meta = {"synthetic_only": all(r["source"] == "synthetic" for r in rows),
            "contains_synthetic": synthetic, "device": str(device), "torch": str(torch.__version__),
            "seed": args.seed, "train_ids": [r["id"] for r in train_rows], "val_ids": [r["id"] for r in val_rows],
            "parameters": sum(p.numel() for p in model.parameters()),
            "teachers": sorted({v for r in rows for v in r["teachers"].values() if v}),
            "validation": "held-out objects" if val_rows else "none; training overfit test only",
            "dataset": str(Path(args.manifest).resolve())}
    (run / "run.json").write_text(json.dumps(meta, indent=2))
    print(json.dumps(meta), flush=True)
    step, best, started = 0, float("inf"), time.monotonic()
    iterator = iter(loader)
    try:
        while step < args.steps:
            try:
                images, target = next(iterator)
            except StopIteration:
                iterator = iter(loader)
                images, target = next(iterator)
            model.train()
            optimizer.zero_grad(set_to_none=True)
            pred = model(images.to(device))
            loss, terms = distillation_loss(pred, on_device(target, device), model.grid)
            if not torch.isfinite(loss):
                raise RuntimeError("Nonfinite training loss; stopping")
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
            optimizer.step()
            step += 1
            if step == 1 or step % args.eval_every == 0 or step == args.steps:
                record = {"step": step, "seconds": round(time.monotonic()-started, 3),
                          "train_loss": float(loss.detach()), **{k: float(v.detach()) for k, v in terms.items()}}
                if val_loader:
                    record["validation"] = evaluate(model, val_loader, device)
                    score = record["validation"]["loss"]
                    if score < best:
                        best = score
                        save_checkpoint(run / "best.pt", model, optimizer, step, meta)
                with (run / "metrics.jsonl").open("a") as handle:
                    handle.write(json.dumps(record) + "\n")
                print(json.dumps(record), flush=True)
                save_checkpoint(run / "last.pt", model, optimizer, step, meta)
    except KeyboardInterrupt:
        print("Interrupted; saving completed optimizer steps.", flush=True)
    finally:
        if step:
            save_checkpoint(run / "last.pt", model, optimizer, step, meta)
    return meta


@torch.no_grad()
def generate(checkpoint, image, out, device="auto", threshold=.5, max_splats=32768):
    if not 0 < threshold < 1 or max_splats < 1:
        raise ValueError("Threshold must be in (0,1) and max-splats positive")
    device = device_for(device)
    state = torch.load(checkpoint, map_location="cpu", weights_only=True)
    if state.get("format") != 1 or state.get("step", 0) < 1:
        raise ValueError("Expected a trained image3d-student checkpoint")
    model = Student(**state["config"])
    model.load_state_dict(state["model"])
    model.to(device).eval()
    if device.type == "mps":
        torch.mps.synchronize()
    started = time.monotonic()
    pred = model(image_tensor(image)[None].to(device))
    score = pred["occ"][0, :, 0].sigmoid()
    valid = torch.nonzero(score >= threshold).flatten()
    if len(valid) == 0:
        raise ValueError("No occupied cells above threshold. Train longer or inspect validation; no fabricated fallback splats were exported.")
    valid = valid[torch.argsort(score[valid], descending=True)[:max_splats]]
    values = {k: pred[k][0, valid].cpu().numpy() for k in ("xyz", "rgb", "scale", "quat", "opacity")}
    elapsed = time.monotonic() - started
    splats = Splats(**values)
    out = Path(out)
    if out.exists() or out.with_suffix(".json").exists():
        raise ValueError("Output already exists; choose a new filename")
    write_ply(out, splats)
    report = {"gaussians": len(valid), "seconds_preprocess_and_forward": elapsed,
              "device": str(device), "checkpoint_step": state["step"], "training": state["metadata"],
              "note": ("Synthetic-trained output is a pipeline demonstration, not a pretrained image-to-3D result."
                       if state["metadata"].get("contains_synthetic") else "Experimental cached-target student; inspect unseen-object quality.")}
    out.with_suffix(".json").write_text(json.dumps(report, indent=2))
    return report