"""Train the binary edge classifier on the HSS-aligned dataset.

Uses the ClassificationPointNet architecture (fast_pointnet_class.py, adapted
from the 2025 first-place solution). Loads the per-sample .npz files produced by
gen_edge_dataset.py, balances the classes with a WeightedRandomSampler,
optionally augments (z-rotation + light isotropic scale), and saves the
best-AUC checkpoint.
"""
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

import argparse
import glob
import json
import math
import time

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
from tqdm import tqdm


def _augment_gpu(x):
    """In-place z-rotation + isotropic scale on a (B, 6, N) tensor."""
    B = x.shape[0]
    theta = torch.rand(B, device=x.device) * (2 * math.pi)
    cos_t = torch.cos(theta).view(B, 1, 1)
    sin_t = torch.sin(theta).view(B, 1, 1)
    x0 = x[:, 0:1, :].clone()
    x1 = x[:, 1:2, :].clone()
    x[:, 0:1, :] = cos_t * x0 - sin_t * x1
    x[:, 1:2, :] = sin_t * x0 + cos_t * x1
    scale = torch.empty(B, 1, 1, device=x.device).uniform_(0.95, 1.05)
    x[:, :3, :] *= scale
    return x


def _make_batch(patches_np, labels_np, idx, device):
    """Slice numpy arrays by idx, transpose to (B, 6, N), send to GPU."""
    x = torch.from_numpy(np.ascontiguousarray(patches_np[idx])).to(device, non_blocking=True)
    x = x.transpose(1, 2).contiguous()  # (B, 6, N)
    y = torch.from_numpy(labels_np[idx]).to(device, non_blocking=True)
    return x, y

from fast_pointnet_class import ClassificationPointNet


class EdgeDataset(Dataset):
    """Pre-loads all patches into RAM (.npz files concatenated)."""

    def __init__(self, files, augment=False, dtype=np.float32):
        self.augment = augment
        patches_chunks, labels_chunks = [], []
        for f in files:
            data = np.load(f)
            patches_chunks.append(data['patches'].astype(dtype))
            labels_chunks.append(data['labels'].astype(np.float32))
        self.patches = np.concatenate(patches_chunks, axis=0)
        self.labels = np.concatenate(labels_chunks, axis=0)
        size_gb = self.patches.nbytes / 1e9
        print(f"Loaded {len(self.labels)} patches, {size_gb:.2f} GB in RAM")

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        patch = self.patches[idx]
        label = self.labels[idx]
        if self.augment:
            patch = patch.copy()
            # Z-axis rotation (preserves vertical)
            theta = np.random.uniform(0, 2 * np.pi)
            cos, sin = np.cos(theta), np.sin(theta)
            xyz = patch[:, :3]
            patch[:, 0] = cos * xyz[:, 0] - sin * xyz[:, 1]
            patch[:, 1] = sin * xyz[:, 0] + cos * xyz[:, 1]
            # Light isotropic scale on positions only
            scale = np.random.uniform(0.95, 1.05)
            patch[:, :3] *= scale
        # PointNet expects (channels, points)
        return (torch.from_numpy(patch.T.astype(np.float32)),
                torch.tensor(label, dtype=torch.float32))


def compute_auc(labels, scores):
    """Simple rank-based AUC (no sklearn dep)."""
    order = np.argsort(scores)
    sorted_labels = labels[order]
    n_pos = sorted_labels.sum()
    n_neg = len(sorted_labels) - n_pos
    if n_pos == 0 or n_neg == 0:
        return 0.5
    ranks = np.arange(1, len(sorted_labels) + 1)
    rank_sum_pos = ranks[sorted_labels == 1].sum()
    return float((rank_sum_pos - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data-dir', required=True,
                        help='Directory of .npz files (gen_edge_dataset.py output)')
    parser.add_argument('--out', default='pnet_class_2026.pth',
                        help='Best-AUC checkpoint path')
    parser.add_argument('--val-frac', type=float, default=0.05)
    parser.add_argument('--epochs', type=int, default=10)
    parser.add_argument('--batch', type=int, default=128)
    parser.add_argument('--lr', type=float, default=1e-3)
    parser.add_argument('--weight-decay', type=float, default=1e-4)
    parser.add_argument('--balanced-sampling', action='store_true',
                        help='Sample positives/negatives equally each batch')
    parser.add_argument('--augment', action='store_true',
                        help='Random z-rotation + small scale on training patches')
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--num-workers', type=int, default=8)
    parser.add_argument('--max-samples', type=int, default=None,
                        help='Limit number of sample files (smoke test)')
    args = parser.parse_args()

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")
    print(f"Args: {vars(args)}\n")

    # File split
    all_files = sorted(glob.glob(os.path.join(args.data_dir, '*.npz')))
    all_files = [f for f in all_files if 'manifest' not in os.path.basename(f)]
    if args.max_samples:
        all_files = all_files[:args.max_samples]
    print(f"Found {len(all_files)} sample files")

    rng = np.random.RandomState(args.seed)
    perm = rng.permutation(len(all_files))
    all_files = [all_files[i] for i in perm]
    n_val = max(1, int(args.val_frac * len(all_files)))
    val_files = all_files[:n_val]
    train_files = all_files[n_val:]
    print(f"Train: {len(train_files)} files, Val: {len(val_files)} files")

    print("\n--- Loading train ---")
    train_ds = EdgeDataset(train_files, augment=args.augment)
    print("\n--- Loading val ---")
    val_ds = EdgeDataset(val_files, augment=False)

    train_pos = int(train_ds.labels.sum())
    train_neg = len(train_ds.labels) - train_pos
    print(f"\nTrain pos/neg: {train_pos}/{train_neg} "
          f"({100*train_pos/len(train_ds.labels):.1f}% positive)")

    # Skip DataLoader entirely -- dataset is in RAM, IPC overhead with workers
    # was the actual bottleneck (~700ms/batch). Direct numpy -> GPU per batch.
    train_patches_np = train_ds.patches    # (N, max_pts, 6) float32
    train_labels_np = train_ds.labels.astype(np.float32)
    val_patches_np = val_ds.patches
    val_labels_np = val_ds.labels.astype(np.float32)
    n_train = len(train_labels_np)
    n_val = len(val_labels_np)
    print(f"Train batches/epoch: {(n_train + args.batch - 1)//args.batch}")

    model = ClassificationPointNet(input_dim=6, max_points=1024).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"\nModel: {n_params:,} params")

    opt = torch.optim.AdamW(model.parameters(), lr=args.lr,
                             weight_decay=args.weight_decay)
    if args.balanced_sampling:
        # Reuse the flag name: now means "weight loss for class balance" (via pos_weight)
        # instead of "use WeightedRandomSampler". Same effect, no per-batch cost.
        pos_weight = torch.tensor([train_neg / max(train_pos, 1)], device=device)
        print(f"BCE pos_weight = {pos_weight.item():.3f} (class balance)")
        bce = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
    else:
        bce = nn.BCEWithLogitsLoss()
    use_amp = (device.type == 'cuda')
    scaler = torch.cuda.amp.GradScaler() if use_amp else None

    metrics_path = args.out.replace('.pth', '_metrics.jsonl')
    # Truncate previous run's metrics if present so `tail -f` shows this run only.
    open(metrics_path, 'w').close()
    print(f"Metrics log (peek with `tail -f {metrics_path}`)\n")

    best_auc = 0.5
    val_batch = args.batch * 2
    for epoch in range(args.epochs):
        # Train
        model.train()
        t0 = time.time()
        perm = np.random.permutation(n_train)
        n_train_batches = (n_train + args.batch - 1) // args.batch
        train_loss, n = 0.0, 0
        pbar = tqdm(range(n_train_batches),
                    desc=f"epoch {epoch+1}/{args.epochs}",
                    leave=False, dynamic_ncols=True)
        for bi in pbar:
            s, e = bi * args.batch, min((bi + 1) * args.batch, n_train)
            idx = perm[s:e]
            x, y = _make_batch(train_patches_np, train_labels_np, idx, device)
            if args.augment:
                x = _augment_gpu(x)
            opt.zero_grad()
            if use_amp:
                with torch.cuda.amp.autocast():
                    logits = model(x).squeeze(-1)
                    loss = bce(logits, y)
                scaler.scale(loss).backward()
                scaler.step(opt)
                scaler.update()
            else:
                logits = model(x).squeeze(-1)
                loss = bce(logits, y)
                loss.backward()
                opt.step()
            bs = e - s
            train_loss += loss.item() * bs
            n += bs
            if n > 0:
                pbar.set_postfix(loss=f"{train_loss / n:.4f}")
        pbar.close()
        train_loss /= max(n, 1)

        # Validate
        model.eval()
        all_s, all_y = [], []
        with torch.no_grad():
            n_val_batches = (n_val + val_batch - 1) // val_batch
            for bi in range(n_val_batches):
                s, e = bi * val_batch, min((bi + 1) * val_batch, n_val)
                idx = np.arange(s, e)
                x, y = _make_batch(val_patches_np, val_labels_np, idx, device)
                logits = model(x).squeeze(-1)
                all_s.append(torch.sigmoid(logits).cpu().numpy())
                all_y.append(y.cpu().numpy())
        all_s = np.concatenate(all_s)
        all_y = np.concatenate(all_y)
        auc = compute_auc(all_y, all_s)
        acc50 = float(((all_s > 0.5) == all_y).mean())
        # Score distribution
        pos_mean = float(all_s[all_y == 1].mean()) if (all_y == 1).any() else 0.0
        neg_mean = float(all_s[all_y == 0].mean()) if (all_y == 0).any() else 0.0

        dt = time.time() - t0
        is_best = auc > best_auc
        print(f"epoch {epoch+1:2d}/{args.epochs} "
              f"loss={train_loss:.4f} val_auc={auc:.4f} acc@0.5={acc50:.3f} "
              f"score(pos)={pos_mean:.3f} score(neg)={neg_mean:.3f} "
              f"t={dt:.0f}s{' *' if is_best else ''}")

        with open(metrics_path, 'a') as fh:
            fh.write(json.dumps({
                'epoch': epoch + 1,
                'loss': train_loss,
                'val_auc': auc,
                'acc_at_0.5': acc50,
                'pos_mean': pos_mean,
                'neg_mean': neg_mean,
                'best': is_best,
                'wall_sec': dt,
            }) + '\n')

        if is_best:
            best_auc = auc
            torch.save({
                'model_state_dict': model.state_dict(),
                'epoch': epoch,
                'auc': auc,
                'config': vars(args),
            }, args.out)
            print(f"  -> saved {args.out} (best AUC)")

    print(f"\nBest val AUC: {best_auc:.4f}")
    print(f"Checkpoint: {args.out}")


if __name__ == '__main__':
    main()