#!/usr/bin/env python3
"""Hard-negative batch sampler based on confusion pairs.

Strategy: each batch contains 50% normal class-balanced samples +
50% samples whose class is in a confusion pair with another. This
forces the loss to separate hardest pairs.

Usage:
    from hard_neg_sampler import ConfusionBatchSampler
    sampler = ConfusionBatchSampler(labels, confusion_pairs, batch_size=64)
    loader = DataLoader(ds, batch_sampler=sampler, ...)
"""
import random
from collections import defaultdict
import numpy as np
import torch
from torch.utils.data import Sampler

class ConfusionBatchSampler(Sampler):
    """Yields batches weighted towards confusion pairs."""
    def __init__(self, labels, confusion_pairs, batch_size=64, hard_frac=0.5,
                 seed=0, num_batches=None):
        self.labels = np.asarray(labels)
        self.batch_size = batch_size
        self.hard_frac = hard_frac
        # Index by class
        self.idx_by_class = defaultdict(list)
        for i, y in enumerate(self.labels):
            self.idx_by_class[int(y)].append(i)
        # Classes in any confusion pair
        self.hard_classes = set()
        for a, b in confusion_pairs:
            self.hard_classes.add(int(a)); self.hard_classes.add(int(b))
        self.hard_classes = [c for c in self.hard_classes if self.idx_by_class[c]]
        self.all_classes = list(self.idx_by_class)
        self.rng = random.Random(seed)
        self.num_batches = num_batches or (len(self.labels) // batch_size)

    def __iter__(self):
        for _ in range(self.num_batches):
            n_hard = int(self.batch_size * self.hard_frac)
            batch = []
            for _ in range(n_hard):
                if not self.hard_classes: break
                c = self.rng.choice(self.hard_classes)
                batch.append(self.rng.choice(self.idx_by_class[c]))
            while len(batch) < self.batch_size:
                c = self.rng.choice(self.all_classes)
                batch.append(self.rng.choice(self.idx_by_class[c]))
            yield batch

    def __len__(self): return self.num_batches

def extract_confusion_pairs(probs, targets, top_n=50):
    """From val probs+targets, find top-N confused (true, pred) pairs."""
    from collections import Counter
    pred = probs.argmax(-1)
    c = Counter()
    for i in range(len(targets)):
        t, p = int(targets[i]), int(pred[i])
        if t != p: c[tuple(sorted((t, p)))] += 1
    return [list(k) for k, _ in c.most_common(top_n)]

if __name__ == '__main__':
    labels = np.random.randint(0, 10, 200)
    pairs = [(0, 1), (2, 3), (4, 5)]
    s = ConfusionBatchSampler(labels, pairs, batch_size=16, num_batches=4)
    for b in s: print(f"batch: {b}")