"""Uncertainty sampling for active learning.

Selects samples where the model has lowest confidence,
indicating areas where human annotation would be most valuable.
"""

import json
import logging
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

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

logger = logging.getLogger(__name__)


class UncertaintyMethod(str, Enum):
    """Methods for calculating uncertainty."""
    LEAST_CONFIDENCE = "least_confidence"  # 1 - max probability
    MARGIN = "margin"  # Difference between top 2 probabilities
    ENTROPY = "entropy"  # Shannon entropy
    RATIO = "ratio"  # Ratio of top to second probability
    VARIANCE = "variance"  # Prediction variance (ensemble)


@dataclass
class UncertaintySample:
    """Sample with uncertainty score."""
    sample_id: str
    text: str
    uncertainty_score: float
    predicted_class: str
    predicted_prob: float
    second_prob: float = 0.0
    metadata: Dict = None
    
    def to_dict(self) -> Dict:
        return {
            "sample_id": self.sample_id,
            "text": self.text,
            "uncertainty_score": self.uncertainty_score,
            "predicted_class": self.predicted_class,
            "predicted_prob": self.predicted_prob,
            "second_prob": self.second_prob,
            "metadata": self.metadata or {},
        }


class PredictionDataset(Dataset):
    """Dataset for model predictions."""
    
    def __init__(self, samples: List[Dict], tokenizer, max_length: int = 128):
        self.samples = samples
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self) -> int:
        return len(self.samples)
    
    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, str]:
        sample = self.samples[idx]
        text = sample.get("text", "")
        
        encoding = self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_length,
            padding="max_length",
            return_tensors="pt",
        )
        
        return (
            encoding["input_ids"].squeeze(0),
            encoding["attention_mask"].squeeze(0),
            sample.get("id", f"sample_{idx}"),
        )


class UncertaintySampler:
    """Sample uncertain instances for active learning."""
    
    def __init__(
        self,
        model: nn.Module,
        method: UncertaintyMethod = UncertaintyMethod.LEAST_CONFIDENCE,
        device: str = "cuda" if torch.cuda.is_available() else "cpu",
    ):
        self.model = model
        self.method = method
        self.device = device
        self.model.to(device)
        self.model.eval()
    
    def _compute_uncertainty(
        self,
        logits: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Compute uncertainty scores from logits.
        
        Returns:
            uncertainty: Uncertainty scores
            predicted_class: Predicted class indices
            probs: Probability distributions
            second_probs: Second highest probabilities (for margin)
        """
        probs = torch.softmax(logits, dim=-1)
        
        if self.method == UncertaintyMethod.ENTROPY:
            # Shannon entropy: -sum(p * log(p))
            entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=-1)
            uncertainty = entropy
        elif self.method == UncertaintyMethod.LEAST_CONFIDENCE:
            # 1 - max probability
            max_prob, _ = probs.max(dim=-1)
            uncertainty = 1 - max_prob
        elif self.method == UncertaintyMethod.MARGIN:
            # Difference between top 2 probabilities
            sorted_probs, _ = torch.sort(probs, dim=-1, descending=True)
            margin = sorted_probs[:, 0] - sorted_probs[:, 1]
            uncertainty = 1 - margin
        elif self.method == UncertaintyMethod.RATIO:
            # Ratio of top to second probability
            sorted_probs, _ = torch.sort(probs, dim=-1, descending=True)
            ratio = sorted_probs[:, 1] / (sorted_probs[:, 0] + 1e-10)
            uncertainty = ratio
        else:
            raise ValueError(f"Unknown method: {self.method}")
        
        predicted_class = probs.argmax(dim=-1)
        max_probs = probs.max(dim=-1).values
        
        # Second highest probability
        sorted_probs_detached = probs.detach().cpu()
        sorted_indices = torch.argsort(sorted_probs_detached, dim=-1, descending=True)
        second_probs = torch.gather(
            probs, 1, sorted_indices[:, 1:2]
        ).squeeze(-1)
        
        return uncertainty, predicted_class, max_probs, second_probs
    
    def score_samples(
        self,
        samples: List[Dict],
        tokenizer,
        batch_size: int = 32,
    ) -> List[UncertaintySample]:
        """Score samples by uncertainty."""
        dataset = PredictionDataset(samples, tokenizer)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
        
        all_uncertainty = []
        all_predicted = []
        all_probs = []
        all_second_probs = []
        all_ids = []
        
        with torch.no_grad():
            for input_ids, attention_mask, sample_ids in tqdm(
                dataloader, desc="Computing uncertainty"
            ):
                input_ids = input_ids.to(self.device)
                attention_mask = attention_mask.to(self.device)
                
                outputs = self.model(input_ids, attention_mask)
                logits = outputs.logits if hasattr(outputs, "logits") else outputs
                
                uncertainty, pred_class, probs, second_probs = self._compute_uncertainty(logits)
                
                all_uncertainty.extend(uncertainty.cpu().tolist())
                all_predicted.extend(pred_class.cpu().tolist())
                all_probs.extend(probs.cpu().tolist())
                all_second_probs.extend(second_probs.cpu().tolist())
                all_ids.extend(sample_ids)
        
        # Create uncertainty samples
        uncertain_samples = []
        class_names = ["negative", "neutral", "positive", "sarcastic"]
        
        for i, sample in enumerate(samples):
            us = UncertaintySample(
                sample_id=all_ids[i],
                text=sample.get("text", ""),
                uncertainty_score=all_uncertainty[i],
                predicted_class=class_names[all_predicted[i]] if all_predicted[i] < len(class_names) else "unknown",
                predicted_prob=all_probs[i],
                second_prob=all_second_probs[i],
                metadata=sample,
            )
            uncertain_samples.append(us)
        
        return uncertain_samples
    
    def select_samples(
        self,
        samples: List[Dict],
        tokenizer,
        n_samples: int = 100,
        batch_size: int = 32,
        exclude_ids: Optional[List[str]] = None,
    ) -> List[UncertaintySample]:
        """Select most uncertain samples for annotation.
        
        Args:
            samples: List of samples to score
            tokenizer: Tokenizer for the model
            n_samples: Number of samples to select
            batch_size: Batch size for inference
            exclude_ids: Sample IDs to exclude (already annotated)
        
        Returns:
            List of selected uncertain samples, sorted by uncertainty
        """
        # Filter out already annotated
        if exclude_ids:
            samples = [s for s in samples if s.get("id") not in exclude_ids]
        
        # Score all samples
        uncertain_samples = self.score_samples(samples, tokenizer, batch_size)
        
        # Sort by uncertainty (highest first)
        uncertain_samples.sort(key=lambda x: x.uncertainty_score, reverse=True)
        
        # Select top n_samples
        selected = uncertain_samples[:n_samples]
        
        logger.info(
            f"Selected {len(selected)} most uncertain samples "
            f"(uncertainty range: {selected[0].uncertainty_score:.4f} - "
            f"{selected[-1].uncertainty_score:.4f})"
        )
        
        return selected
    
    def diversity_sample(
        self,
        samples: List[Dict],
        tokenizer,
        n_samples: int = 100,
        batch_size: int = 32,
        n_clusters: int = 10,
    ) -> List[UncertaintySample]:
        """Select diverse uncertain samples using clustering.
        
        Combines uncertainty with diversity to avoid selecting
        similar samples.
        """
        from sklearn.cluster import MiniBatchKMeans
        from sklearn.feature_extraction.text import TfidfVectorizer
        
        # Score samples
        uncertain_samples = self.score_samples(samples, tokenizer, batch_size)
        
        # Create embeddings for clustering
        vectorizer = TfidfVectorizer(max_features=1000)
        texts = [s.text for s in samples]
        embeddings = vectorizer.fit_transform(texts)
        
        # Cluster
        kmeans = MiniBatchKMeans(n_clusters=n_clusters, random_state=42)
        cluster_labels = kmeans.fit_predict(embeddings)
        
        # Select from each cluster
        selected = []
        for cluster_id in range(n_clusters):
            cluster_indices = [
                i for i, label in enumerate(cluster_labels)
                if label == cluster_id
            ]
            cluster_uncertain = [
                uncertain_samples[i] for i in cluster_indices
            ]
            cluster_uncertain.sort(key=lambda x: x.uncertainty_score, reverse=True)
            
            # Take top samples from each cluster
            n_per_cluster = max(1, n_samples // n_clusters)
            selected.extend(cluster_uncertain[:n_per_cluster])
        
        # Sort by uncertainty
        selected.sort(key=lambda x: x.uncertainty_score, reverse=True)
        
        return selected[:n_samples]
    
    def batch_sample(
        self,
        samples: List[Dict],
        tokenizer,
        strategy: str = "greedy",
        n_samples: int = 100,
        batch_size: int = 32,
    ) -> List[UncertaintySample]:
        """Sample using batch mode for efficiency.
        
        Strategies:
            - greedy: Select top n_samples by uncertainty
            - diverse: Cluster-based diverse sampling
            - random: Random baseline
        """
        if strategy == "random":
            import random
            random.seed(42)
            indices = random.sample(range(len(samples)), min(n_samples, len(samples)))
            return [UncertaintySample(
                sample_id=samples[i].get("id", f"sample_{i}"),
                text=samples[i].get("text", ""),
                uncertainty_score=0.0,
                predicted_class="unknown",
                predicted_prob=0.0,
            ) for i in indices]
        
        elif strategy == "diverse":
            return self.diversity_sample(
                samples, tokenizer, n_samples, batch_size
            )
        
        else:  # greedy
            return self.select_samples(
                samples, tokenizer, n_samples, batch_size
            )


def save_selected_samples(
    samples: List[UncertaintySample],
    output_path: str,
) -> None:
    """Save selected samples to JSON file."""
    output_data = [s.to_dict() for s in samples]
    
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(output_data, f, indent=2, ensure_ascii=False)
    
    logger.info(f"Saved {len(samples)} samples to {output_path}")


if __name__ == "__main__":
    print("UncertaintySampler module loaded")
    print(f"Available methods: {[m.value for m in UncertaintyMethod]}")