"""Threshold optimizer with multi-objective support.

This module implements the main calibration workflow including
Pareto frontier analysis, cross-validation, and result aggregation.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from .config import CalibrationConfig
    from .dataset import DatasetSplit, OfflineDataset

from embergraphrouter.inference.tiered_router import TieredThresholds

from .config import OptimizationTarget
from .metrics import CalibrationMetrics, RoutingSimulator
from .search import BayesianSearch, GridSearch, RandomSearch, SearchStrategy

logger = logging.getLogger(__name__)


@dataclass
class EvaluationResult:
    """Result of evaluating a single threshold configuration.

    Attributes:
        thresholds: Threshold configuration.
        metrics: Full metrics on evaluation set.
        cv_metrics: Cross-validation metrics (if CV enabled).
        combined_score: Weighted combined score.
        quality: Quality component (frac_of_oracle).
        cost: Cost per query.
        latency: Latency per query.
        is_pareto_optimal: Whether this is on Pareto frontier.
    """

    thresholds: TieredThresholds
    metrics: CalibrationMetrics
    cv_metrics: list[CalibrationMetrics] | None = None
    combined_score: float = 0.0
    quality: float = 0.0
    cost: float = 0.0
    latency: float = 0.0
    is_pareto_optimal: bool = False


@dataclass
class ParetoFrontier:
    """Pareto frontier of threshold configurations.

    A configuration is Pareto-optimal if no other configuration
    dominates it (i.e., is better on ALL objectives).

    Attributes:
        points: List of Pareto-optimal evaluation results.
        quality_range: (min, max) quality on frontier.
        cost_range: (min, max) cost on frontier.
        latency_range: (min, max) latency on frontier.
    """

    points: list[EvaluationResult]
    quality_range: tuple[float, float] = (0.0, 1.0)
    cost_range: tuple[float, float] = (0.0, 1.0)
    latency_range: tuple[float, float] = (0.0, 1.0)

    def __len__(self) -> int:
        return len(self.points)

    def __iter__(self):
        return iter(self.points)

    def best_for_quality(self) -> EvaluationResult | None:
        """Return point with highest quality."""
        if not self.points:
            return None
        return max(self.points, key=lambda p: p.quality)

    def best_for_cost(self) -> EvaluationResult | None:
        """Return point with lowest cost."""
        if not self.points:
            return None
        return min(self.points, key=lambda p: p.cost)

    def best_for_latency(self) -> EvaluationResult | None:
        """Return point with lowest latency."""
        if not self.points:
            return None
        return min(self.points, key=lambda p: p.latency)

    def best_balanced(
        self,
        quality_weight: float = 0.5,
        cost_weight: float = 0.3,
        latency_weight: float = 0.2,
    ) -> EvaluationResult | None:
        """Return best balanced point on frontier."""
        if not self.points:
            return None

        def score(p: EvaluationResult) -> float:
            # Normalize to [0, 1]
            q = (p.quality - self.quality_range[0]) / max(
                self.quality_range[1] - self.quality_range[0], 1e-8
            )
            c = 1.0 - (p.cost - self.cost_range[0]) / max(
                self.cost_range[1] - self.cost_range[0], 1e-8
            )
            lat = 1.0 - (p.latency - self.latency_range[0]) / max(
                self.latency_range[1] - self.latency_range[0], 1e-8
            )
            return quality_weight * q + cost_weight * c + latency_weight * lat

        return max(self.points, key=score)


@dataclass
class CalibrationResult:
    """Complete result of threshold calibration.

    Attributes:
        optimal_thresholds: Best thresholds according to config.
        optimal_metrics: Metrics for optimal thresholds.
        pareto_frontier: Pareto frontier of all configurations.
        all_results: All evaluated configurations.
        config: Calibration configuration used.
        n_iterations: Total iterations performed.
        cv_summary: Cross-validation summary if CV was used.
    """

    optimal_thresholds: TieredThresholds
    optimal_metrics: CalibrationMetrics
    pareto_frontier: ParetoFrontier
    all_results: list[EvaluationResult]
    config: "CalibrationConfig"
    n_iterations: int
    cv_summary: dict | None = None

    def to_dict(self) -> dict:
        """Convert to dictionary for serialization."""
        return {
            "optimal_thresholds": {
                "trivial_mean_logit": self.optimal_thresholds.trivial_mean_logit,
                "trivial_min_prob": self.optimal_thresholds.trivial_min_prob,
                "uncertain_mean_logit": self.optimal_thresholds.uncertain_mean_logit,
                "uncertain_max_prob": self.optimal_thresholds.uncertain_max_prob,
                "uncertain_logit_margin": self.optimal_thresholds.uncertain_logit_margin,
            },
            "optimal_metrics": self.optimal_metrics.to_dict(),
            "pareto_frontier_size": len(self.pareto_frontier),
            "n_iterations": self.n_iterations,
            "cv_summary": self.cv_summary,
        }


class ThresholdCalibrator:
    """Main threshold calibration orchestrator.

    This class coordinates the calibration workflow:
    1. Initialize search strategy
    2. Evaluate threshold configurations
    3. Apply cross-validation (optional)
    4. Compute Pareto frontier
    5. Select optimal thresholds

    Example:
        config = CalibrationConfig(
            optimization_target=OptimizationTarget.BALANCED,
            quality_weight=0.5,
            cost_weight=0.3,
            latency_weight=0.2,
        )
        calibrator = ThresholdCalibrator(config)
        result = calibrator.calibrate(dataset)
        print(result.optimal_thresholds)
    """

    def __init__(
        self,
        config: "CalibrationConfig",
        search_strategy: str = "bayesian",
    ):
        """Initialize calibrator.

        Args:
            config: Calibration configuration.
            search_strategy: Search strategy ("grid", "random", "bayesian").
        """
        self.config = config
        self.search_strategy_name = search_strategy

        self._simulator = RoutingSimulator(
            cost_model=config.cost_model,
            latency_model=config.latency_model,
            min_tier_fraction=config.min_tier_fraction,
            quality_floor=config.quality_floor,
        )

    def calibrate(
        self,
        dataset: "OfflineDataset",
        search_strategy: SearchStrategy | None = None,
    ) -> CalibrationResult:
        """Run threshold calibration.

        Args:
            dataset: Offline dataset for calibration.
            search_strategy: Optional pre-configured search strategy.

        Returns:
            CalibrationResult with optimal thresholds.
        """
        logger.info(
            "Starting calibration: %d samples, %d iterations, %s strategy",
            len(dataset),
            self.config.n_search_iterations,
            self.search_strategy_name,
        )

        # Create search strategy
        if search_strategy is None:
            search_strategy = self._create_search_strategy()

        # Split data for evaluation
        if self.config.n_cv_folds > 1:
            cv_splits = dataset.cross_validation_splits(
                n_folds=self.config.n_cv_folds,
                random_state=self.config.random_seed,
            )
            # Use first split's training data for main evaluation
            train_split, _ = cv_splits[0]
        else:
            train_split, _ = dataset.train_test_split(
                test_fraction=0.2,
                random_state=self.config.random_seed,
            )
            cv_splits = None

        # Run search
        all_results: list[EvaluationResult] = []
        for iteration in range(self.config.n_search_iterations):
            candidates = search_strategy.suggest(n=1)
            if not candidates:
                logger.info("Search exhausted at iteration %d", iteration)
                break

            for candidate in candidates:
                result = self._evaluate_candidate(
                    candidate.thresholds,
                    train_split,
                    cv_splits,
                )
                all_results.append(result)

                # Update search with feedback
                search_strategy.update(candidate, result.combined_score)

            if (iteration + 1) % 10 == 0:
                best_so_far = max(all_results, key=lambda r: r.combined_score)
                logger.info(
                    "Iteration %d/%d: best score=%.4f, quality=%.3f, cost=%.4f",
                    iteration + 1,
                    self.config.n_search_iterations,
                    best_so_far.combined_score,
                    best_so_far.quality,
                    best_so_far.cost,
                )

        # Compute Pareto frontier
        pareto_frontier = self._compute_pareto_frontier(all_results)

        # Select optimal based on configuration
        optimal = self._select_optimal(all_results, pareto_frontier)

        # Compute CV summary if applicable
        cv_summary = None
        if cv_splits is not None and optimal.cv_metrics:
            cv_summary = self._compute_cv_summary(optimal.cv_metrics)

        return CalibrationResult(
            optimal_thresholds=optimal.thresholds,
            optimal_metrics=optimal.metrics,
            pareto_frontier=pareto_frontier,
            all_results=all_results,
            config=self.config,
            n_iterations=len(all_results),
            cv_summary=cv_summary,
        )

    def _create_search_strategy(self) -> SearchStrategy:
        """Create search strategy from config."""
        name = self.search_strategy_name.lower()

        if name == "grid":
            return GridSearch(
                bounds=self.config.threshold_bounds,
                random_seed=self.config.random_seed,
            )
        elif name == "random":
            return RandomSearch(
                bounds=self.config.threshold_bounds,
                random_seed=self.config.random_seed,
            )
        elif name == "bayesian":
            return BayesianSearch(
                bounds=self.config.threshold_bounds,
                random_seed=self.config.random_seed,
            )
        else:
            raise ValueError(f"Unknown search strategy: {name}")

    def _evaluate_candidate(
        self,
        thresholds: TieredThresholds,
        train_split: "DatasetSplit",
        cv_splits: list[tuple["DatasetSplit", "DatasetSplit"]] | None,
    ) -> EvaluationResult:
        """Evaluate a threshold configuration."""
        # Main evaluation
        metrics = self._simulator.evaluate(train_split, thresholds)

        # Cross-validation if enabled
        cv_metrics = None
        if cv_splits is not None:
            cv_metrics = []
            for _train_fold, val_fold in cv_splits:
                fold_metrics = self._simulator.evaluate(val_fold, thresholds)
                cv_metrics.append(fold_metrics)

            # Use CV average for quality
            cv_quality = np.mean([m.frac_of_oracle for m in cv_metrics])
            cv_cost = np.mean([m.avg_cost_per_query for m in cv_metrics])
            cv_latency = np.mean([m.avg_latency_ms for m in cv_metrics])
        else:
            cv_quality = metrics.frac_of_oracle
            cv_cost = metrics.avg_cost_per_query
            cv_latency = metrics.avg_latency_ms

        # Compute combined score
        combined_score = metrics.combined_score(
            quality_weight=self.config.quality_weight,
            cost_weight=self.config.cost_weight,
            latency_weight=self.config.latency_weight,
        )

        return EvaluationResult(
            thresholds=thresholds,
            metrics=metrics,
            cv_metrics=cv_metrics,
            combined_score=combined_score,
            quality=cv_quality,
            cost=cv_cost,
            latency=cv_latency,
        )

    def _compute_pareto_frontier(
        self,
        results: list[EvaluationResult],
    ) -> ParetoFrontier:
        """Compute Pareto frontier of results.

        A point is Pareto-optimal if no other point dominates it
        (better on ALL objectives: quality, -cost, -latency).
        """
        # Filter valid results
        valid_results = [r for r in results if r.metrics.is_valid]
        if not valid_results:
            valid_results = results  # Fall back to all if none valid

        # Extract objective values (maximize quality, minimize cost/latency)
        points = []
        for r in valid_results:
            points.append({
                "result": r,
                "quality": r.quality,
                "cost": r.cost,
                "latency": r.latency,
            })

        # Find Pareto-optimal points
        pareto_points = []
        for i, p1 in enumerate(points):
            is_dominated = False
            for j, p2 in enumerate(points):
                if i == j:
                    continue
                # p2 dominates p1 if p2 is better on all objectives
                if (
                    p2["quality"] >= p1["quality"]
                    and p2["cost"] <= p1["cost"]
                    and p2["latency"] <= p1["latency"]
                    and (
                        p2["quality"] > p1["quality"]
                        or p2["cost"] < p1["cost"]
                        or p2["latency"] < p1["latency"]
                    )
                ):
                    is_dominated = True
                    break

            if not is_dominated:
                p1["result"].is_pareto_optimal = True
                pareto_points.append(p1["result"])

        # Compute ranges
        qualities = [p.quality for p in pareto_points] if pareto_points else [0.0]
        costs = [p.cost for p in pareto_points] if pareto_points else [0.0]
        latencies = [p.latency for p in pareto_points] if pareto_points else [0.0]

        return ParetoFrontier(
            points=pareto_points,
            quality_range=(min(qualities), max(qualities)),
            cost_range=(min(costs), max(costs)),
            latency_range=(min(latencies), max(latencies)),
        )

    def _select_optimal(
        self,
        results: list[EvaluationResult],
        pareto_frontier: ParetoFrontier,
    ) -> EvaluationResult:
        """Select optimal result based on configuration."""
        target = self.config.optimization_target

        if target == OptimizationTarget.QUALITY:
            return max(results, key=lambda r: r.quality)

        elif target == OptimizationTarget.COST:
            # Among valid results, find lowest cost
            valid = [r for r in results if r.quality >= self.config.quality_floor]
            if valid:
                return min(valid, key=lambda r: r.cost)
            return min(results, key=lambda r: r.cost)

        elif target == OptimizationTarget.LATENCY:
            valid = [r for r in results if r.quality >= self.config.quality_floor]
            if valid:
                return min(valid, key=lambda r: r.latency)
            return min(results, key=lambda r: r.latency)

        elif target == OptimizationTarget.PARETO:
            return pareto_frontier.best_balanced(
                quality_weight=self.config.quality_weight,
                cost_weight=self.config.cost_weight,
                latency_weight=self.config.latency_weight,
            ) or max(results, key=lambda r: r.combined_score)

        else:  # BALANCED
            return max(results, key=lambda r: r.combined_score)

    def _compute_cv_summary(
        self,
        cv_metrics: list[CalibrationMetrics],
    ) -> dict:
        """Compute cross-validation summary statistics."""
        qualities = [m.frac_of_oracle for m in cv_metrics]
        costs = [m.avg_cost_per_query for m in cv_metrics]
        latencies = [m.avg_latency_ms for m in cv_metrics]

        return {
            "quality": {
                "mean": float(np.mean(qualities)),
                "std": float(np.std(qualities)),
                "min": float(np.min(qualities)),
                "max": float(np.max(qualities)),
            },
            "cost": {
                "mean": float(np.mean(costs)),
                "std": float(np.std(costs)),
            },
            "latency_ms": {
                "mean": float(np.mean(latencies)),
                "std": float(np.std(latencies)),
            },
            "n_folds": len(cv_metrics),
        }


def quick_calibrate(
    dataset: "OfflineDataset",
    optimization_target: str = "balanced",
    n_iterations: int = 50,
) -> CalibrationResult:
    """Quick calibration with sensible defaults.

    Args:
        dataset: Offline dataset for calibration.
        optimization_target: One of "quality", "cost", "latency", "balanced".
        n_iterations: Number of search iterations.

    Returns:
        CalibrationResult.
    """
    from .config import CalibrationConfig

    target_map = {
        "quality": OptimizationTarget.QUALITY,
        "cost": OptimizationTarget.COST,
        "latency": OptimizationTarget.LATENCY,
        "balanced": OptimizationTarget.BALANCED,
        "pareto": OptimizationTarget.PARETO,
    }

    config = CalibrationConfig(
        optimization_target=target_map.get(optimization_target, OptimizationTarget.BALANCED),
        n_search_iterations=n_iterations,
        n_cv_folds=3,  # Quick CV
    )

    calibrator = ThresholdCalibrator(config, search_strategy="bayesian")
    return calibrator.calibrate(dataset)
