"""Metrics computation utilities for Myanmar Ghost project."""

from dataclasses import dataclass
from typing import Any, Dict, List, Optional

import numpy as np
import torch
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    precision_score,
    recall_score,
    confusion_matrix,
    classification_report,
)


@dataclass
class MetricResult:
    """Result of metric computation."""
    name: str
    value: float
    std: Optional[float] = None


def compute_accuracy(predictions: List[int], targets: List[int]) -> float:
    """Compute accuracy."""
    return accuracy_score(targets, predictions)


def compute_f1(
    predictions: List[int],
    targets: List[int],
    average: str = "weighted",
) -> float:
    """Compute F1 score."""
    return f1_score(targets, predictions, average=average, zero_division=0)


def compute_precision(
    predictions: List[int],
    targets: List[int],
    average: str = "weighted",
) -> float:
    """Compute precision score."""
    return precision_score(targets, predictions, average=average, zero_division=0)


def compute_recall(
    predictions: List[int],
    targets: List[int],
    average: str = "weighted",
) -> float:
    """Compute recall score."""
    return recall_score(targets, predictions, average=average, zero_division=0)


def compute_confusion_matrix(
    predictions: List[int],
    targets: List[int],
) -> np.ndarray:
    """Compute confusion matrix."""
    return confusion_matrix(targets, predictions)


def compute_metrics(
    predictions: List[int],
    targets: List[int],
    class_names: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """Compute all metrics.
    
    Args:
        predictions: List of predicted labels
        targets: List of ground truth labels
        class_names: Optional list of class names
    
    Returns:
        Dictionary of metrics
    """
    metrics = {
        "accuracy": compute_accuracy(predictions, targets),
        "f1_weighted": compute_f1(predictions, targets, "weighted"),
        "f1_macro": compute_f1(predictions, targets, "macro"),
        "f1_micro": compute_f1(predictions, targets, "micro"),
        "precision_weighted": compute_precision(predictions, targets, "weighted"),
        "precision_macro": compute_precision(predictions, targets, "macro"),
        "recall_weighted": compute_recall(predictions, targets, "weighted"),
        "recall_macro": compute_recall(predictions, targets, "macro"),
    }
    
    # Per-class metrics
    labels = list(range(len(class_names))) if class_names else None
    per_class_f1 = f1_score(targets, predictions, labels=labels, average=None, zero_division=0)
    
    if class_names:
        for i, name in enumerate(class_names):
            metrics[f"f1_{name}"] = per_class_f1[i]
    
    return metrics


class MetricsTracker:
    """Track metrics during training."""
    
    def __init__(
        self,
        metrics: List[str] = None,
        class_names: Optional[List[str]] = None,
    ):
        self.metrics = metrics or ["loss", "accuracy", "f1"]
        self.class_names = class_names or ["negative", "neutral", "positive", "sarcastic"]
        
        self.history = {m: [] for m in self.metrics}
        self.best_values = {m: float("-inf") for m in self.metrics}
        self.best_epochs = {m: 0 for m in self.metrics}
    
    def update(self, metrics: Dict[str, float], step: int) -> None:
        """Update metrics at current step."""
        for name, value in metrics.items():
            if name in self.metrics:
                self.history[name].append((step, value))
                
                # Track best
                if value > self.best_values[name]:
                    self.best_values[name] = value
                    self.best_epochs[name] = step
    
    def get_current(self, metric_name: str) -> float:
        """Get current value of a metric."""
        if metric_name in self.history and self.history[metric_name]:
            return self.history[metric_name][-1][1]
        return 0.0
    
    def get_best(self, metric_name: str) -> tuple:
        """Get best value and epoch of a metric."""
        return self.best_values.get(metric_name, 0), self.best_epochs.get(metric_name, 0)
    
    def get_summary(self) -> Dict[str, Any]:
        """Get summary of all metrics."""
        return {
            "best": self.best_values,
            "best_epochs": self.best_epochs,
            "current": {m: self.get_current(m) for m in self.metrics},
        }


def compute_bleu(
    predictions: List[str],
    references: List[str],
) -> float:
    """Compute BLEU score for text generation."""
    from sacrebleu import sentence_bleu
    
    scores = []
    for pred, ref in zip(predictions, references):
        score = sentence_bleu(pred, [ref])
        scores.append(score.score)
    
    return np.mean(scores)


def compute_perplexity(
    loss: float,
) -> float:
    """Compute perplexity from cross-entropy loss."""
    return np.exp(loss)


if __name__ == "__main__":
    # Test metrics computation
    predictions = [0, 1, 2, 0, 1, 2, 0, 1, 2]
    targets = [0, 1, 2, 0, 1, 1, 0, 0, 2]
    
    metrics = compute_metrics(predictions, targets)
    print("Metrics:", metrics)