#!/usr/bin/env python3
"""Export margin-low samples for expert review.

Rows to review (margin = p_top1 - p_top2):
  margin < 0.2 → most informative
  entropy > 2.5 → high uncertainty
  flipped by pair-MLP → suspected misclassification

Output: CSV + image symlinks in review folder.
"""
import os, json, csv, argparse, shutil
from pathlib import Path

import torch

ROOT = Path("/arf/scratch/stakan/hitit-proje")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--probs', required=True)
    ap.add_argument('--manifest', required=True)
    ap.add_argument('--label-to-idx', required=True,
                    help='any train ckpt providing label_to_idx')
    ap.add_argument('--val-fold', type=int, default=0)
    ap.add_argument('--margin-threshold', type=float, default=0.2)
    ap.add_argument('--entropy-threshold', type=float, default=2.5)
    ap.add_argument('--output-csv', required=True)
    ap.add_argument('--output-dir', default=None,
                    help='If set, symlink low-margin images here for UI review')
    args = ap.parse_args()

    d = torch.load(args.probs, map_location='cpu', weights_only=False)
    probs, targets = d['probs'], d['targets']
    ck = torch.load(args.label_to_idx, map_location='cpu', weights_only=False)
    label_to_idx = ck['label_to_idx']; idx_to_label = {v: k for k, v in label_to_idx.items()}

    # Respect min_samples filter (same as prototype_net / training dataset)
    from collections import Counter
    cls_count = Counter()
    for line in open(args.manifest):
        r = json.loads(line)
        if r.get('task') != 'classification' or not r.get('unified_label'): continue
        cls_count[r['unified_label']] += 1
    MIN_SAMPLES = 10  # same default used by training/prototype_net
    records = []
    with open(args.manifest) as f:
        for line in f:
            r = json.loads(line)
            if r.get('task') != 'classification': continue
            if not r.get('unified_label') or r['unified_label'] not in label_to_idx: continue
            if not r.get('path') or r.get('storage') != 'fs': continue
            if r.get('integrity_ok') is False: continue
            if cls_count[r['unified_label']] < MIN_SAMPLES: continue
            if r.get('tablet_view_fold', 0) != args.val_fold: continue
            records.append(r)
    if len(records) != probs.size(0):
        print(f"WARNING: records={len(records)} vs probs={probs.size(0)}; truncating to min")
        records = records[:probs.size(0)]
        probs = probs[:len(records)]
        targets = targets[:len(records)]

    top2 = probs.topk(2, dim=-1)
    margin = top2.values[:, 0] - top2.values[:, 1]
    entropy = -(probs * probs.clamp_min(1e-9).log()).sum(-1)

    rows = []
    for i in range(len(records)):
        if margin[i] < args.margin_threshold or entropy[i] > args.entropy_threshold:
            rows.append({
                'idx': i,
                'path': records[i]['path'],
                'true_label': records[i]['unified_label'],
                'pred_label': idx_to_label[int(top2.indices[i, 0])],
                'alt_label': idx_to_label[int(top2.indices[i, 1])],
                'p_pred': float(top2.values[i, 0]),
                'p_alt': float(top2.values[i, 1]),
                'margin': float(margin[i]),
                'entropy': float(entropy[i]),
                'misclassified': int(top2.indices[i, 0]) != int(targets[i]),
            })

    rows.sort(key=lambda r: r['margin'])
    Path(args.output_csv).parent.mkdir(parents=True, exist_ok=True)
    with open(args.output_csv, 'w') as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()) if rows else ['idx'])
        w.writeheader()
        w.writerows(rows)
    print(f"Exported {len(rows)} rows for review → {args.output_csv}")

    if args.output_dir:
        d = Path(args.output_dir); d.mkdir(parents=True, exist_ok=True)
        for r in rows[:300]:  # cap for sanity
            src = Path(r['path']); dst = d / f"{r['idx']:05d}_{r['true_label']}_vs_{r['pred_label']}.png"
            if src.exists() and not dst.exists():
                try: shutil.copy(src, dst)
                except Exception: pass

if __name__ == '__main__':
    main()