import os
import random
import numpy as np
import torch
import torch.nn as nn


class CropCELoss(nn.Module):
    """
    Cross entropy with hard sample mining + optional class weights.
    """

    def __init__(self, ignore_index=19, hard_mining_ratio=0.5,
                 class_weights=None):
        super().__init__()
        self.ignore_index      = ignore_index
        self.hard_mining_ratio = hard_mining_ratio
        self.criterion = nn.CrossEntropyLoss(
            reduction='none',
            ignore_index=ignore_index,
            weight=class_weights,
        )

    def forward(self, pred, label):
        """
        pred:  (B, C, H, W)
        label: (B, H, W) long
        """
        loss = self.criterion(pred, label)   # (B, H, W)

        if self.hard_mining_ratio >= 1.0:
            valid = label != self.ignore_index
            if valid.sum() == 0:
                return {'loss': loss.mean()}
            return {'loss': loss[valid].mean()}

        # Hard sample mining — top-k% highest loss pixels per image
        B          = loss.shape[0]
        total_loss = 0.
        n_valid    = 0
        for b in range(B):
            valid = label[b] != self.ignore_index
            px    = loss[b][valid]
            if px.numel() == 0:
                continue
            k = max(1, int(px.numel() * self.hard_mining_ratio))
            hard, _ = px.topk(k, largest=True, sorted=False)
            total_loss += hard.mean()
            n_valid    += 1

        if n_valid == 0:
            return {'loss': loss.mean()}
        return {'loss': total_loss / n_valid}


def compute_class_weights(data_root, fold=1, label_channel=0,
                          num_classes=20, ignore_index=19,
                          max_samples=300):
    """
    Compute inverse-frequency class weights from training labels.
    Handles severe class imbalance in PASTIS.
    """
    ann_dir  = os.path.join(data_root, 'ANNOTATIONS')
    all_ids  = sorted([
        f.replace('TARGET_','').replace('.npy','')
        for f in os.listdir(ann_dir)
        if f.startswith('TARGET_')
    ])

    # Build fold map
    fold_map = {}
    for pid in all_ids:
        d = int(pid[0])
        fold_map.setdefault(d, []).append(pid)

    val_fold   = (fold % 4) + 1
    train_ids  = []
    for d in [1,2,3,4]:
        if d != fold and d != val_fold:
            train_ids.extend(fold_map.get(d, []))

    random.shuffle(train_ids)
    train_ids = train_ids[:max_samples]

    counts = np.zeros(num_classes, dtype=np.float64)
    for pid in train_ids:
        label = np.load(os.path.join(ann_dir, f'TARGET_{pid}.npy'))
        label = label[label_channel]
        for c in range(num_classes):
            if c == ignore_index:
                continue
            counts[c] += (label == c).sum()

    counts           = np.maximum(counts, 1)
    weights          = 1.0 / counts
    weights[ignore_index] = 0.0

    # Normalize: mean weight of valid classes = 1
    valid           = np.array([i != ignore_index for i in range(num_classes)])
    weights[valid]  = weights[valid] / weights[valid].mean()

    print(f"\nClass weights (fold={fold}, sampled {len(train_ids)} patches):")
    print(f"  {'Class':<30} {'Count':>10} {'Weight':>8}")
    print(f"  {'─'*30} {'─'*10} {'─'*8}")
    for c in range(num_classes):
        if c == ignore_index:
            continue
        from datasets.pastis_dataset import PASTIS_CLASSES
        name = PASTIS_CLASSES.get(c, f'Class_{c}')
        print(f"  {name:<30} {counts[c]:>10.0f} {weights[c]:>8.3f}")

    return torch.tensor(weights, dtype=torch.float32)