"""Configuration for threshold calibration.

This module defines the configuration structures for calibration experiments,
including cost models, latency models, and optimization objectives.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Mapping


class OptimizationTarget(Enum):
    """Optimization objective for threshold calibration."""

    QUALITY = "quality"  # Maximize accuracy/quality only
    COST = "cost"  # Minimize cost only
    LATENCY = "latency"  # Minimize latency only
    BALANCED = "balanced"  # Weighted combination
    PARETO = "pareto"  # Find Pareto frontier


@dataclass(frozen=True, slots=True)
class CostModel:
    """Cost model for different inference tiers.

    Costs are in USD per 1000 queries. These should be calibrated to your
    actual provider pricing.

    Attributes:
        trivial_cost: Cost per 1K queries for TRIVIAL tier (cheap model).
        discriminative_cost: Cost per 1K queries for DISCRIMINATIVE tier.
        uncertain_cost: Cost per 1K queries for UNCERTAIN tier (ensemble).
        model_costs: Optional per-model cost overrides.
    """

    trivial_cost: float = 0.10  # $0.10 per 1K (Gemini Flash)
    discriminative_cost: float = 1.00  # $1.00 per 1K (Gemini Pro)
    uncertain_cost: float = 10.00  # $10.00 per 1K (EB1 ensemble)
    model_costs: Mapping[str, float] = field(default_factory=dict)

    def cost_per_query(self, tier: str) -> float:
        """Get cost per query for a tier."""
        tier_lower = tier.lower()
        if tier_lower == "trivial":
            return self.trivial_cost / 1000
        elif tier_lower == "discriminative":
            return self.discriminative_cost / 1000
        elif tier_lower == "uncertain":
            return self.uncertain_cost / 1000
        else:
            raise ValueError(f"Unknown tier: {tier}")


@dataclass(frozen=True, slots=True)
class LatencyModel:
    """Latency model for different inference tiers.

    Latencies are in milliseconds (p50). These should be calibrated to your
    actual infrastructure.

    Attributes:
        trivial_latency_ms: P50 latency for TRIVIAL tier.
        discriminative_latency_ms: P50 latency for DISCRIMINATIVE tier.
        uncertain_latency_ms: P50 latency for UNCERTAIN tier.
        routing_overhead_ms: Additional latency for routing decision.
    """

    trivial_latency_ms: float = 80.0  # Fast model
    discriminative_latency_ms: float = 150.0  # Strong model
    uncertain_latency_ms: float = 1500.0  # Ensemble (parallel candidates + judge)
    routing_overhead_ms: float = 50.0  # GNN inference + embedding

    def latency_per_query(self, tier: str) -> float:
        """Get latency per query for a tier (including routing overhead)."""
        tier_lower = tier.lower()
        if tier_lower == "trivial":
            return self.trivial_latency_ms + self.routing_overhead_ms
        elif tier_lower == "discriminative":
            return self.discriminative_latency_ms + self.routing_overhead_ms
        elif tier_lower == "uncertain":
            return self.uncertain_latency_ms + self.routing_overhead_ms
        else:
            raise ValueError(f"Unknown tier: {tier}")


@dataclass(frozen=True, slots=True)
class ThresholdBounds:
    """Bounds for threshold search space.

    Defines the valid ranges for each threshold parameter during optimization.
    These bounds are intentionally WIDE to explore the full cost-quality tradeoff:
    - Cost-minimizing: low trivial threshold → route most queries to TRIVIAL
    - Quality-maximizing: high uncertain threshold → route most to UNCERTAIN

    Signal distribution (from eb1-gnn-preview-12-19):
    - mean_logit: mean=0.15, std=0.89, range=[-2.2, 1.6]
    - logit_margin: mean=0.10, std=0.10, range=[0, 0.5+]

    Threshold effects on tier distribution:
    - trivial_mean_logit = -1.5 → 97% TRIVIAL (cost: ~$0.10/1K)
    - trivial_mean_logit = +1.0 → 17% TRIVIAL (cost: ~$5-10/1K)
    """

    # TRIVIAL tier: mean_logit threshold
    # Lower = more TRIVIAL routing = lower cost
    # Range spans from "almost everything TRIVIAL" to "almost nothing TRIVIAL"
    trivial_mean_logit_min: float = -1.5  # 97% TRIVIAL (aggressive cost saving)
    trivial_mean_logit_max: float = 1.5   # 6% TRIVIAL (quality focused)

    # UNCERTAIN tier: mean_logit threshold
    # Higher = more UNCERTAIN routing = higher quality
    uncertain_mean_logit_min: float = -2.0  # Almost nothing UNCERTAIN
    uncertain_mean_logit_max: float = 1.0   # Most queries UNCERTAIN

    # logit_margin for uncertain: lower = more UNCERTAIN
    uncertain_logit_margin_min: float = 0.0   # Everything with any uncertainty
    uncertain_logit_margin_max: float = 0.5   # Only very uncertain queries


@dataclass(frozen=True, slots=True)
class CalibrationConfig:
    """Configuration for threshold calibration.

    Attributes:
        optimization_target: What to optimize for.
        quality_weight: Weight for quality in balanced optimization (0-1).
        cost_weight: Weight for cost in balanced optimization (0-1).
        latency_weight: Weight for latency in balanced optimization (0-1).
        cost_model: Cost model for inference tiers.
        latency_model: Latency model for inference tiers.
        threshold_bounds: Bounds for threshold search.
        n_search_iterations: Number of search iterations.
        n_cv_folds: Number of cross-validation folds (0 = no CV).
        random_seed: Random seed for reproducibility.
        min_tier_fraction: Minimum fraction of queries per tier.
        quality_floor: Minimum acceptable quality (accuracy).
        cost_ceiling: Maximum acceptable cost per query.
        latency_ceiling_ms: Maximum acceptable latency per query.
    """

    optimization_target: OptimizationTarget = OptimizationTarget.BALANCED
    quality_weight: float = 0.5
    cost_weight: float = 0.3
    latency_weight: float = 0.2
    cost_model: CostModel = field(default_factory=CostModel)
    latency_model: LatencyModel = field(default_factory=LatencyModel)
    threshold_bounds: ThresholdBounds = field(default_factory=ThresholdBounds)
    n_search_iterations: int = 100
    n_cv_folds: int = 5
    random_seed: int = 42
    min_tier_fraction: float = 0.05
    quality_floor: float = 0.80
    cost_ceiling: float | None = None
    latency_ceiling_ms: float | None = None

    def __post_init__(self) -> None:
        """Validate configuration."""
        total_weight = self.quality_weight + self.cost_weight + self.latency_weight
        if abs(total_weight - 1.0) > 1e-6:
            raise ValueError(
                f"Weights must sum to 1.0, got {total_weight:.4f} "
                f"(quality={self.quality_weight}, cost={self.cost_weight}, "
                f"latency={self.latency_weight})"
            )

        if self.min_tier_fraction < 0 or self.min_tier_fraction > 0.33:
            raise ValueError(
                f"min_tier_fraction must be in [0, 0.33], got {self.min_tier_fraction}"
            )


# Preset configurations for common use cases
QUALITY_FOCUSED_CONFIG = CalibrationConfig(
    optimization_target=OptimizationTarget.QUALITY,
    quality_weight=1.0,
    cost_weight=0.0,
    latency_weight=0.0,
)

COST_FOCUSED_CONFIG = CalibrationConfig(
    optimization_target=OptimizationTarget.COST,
    quality_weight=0.0,
    cost_weight=1.0,
    latency_weight=0.0,
    quality_floor=0.85,  # Maintain minimum quality
)

LATENCY_FOCUSED_CONFIG = CalibrationConfig(
    optimization_target=OptimizationTarget.LATENCY,
    quality_weight=0.0,
    cost_weight=0.0,
    latency_weight=1.0,
    quality_floor=0.85,
)

BALANCED_CONFIG = CalibrationConfig(
    optimization_target=OptimizationTarget.BALANCED,
    quality_weight=0.5,
    cost_weight=0.3,
    latency_weight=0.2,
)

PARETO_CONFIG = CalibrationConfig(
    optimization_target=OptimizationTarget.PARETO,
    quality_weight=0.33,
    cost_weight=0.34,
    latency_weight=0.33,
    n_search_iterations=200,  # More iterations for Pareto frontier
)
