"""Metrics computation for threshold calibration.

This module provides utilities for computing quality, cost, and latency
metrics for a given threshold configuration applied to calibration data.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from .config import CalibrationConfig, CostModel, LatencyModel
    from .dataset import CalibrationSample, DatasetSplit

from embergraphrouter.inference.tiered_router import (
    DifficultySignals,
    InferenceTier,
    TieredThresholds,
)


@dataclass
class TierMetrics:
    """Metrics for a single inference tier.

    Attributes:
        tier: The inference tier.
        n_queries: Number of queries routed to this tier.
        fraction: Fraction of total queries.
        accuracy: Accuracy on queries routed to this tier.
        avg_effect: Average effect (quality) achieved.
        oracle_effect: Average oracle effect for these queries.
        regret: Average regret (oracle - achieved).
        avg_cost: Average cost per query.
        avg_latency_ms: Average latency per query.
        top1_match_rate: Rate at which GNN top1 matches oracle.
    """

    tier: InferenceTier
    n_queries: int
    fraction: float
    accuracy: float
    avg_effect: float
    oracle_effect: float
    regret: float
    avg_cost: float
    avg_latency_ms: float
    top1_match_rate: float


@dataclass
class CalibrationMetrics:
    """Complete metrics for a threshold configuration.

    Attributes:
        thresholds: The threshold configuration evaluated.
        n_queries: Total number of queries.
        tier_metrics: Metrics per tier.
        overall_accuracy: Overall accuracy across all tiers.
        overall_effect: Overall average effect.
        overall_regret: Overall regret.
        frac_of_oracle: Fraction of oracle quality achieved.
        total_cost: Total cost for all queries.
        avg_cost_per_query: Average cost per query.
        avg_latency_ms: Average latency per query.
        tier_distribution: Distribution of queries across tiers.
        is_valid: Whether thresholds meet validity constraints.
        validity_issues: List of constraint violations.
    """

    thresholds: TieredThresholds
    n_queries: int
    tier_metrics: dict[InferenceTier, TierMetrics]
    overall_accuracy: float
    overall_effect: float
    overall_regret: float
    frac_of_oracle: float
    total_cost: float
    avg_cost_per_query: float
    avg_latency_ms: float
    tier_distribution: dict[InferenceTier, float]
    is_valid: bool = True
    validity_issues: list[str] = field(default_factory=list)

    def combined_score(
        self,
        quality_weight: float = 0.5,
        cost_weight: float = 0.3,
        latency_weight: float = 0.2,
        normalize: bool = True,
    ) -> float:
        """Compute weighted combined score.

        Higher is better. Quality is used directly, cost and latency
        are inverted (lower is better).

        Args:
            quality_weight: Weight for quality (frac_of_oracle).
            cost_weight: Weight for cost efficiency.
            latency_weight: Weight for latency efficiency.
            normalize: Whether to normalize cost/latency to [0, 1].

        Returns:
            Combined score.
        """
        quality_score = self.frac_of_oracle

        # Cost efficiency: 1 - (cost / max_cost)
        # Assuming max cost is uncertain tier cost
        max_cost = 0.01  # $10 per 1K = $0.01 per query
        if normalize:
            cost_score = 1.0 - min(self.avg_cost_per_query / max_cost, 1.0)
        else:
            cost_score = -self.avg_cost_per_query

        # Latency efficiency: 1 - (latency / max_latency)
        max_latency = 2000.0  # 2 seconds
        if normalize:
            latency_score = 1.0 - min(self.avg_latency_ms / max_latency, 1.0)
        else:
            latency_score = -self.avg_latency_ms

        return (
            quality_weight * quality_score
            + cost_weight * cost_score
            + latency_weight * latency_score
        )

    def to_dict(self) -> dict:
        """Convert to dictionary for serialization."""
        return {
            "thresholds": {
                "trivial_mean_logit": self.thresholds.trivial_mean_logit,
                "trivial_min_prob": self.thresholds.trivial_min_prob,
                "uncertain_mean_logit": self.thresholds.uncertain_mean_logit,
                "uncertain_max_prob": self.thresholds.uncertain_max_prob,
                "uncertain_logit_margin": self.thresholds.uncertain_logit_margin,
            },
            "n_queries": self.n_queries,
            "overall_accuracy": self.overall_accuracy,
            "overall_effect": self.overall_effect,
            "overall_regret": self.overall_regret,
            "frac_of_oracle": self.frac_of_oracle,
            "total_cost": self.total_cost,
            "avg_cost_per_query": self.avg_cost_per_query,
            "avg_latency_ms": self.avg_latency_ms,
            "tier_distribution": {
                tier.name: frac for tier, frac in self.tier_distribution.items()
            },
            "tier_metrics": {
                tier.name: {
                    "n_queries": tm.n_queries,
                    "fraction": tm.fraction,
                    "accuracy": tm.accuracy,
                    "avg_effect": tm.avg_effect,
                    "regret": tm.regret,
                    "avg_cost": tm.avg_cost,
                    "avg_latency_ms": tm.avg_latency_ms,
                }
                for tier, tm in self.tier_metrics.items()
            },
            "is_valid": self.is_valid,
            "validity_issues": self.validity_issues,
        }


class RoutingSimulator:
    """Simulates routing decisions with different thresholds.

    This class replays routing decisions on offline data using
    different threshold configurations without needing to run
    actual inference.
    """

    def __init__(
        self,
        cost_model: "CostModel",
        latency_model: "LatencyModel",
        min_tier_fraction: float = 0.05,
        quality_floor: float = 0.80,
    ):
        """Initialize simulator.

        Args:
            cost_model: Cost model for computing costs.
            latency_model: Latency model for computing latencies.
            min_tier_fraction: Minimum fraction per tier for validity.
            quality_floor: Minimum acceptable quality.
        """
        self.cost_model = cost_model
        self.latency_model = latency_model
        self.min_tier_fraction = min_tier_fraction
        self.quality_floor = quality_floor

    def classify_sample(
        self,
        sample: "CalibrationSample",
        thresholds: TieredThresholds,
    ) -> InferenceTier:
        """Classify a sample into a tier using the given thresholds."""
        signals = DifficultySignals(
            mean_logit=sample.mean_logit,
            max_prob=sample.max_prob,
            logit_margin=sample.logit_margin,
            logit_std=sample.logit_std,
            entropy=sample.entropy,
            top1_model=sample.top1_model,
            top1_idx=0,
            top2_model="",
            top2_idx=1,
            all_probs={},
            all_logits={},
        )
        return signals.classify_tier(thresholds)

    def simulate_query(
        self,
        sample: "CalibrationSample",
        tier: InferenceTier,
    ) -> tuple[float, float, float]:
        """Simulate query execution and return (effect, cost, latency).

        For TRIVIAL: uses cheapest model (may not match top1)
        For DISCRIMINATIVE: uses GNN's top1 model
        For UNCERTAIN: uses ensemble (approximated as oracle effect)

        Args:
            sample: Calibration sample with ground truth.
            tier: Assigned tier.

        Returns:
            Tuple of (achieved_effect, cost, latency_ms).
        """
        tier_name = tier.name.lower()
        cost = self.cost_model.cost_per_query(tier_name)
        latency = self.latency_model.latency_per_query(tier_name)

        if tier == InferenceTier.TRIVIAL:
            # Use cheapest model - assume it's correct for trivial queries
            # Approximate with average effect for "trivial" ground truth queries
            if sample.is_trivial_ground_truth:
                # All models work, so any effect is good
                effect = np.mean(list(sample.ground_truth_effects.values()))
            else:
                # Misclassified as trivial - use worst model as conservative estimate
                effect = min(sample.ground_truth_effects.values())

        elif tier == InferenceTier.DISCRIMINATIVE:
            # Use GNN's top1 model
            effect = sample.ground_truth_effects.get(sample.top1_model, 0.0)

        else:  # UNCERTAIN - ensemble
            # Ensemble approximately achieves oracle or slightly below
            # Use 95% of oracle as conservative estimate
            effect = sample.oracle_effect * 0.95

        return effect, cost, latency

    def evaluate(
        self,
        split: "DatasetSplit",
        thresholds: TieredThresholds,
    ) -> CalibrationMetrics:
        """Evaluate thresholds on a dataset split.

        Args:
            split: Dataset split to evaluate on.
            thresholds: Threshold configuration.

        Returns:
            CalibrationMetrics with full evaluation results.
        """
        # Classify all samples
        tier_assignments: dict[InferenceTier, list["CalibrationSample"]] = {
            tier: [] for tier in InferenceTier
        }
        for sample in split.samples:
            tier = self.classify_sample(sample, thresholds)
            tier_assignments[tier].append(sample)

        # Compute metrics per tier
        tier_metrics = {}
        total_effect = 0.0
        total_oracle = 0.0
        total_cost = 0.0
        total_latency = 0.0
        total_correct = 0

        for tier, samples in tier_assignments.items():
            if not samples:
                tier_metrics[tier] = TierMetrics(
                    tier=tier,
                    n_queries=0,
                    fraction=0.0,
                    accuracy=0.0,
                    avg_effect=0.0,
                    oracle_effect=0.0,
                    regret=0.0,
                    avg_cost=0.0,
                    avg_latency_ms=0.0,
                    top1_match_rate=0.0,
                )
                continue

            effects = []
            oracle_effects = []
            costs = []
            latencies = []
            correct = 0
            top1_matches = 0

            for sample in samples:
                effect, cost, latency = self.simulate_query(sample, tier)
                effects.append(effect)
                oracle_effects.append(sample.oracle_effect)
                costs.append(cost)
                latencies.append(latency)

                # Count as correct if effect >= 0.5
                if effect >= 0.5:
                    correct += 1

                # Track top1 match
                if sample.top1_model == sample.oracle_model:
                    top1_matches += 1

            n = len(samples)
            avg_effect = float(np.mean(effects))
            avg_oracle = float(np.mean(oracle_effects))
            avg_cost = float(np.mean(costs))
            avg_latency = float(np.mean(latencies))

            tier_metrics[tier] = TierMetrics(
                tier=tier,
                n_queries=n,
                fraction=n / len(split),
                accuracy=correct / n,
                avg_effect=avg_effect,
                oracle_effect=avg_oracle,
                regret=avg_oracle - avg_effect,
                avg_cost=avg_cost,
                avg_latency_ms=avg_latency,
                top1_match_rate=top1_matches / n,
            )

            total_effect += sum(effects)
            total_oracle += sum(oracle_effects)
            total_cost += sum(costs)
            total_latency += sum(latencies)
            total_correct += correct

        # Compute overall metrics
        n_total = len(split)
        overall_effect = total_effect / n_total if n_total > 0 else 0.0
        overall_oracle = total_oracle / n_total if n_total > 0 else 0.0
        overall_accuracy = total_correct / n_total if n_total > 0 else 0.0
        overall_regret = overall_oracle - overall_effect
        frac_of_oracle = overall_effect / overall_oracle if overall_oracle > 0 else 0.0

        tier_distribution = {
            tier: tm.fraction for tier, tm in tier_metrics.items()
        }

        # Check validity constraints
        validity_issues = []
        for tier, fraction in tier_distribution.items():
            if fraction < self.min_tier_fraction and fraction > 0:
                validity_issues.append(
                    f"{tier.name} tier has only {fraction:.1%} of queries "
                    f"(minimum: {self.min_tier_fraction:.1%})"
                )

        if overall_accuracy < self.quality_floor:
            validity_issues.append(
                f"Overall accuracy {overall_accuracy:.1%} below floor "
                f"{self.quality_floor:.1%}"
            )

        return CalibrationMetrics(
            thresholds=thresholds,
            n_queries=n_total,
            tier_metrics=tier_metrics,
            overall_accuracy=overall_accuracy,
            overall_effect=overall_effect,
            overall_regret=overall_regret,
            frac_of_oracle=frac_of_oracle,
            total_cost=total_cost,
            avg_cost_per_query=total_cost / n_total if n_total > 0 else 0.0,
            avg_latency_ms=total_latency / n_total if n_total > 0 else 0.0,
            tier_distribution=tier_distribution,
            is_valid=len(validity_issues) == 0,
            validity_issues=validity_issues,
        )


def compute_metrics(
    split: "DatasetSplit",
    thresholds: TieredThresholds,
    config: "CalibrationConfig",
) -> CalibrationMetrics:
    """Convenience function to compute metrics.

    Args:
        split: Dataset split to evaluate.
        thresholds: Threshold configuration.
        config: Calibration configuration.

    Returns:
        CalibrationMetrics.
    """
    simulator = RoutingSimulator(
        cost_model=config.cost_model,
        latency_model=config.latency_model,
        min_tier_fraction=config.min_tier_fraction,
        quality_floor=config.quality_floor,
    )
    return simulator.evaluate(split, thresholds)
