#!/usr/bin/env python3
"""Auto-merge highly confused class pairs.

Reads probs from v4 eval; finds pairs where >=30% of true-class-A samples
are predicted class-B (or vice versa). Merges A+B into A|B compound class.

Output: manifest with unified_label replaced by merged form, plus label_merge_map.json.
"""
import json, argparse
from pathlib import Path
from collections import defaultdict, Counter
import torch


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--probs', required=True, help='pt with probs/targets/label_to_idx')
    ap.add_argument('--manifest', required=True)
    ap.add_argument('--output-manifest', required=True)
    ap.add_argument('--output-map', required=True)
    ap.add_argument('--confusion-frac', type=float, default=0.25,
                    help='Merge if A→B confusion >= frac of A samples')
    ap.add_argument('--min-sym-frac', type=float, default=0.15,
                    help='Also require B→A confusion >= frac (symmetric)')
    args = ap.parse_args()

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

    pred = probs.argmax(-1)
    # Build confusion matrix
    C = len(label_to_idx)
    confusion = torch.zeros(C, C, dtype=torch.long)
    for i in range(len(targets)):
        confusion[int(targets[i]), int(pred[i])] += 1

    # Row sums (per-class n)
    n_per = confusion.sum(-1).tolist()

    # Find confusion pairs
    merge_pairs = []
    for a in range(C):
        if n_per[a] < 5: continue
        for b in range(a+1, C):
            if n_per[b] < 5: continue
            ab = confusion[a, b].item() / max(1, n_per[a])
            ba = confusion[b, a].item() / max(1, n_per[b])
            if ab >= args.confusion_frac and ba >= args.min_sym_frac:
                merge_pairs.append((a, b, ab, ba))
            elif ba >= args.confusion_frac and ab >= args.min_sym_frac:
                merge_pairs.append((b, a, ba, ab))

    # Group into merge groups via union-find
    parent = list(range(C))
    def find(x):
        while parent[x] != x: x = parent[x]
        return x
    for a, b, _, _ in merge_pairs:
        ra, rb = find(a), find(b)
        if ra != rb: parent[min(ra, rb)] = max(ra, rb)  # keep larger idx as root
    groups = defaultdict(list)
    for i in range(C): groups[find(i)].append(i)

    # Build merge_map: label → merged_label
    merge_map = {}
    for root, members in groups.items():
        if len(members) == 1:
            lab = idx_to_label[root]
            merge_map[lab] = lab
        else:
            # Sort by n desc, compound name = "A|B|C" (alphabetic stable)
            sorted_labs = sorted([idx_to_label[m] for m in members])
            merged = '|'.join(sorted_labs)
            for m in members:
                merge_map[idx_to_label[m]] = merged

    # Apply to manifest
    n_merged = sum(1 for v in merge_map.values() if '|' in v)
    n_groups = len(set(merge_map.values()))
    print(f"Confusion pairs found: {len(merge_pairs)}")
    print(f"Original classes: {C}, After merge: {n_groups}, Labels in merged groups: {n_merged}")

    records = [json.loads(l) for l in open(args.manifest)]
    with open(args.output_manifest, 'w') as f:
        for r in records:
            lab = r.get('unified_label')
            if lab and lab in merge_map:
                r['unified_label'] = merge_map[lab]
                r['original_label'] = lab
            f.write(json.dumps(r) + '\n')

    Path(args.output_map).parent.mkdir(parents=True, exist_ok=True)
    json.dump({'merge_map': merge_map, 'n_original': C, 'n_merged': n_groups,
                'pairs': [{'a': idx_to_label[a], 'b': idx_to_label[b],
                            'ab': ab, 'ba': ba} for a, b, ab, ba in merge_pairs]},
               open(args.output_map, 'w'), indent=2)
    print(f"Saved: {args.output_manifest}, map: {args.output_map}")

if __name__ == '__main__':
    main()