"""
Train each CNN from scratch on the small demo dataset, a few epochs, to
prove the training pipeline runs end-to-end -- not to reproduce the shipped
checkpoint's accuracy (production trains on ~1e5-1e6 events/thread for
hundreds of epochs; see DeepLearning/Train_{1,2,3}Photon.py).

Pile-up loss matches predictions to ground truth via the permutation
minimizing L2 distance, then applies Smooth L1 -- same idea as
Train_2Photon.py / Train_3Photon.py's set loss, generalized to N=1..3.

Usage: python 04_train_demo.py --photons {1,2,3} [--epochs N]

Needs: torch, numpy. Use the `mlxid_demo` conda env.
"""
import sys
import itertools
import argparse
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F

HERE = Path(__file__).resolve().parent.parent
sys.path.append(str(HERE / 'src'))
from datasets import singlePhotonDataset, doublePhotonDataset, triplePhotonDataset
from model_zoo import new_model

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

DATASET_FOR = {
    1: dict(cls=singlePhotonDataset,
            files=lambda: sorted(str(f) for f in (HERE / 'data/mc_samples').glob('*.npz'))),
    2: dict(cls=doublePhotonDataset,
            files=lambda: [str(HERE / 'data/pileup_samples/2photon_pileup_12keV_demo.npz')]),
    3: dict(cls=triplePhotonDataset,
            files=lambda: [str(HERE / 'data/pileup_samples/3photon_pileup_12keV_demo.npz')]),
}


def set_matched_smooth_l1(pred, gt_xy, beta=0.01):
    """pred: (B, 2*N); gt_xy: (B, N, 2), any order. Matches pred to gt via the
    permutation minimizing L2 distance, then returns Smooth L1 on that matching."""
    B = pred.shape[0]
    n = gt_xy.shape[1]
    pred_xy = pred.view(B, n, 2)
    if n == 1:
        return F.smooth_l1_loss(pred_xy, gt_xy, beta=beta)
    with torch.no_grad():
        best_perm, best_dist = None, None
        for perm in itertools.permutations(range(n)):
            d = ((pred_xy - gt_xy[:, perm, :]) ** 2).sum(dim=(1, 2))
            if best_dist is None:
                best_dist, best_perm = d, torch.tensor(perm, device=pred.device).expand(B, n).clone()
            else:
                better = d < best_dist
                best_dist = torch.where(better, d, best_dist)
                best_perm[better] = torch.tensor(perm, device=pred.device)
    gt_matched = torch.gather(gt_xy, 1, best_perm.unsqueeze(-1).expand(-1, -1, 2))
    return F.smooth_l1_loss(pred_xy, gt_matched, beta=beta)


def get_labels_xy(label, n_photons):
    """Normalize the three dataset classes' differing label layouts to (B, n_photons, 2)."""
    if n_photons == 1:
        return label[:, :2].unsqueeze(1)
    if label.dim() == 3:  ### triplePhotonDataset: (B, n_photons, 4)
        return label[:, :, :2]
    return label.view(label.shape[0], n_photons, -1)[:, :, :2]  ### doublePhotonDataset: (B, n_photons*4)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--photons', type=int, choices=[1, 2, 3], required=True)
    parser.add_argument('--epochs', type=int, default=20)
    parser.add_argument('--batch-size', type=int, default=512)
    parser.add_argument('--lr', type=float, default=5e-4)
    args = parser.parse_args()

    info = DATASET_FOR[args.photons]
    dataset = info['cls'](info['files'](), sampleRatio=1.0, datasetName=f'{args.photons}ph-train-demo')
    n_val = max(1, int(0.1 * len(dataset)))
    train_set, val_set = torch.utils.data.random_split(
        dataset, [len(dataset) - n_val, n_val], generator=torch.Generator().manual_seed(0))
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True)
    val_loader = torch.utils.data.DataLoader(val_set, batch_size=args.batch_size, shuffle=False)

    model = new_model(args.photons, DEVICE)
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.7, patience=5)

    out_dir = HERE / 'models/retrained'
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f'[04] Training {args.photons}-photon model from scratch on {len(train_set)} events '
          f'({n_val} held out for validation), {args.epochs} epochs, device={DEVICE}')
    for epoch in range(1, args.epochs + 1):
        model.train()
        train_loss = 0.0
        for sample, label in train_loader:
            sample, label = sample.to(DEVICE), label.to(DEVICE)
            gt_xy = get_labels_xy(label, args.photons)
            optimizer.zero_grad()
            pred = model(sample)
            loss = set_matched_smooth_l1(pred, gt_xy)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            train_loss += loss.item() * sample.size(0)
        train_loss /= len(train_set)

        model.eval()
        val_sq_err, n_val_pts = 0.0, 0
        with torch.no_grad():
            for sample, label in val_loader:
                sample, label = sample.to(DEVICE), label.to(DEVICE)
                gt_xy = get_labels_xy(label, args.photons)
                pred_xy = model(sample).view(sample.size(0), args.photons, 2)
                if args.photons == 1:
                    d = ((pred_xy - gt_xy) ** 2).sum(dim=(1, 2))
                else:
                    best = None
                    for perm in itertools.permutations(range(args.photons)):
                        dd = ((pred_xy - gt_xy[:, perm, :]) ** 2).sum(dim=(1, 2))
                        best = dd if best is None else torch.minimum(best, dd)
                    d = best
                val_sq_err += d.sum().item()
                n_val_pts += sample.size(0) * args.photons
        val_rms = (val_sq_err / n_val_pts) ** 0.5
        scheduler.step(val_rms)
        print(f'  epoch {epoch:3d}/{args.epochs}  train_loss={train_loss:.5f}  val_RMS={val_rms:.4f} px')

    out_file = out_dir / f'{args.photons}photon_12keV_demo_retrained.pth'
    torch.save(model.state_dict(), out_file)
    print(f'[04] Saved retrained model -> {out_file}')
    print('[04] Compare val_RMS above to 03_eval_mc_truth.py\'s pretrained-checkpoint numbers -- '
          'this demo model is not expected to match production accuracy.')
    if args.photons > 1:
        print(f'[04] Note: val_RMS for {args.photons}-photon pile-up is expected to look noisy for '
              f'many epochs -- production (Train_{args.photons}Photon.py) runs up to 1000 epochs '
              f'on far more data for exactly this reason.')


if __name__ == '__main__':
    main()