#!/usr/bin/env python3
"""Train CellQuant on CVAT-annotated images.

Examples
--------
Final model on every annotated image:
    python train.py --data-dir /path/to/images_and_zips --out weights/cellquant.pt
Leave-one-image-out fold (held-out image is only used for monitoring):
    python train.py --data-dir ... --holdout C110 --out runs/fold_C110.pt
"""
from __future__ import annotations

import argparse
import json
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 cellquant.annotations import discover, load_image, rasterize  # noqa: E402
from cellquant.data import AugConfig, Prefetcher, make_targets, negative_nucleus_points  # noqa: E402
from cellquant.infer import Predictor, pick_device  # noqa: E402
from cellquant.losses import multitask_loss  # noqa: E402
from cellquant.metrics import evaluate_cells, evaluate_nostrils  # noqa: E402
from cellquant.model import CellUNet, save_checkpoint  # noqa: E402
from cellquant.postprocess import PostConfig, analyze  # noqa: E402


def parse_args(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-dir", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--holdout", default="", help="image key to leave out (e.g. C110); empty = train on all")
    ap.add_argument("--iters", type=int, default=3000)
    ap.add_argument("--batch", type=int, default=4)
    ap.add_argument("--crop", type=int, default=512)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--weight-decay", type=float, default=1e-4)
    ap.add_argument("--warmup", type=int, default=100)
    ap.add_argument("--encoder", default="resnet34")
    ap.add_argument("--no-pretrained", action="store_true")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--device", default="auto")
    ap.add_argument("--log-every", type=int, default=50)
    ap.add_argument("--val-every", type=int, default=500)
    ap.add_argument("--workers", type=int, default=3, help="augmentation threads")
    ap.add_argument("--p-focus", type=float, default=AugConfig.p_focus,
                    help="share of crops forced to contain an antibody-negative nucleus (0 = off)")
    ap.add_argument("--ab-neg-weight", default="auto",
                    help="loss weight of antibody-negative nucleus pixels: 'auto' = sqrt(pos px / neg px) clipped to [1, 5], or a number")
    return ap.parse_args(argv)


def antibody_negative_weight(targets: list[dict], spec: str = "auto") -> float:
    """Class-balancing weight for the antibody head (see losses.masked_bce)."""
    if spec != "auto":
        return float(spec)
    pos = sum(float(((t["antibody"] > 0.5) & (t["cellmask"] > 0)).sum()) for t in targets)
    neg = sum(float(((t["antibody"] < 0.5) & (t["cellmask"] > 0)).sum()) for t in targets)
    if neg == 0:
        return 1.0
    return float(np.clip(np.sqrt(pos / neg), 1.0, 5.0))


def lr_at(it: int, total: int, base: float, warmup: int) -> float:
    if it < warmup:
        return base * (it + 1) / warmup
    t = (it - warmup) / max(1, total - warmup)
    return base * (0.02 + 0.98 * 0.5 * (1 + math.cos(math.pi * t)))


def quick_val(model: CellUNet, img: np.ndarray, R, device) -> dict:
    pred = Predictor(model, device=str(device), tta=1)
    res, _ = pred.predict(img, PostConfig())
    ce = evaluate_cells(R.cell_inst, R.cell_labels, res.cell_inst, res.cell_labels)
    ne = evaluate_nostrils(R.nostril_inst, R.nostril_labels, res.nostril_inst, res.nostril_labels)
    return {
        "cell_det_f1": ce["detection"]["f1"],
        "cell_attr_both": ce["attribute_accuracy"]["both"],
        "nos_det_f1": ne["detection"]["f1"],
        "nos_type_acc": ne["type_accuracy"],
        "n_pred_cells": ce["n_pred"], "n_gt_cells": ce["n_gt"],
    }


def main(argv=None):
    a = parse_args(argv)
    torch.manual_seed(a.seed)
    np.random.seed(a.seed)
    device = pick_device(a.device)
    samples, _ = discover(a.data_dir)
    if not samples:
        raise SystemExit("no annotated images found")
    train = [s for s in samples if s.key != a.holdout]
    held = [s for s in samples if s.key == a.holdout]
    if a.holdout and not held:
        raise SystemExit(f"holdout {a.holdout!r} not among {[s.key for s in samples]}")
    print(f"device={device} train={[s.key for s in train]} holdout={[s.key for s in held]}")
    images, targets, focus = [], [], []
    for s in train:
        img = load_image(s.image_path)
        R = rasterize(s.annotation, img.shape[0], img.shape[1])
        images.append(img)
        targets.append(make_targets(R))
        focus.append(negative_nucleus_points(R))
    ab_neg_weight = antibody_negative_weight(targets, a.ab_neg_weight)
    print(f"antibody-negative nuclei in training set: {sum(len(f) for f in focus)} "
          f"(loss weight {ab_neg_weight:.2f}, focus share {a.p_focus:.2f})")
    sampler = Prefetcher(images, targets, AugConfig(crop=a.crop, p_focus=a.p_focus), batch=a.batch, seed=a.seed,
                         workers=a.workers, focus=focus if a.p_focus > 0 else None)
    held_img = held_R = None
    if held:
        held_img = load_image(held[0].image_path)
        held_R = rasterize(held[0].annotation, held_img.shape[0], held_img.shape[1])

    model = CellUNet(a.encoder, pretrained=not a.no_pretrained).to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=a.weight_decay)
    history = {"train": [], "val": []}
    t0 = time.time()
    run = {}
    for it in range(a.iters):
        lr = lr_at(it, a.iters, a.lr, a.warmup)
        for g in opt.param_groups:
            g["lr"] = lr
        model.train()
        x, t = sampler.next()
        out = model(x.to(device))
        loss, parts = multitask_loss(out, t, ab_neg_weight=ab_neg_weight)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
        opt.step()
        for k, v in parts.items():
            run[k] = run.get(k, 0.0) + v
        if (it + 1) % a.log_every == 0:
            avg = {k: v / a.log_every for k, v in run.items()}
            run = {}
            avg["iter"] = it + 1
            avg["lr"] = lr
            avg["elapsed_s"] = time.time() - t0
            history["train"].append(avg)
            print(f"it {it+1:5d} loss {avg['total']:.3f} cell {avg['ce_cell']:.3f}/{avg['dice_cell']:.3f} "
                  f"nos {avg['ce_nos']:.3f}/{avg['dice_nos']:.3f} ab {avg['bce_ab']:.3f} cn {avg['bce_cn']:.3f} "
                  f"dist {avg.get('mse_dist', 0.0):.4f} lr {lr:.2e} {avg['elapsed_s']:.0f}s", flush=True)
        if held and ((it + 1) % a.val_every == 0 or it + 1 == a.iters):
            model.eval()
            v = quick_val(model, held_img, held_R, device)
            v["iter"] = it + 1
            history["val"].append(v)
            print(f"   val[{a.holdout}] cell F1 {v['cell_det_f1']:.3f} attr {v['cell_attr_both']:.3f} "
                  f"nos F1 {v['nos_det_f1']:.3f} type {v['nos_type_acc']:.3f} (pred {v['n_pred_cells']} / gt {v['n_gt_cells']} cells)", flush=True)
    sampler.close()
    model.eval().cpu()
    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    save_checkpoint(model, a.out, {
        "config": {**vars(a), "ab_neg_weight_used": ab_neg_weight},
        "train_keys": [s.key for s in train],
        "holdout": a.holdout,
        "history": history,
        "labels": list(__import__("cellquant.labels", fromlist=["ALL_LABELS"]).ALL_LABELS),
    })
    with open(os.path.splitext(a.out)[0] + "_history.json", "w") as f:
        json.dump(history, f, indent=1)
    print(f"saved {a.out} ({time.time()-t0:.0f}s)")


if __name__ == "__main__":
    main()