"""Budget-aware threshold optimization using MIP/LP.

This module provides mathematically optimal threshold calibration subject to
budget constraints on cost and/or latency. Unlike the Bayesian search approach
which explores empirically, this formulation finds provably optimal thresholds.

The key insight is that tier assignment is a step function of thresholds:
    TRIVIAL if mean_logit > t_trivial
    UNCERTAIN if mean_logit < t_uncertain_ml OR logit_margin < t_uncertain_margin
    DISCRIMINATIVE otherwise

Given fixed thresholds, the problem becomes linear in assignments, making it
amenable to MIP (Mixed Integer Programming) or LP relaxation.

Example:
    from embergraphrouter.calibration.budget_optimizer import BudgetOptimizer

    optimizer = BudgetOptimizer.from_calibration_data(dataset)

    # Find optimal thresholds for $2/1K query budget
    result = optimizer.optimize_for_budget(cost_budget_per_1k=2.0)
    print(f"Optimal thresholds: {result.thresholds}")
    print(f"Expected quality: {result.quality:.1%}")

    # Generate Pareto frontier
    frontier = optimizer.compute_pareto_frontier(n_points=20)
    frontier.save("pareto_frontier.json")
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from .dataset import OfflineDataset

from embergraphrouter.inference.tiered_router import InferenceTier, TieredThresholds

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class TierCosts:
    """Cost and latency model for each tier.

    Attributes:
        trivial_cost: Cost per query for TRIVIAL tier (USD).
        discriminative_cost: Cost per query for DISCRIMINATIVE tier (USD).
        uncertain_cost: Cost per query for UNCERTAIN tier (USD).
        trivial_latency_ms: Latency for TRIVIAL tier (ms).
        discriminative_latency_ms: Latency for DISCRIMINATIVE tier (ms).
        uncertain_latency_ms: Latency for UNCERTAIN tier (ms).
    """

    trivial_cost: float = 0.0001  # $0.10 per 1K
    discriminative_cost: float = 0.001  # $1.00 per 1K
    uncertain_cost: float = 0.01  # $10.00 per 1K
    trivial_latency_ms: float = 100.0
    discriminative_latency_ms: float = 200.0
    uncertain_latency_ms: float = 1500.0


@dataclass
class QueryProfile:
    """Profile of a query for budget optimization.

    Contains all information needed to optimize tier assignment.

    Attributes:
        query_id: Unique identifier for the query.
        mean_logit: GNN mean logit signal.
        logit_margin: GNN logit margin signal.
        quality_trivial: Expected quality if routed to TRIVIAL.
        quality_discriminative: Expected quality if routed to DISCRIMINATIVE.
        quality_uncertain: Expected quality if routed to UNCERTAIN.
    """

    query_id: str
    mean_logit: float
    logit_margin: float
    quality_trivial: float
    quality_discriminative: float
    quality_uncertain: float


@dataclass
class OptimizationResult:
    """Result of budget-constrained optimization.

    Attributes:
        thresholds: Optimal tier thresholds.
        quality: Expected quality (fraction of oracle).
        cost_per_query: Expected cost per query (USD).
        cost_per_1k: Expected cost per 1K queries (USD).
        latency_ms: Expected latency per query (ms).
        tier_distribution: Fraction of queries per tier.
        n_queries: Number of queries in optimization.
        solver_status: Solver status string.
    """

    thresholds: TieredThresholds
    quality: float
    cost_per_query: float
    cost_per_1k: float
    latency_ms: float
    tier_distribution: dict[str, float]
    n_queries: int
    solver_status: str = "optimal"


@dataclass
class ParetoPoint:
    """A point on the Pareto frontier."""

    quality: float
    cost_per_1k: float
    latency_ms: float
    thresholds: TieredThresholds
    tier_distribution: dict[str, float]


@dataclass
class ParetoFrontier:
    """Pareto frontier of quality-cost-latency tradeoffs.

    Attributes:
        points: List of Pareto-optimal points.
        cost_range: (min, max) cost per 1K queries.
        quality_range: (min, max) quality.
    """

    points: list[ParetoPoint]
    cost_range: tuple[float, float] = (0.0, 10.0)
    quality_range: tuple[float, float] = (0.0, 1.0)

    def get_thresholds_for_budget(
        self, cost_budget_per_1k: float
    ) -> TieredThresholds | None:
        """Get optimal thresholds for a given budget.

        Args:
            cost_budget_per_1k: Maximum cost per 1K queries (USD).

        Returns:
            Optimal thresholds, or None if no feasible point exists.
        """
        feasible = [p for p in self.points if p.cost_per_1k <= cost_budget_per_1k]
        if not feasible:
            return None
        best = max(feasible, key=lambda p: p.quality)
        return best.thresholds

    def to_dict(self) -> dict:
        """Convert to dictionary for serialization."""
        return {
            "points": [
                {
                    "quality": p.quality,
                    "cost_per_1k": p.cost_per_1k,
                    "latency_ms": p.latency_ms,
                    "thresholds": {
                        "trivial_mean_logit": p.thresholds.trivial_mean_logit,
                        "uncertain_mean_logit": p.thresholds.uncertain_mean_logit,
                        "uncertain_logit_margin": p.thresholds.uncertain_logit_margin,
                    },
                    "tier_distribution": p.tier_distribution,
                }
                for p in self.points
            ],
            "cost_range": list(self.cost_range),
            "quality_range": list(self.quality_range),
        }

    def save(self, path: str | Path) -> None:
        """Save frontier to JSON file."""
        with open(path, "w") as f:
            json.dump(self.to_dict(), f, indent=2)

    @classmethod
    def load(cls, path: str | Path) -> "ParetoFrontier":
        """Load frontier from JSON file."""
        with open(path) as f:
            data = json.load(f)

        points = []
        for p in data["points"]:
            points.append(
                ParetoPoint(
                    quality=p["quality"],
                    cost_per_1k=p["cost_per_1k"],
                    latency_ms=p["latency_ms"],
                    thresholds=TieredThresholds(
                        trivial_mean_logit=p["thresholds"]["trivial_mean_logit"],
                        uncertain_mean_logit=p["thresholds"]["uncertain_mean_logit"],
                        uncertain_logit_margin=p["thresholds"]["uncertain_logit_margin"],
                    ),
                    tier_distribution=p["tier_distribution"],
                )
            )

        return cls(
            points=points,
            cost_range=tuple(data["cost_range"]),
            quality_range=tuple(data["quality_range"]),
        )


class BudgetOptimizer:
    """Budget-aware threshold optimizer using grid search with LP evaluation.

    This optimizer finds thresholds that maximize quality subject to budget
    constraints. It uses a grid search over threshold space combined with
    exact evaluation of each threshold configuration.

    For large-scale optimization, consider using OR-Tools MIP solver for
    the inner problem (see optimize_with_ortools method).
    """

    def __init__(
        self,
        query_profiles: list[QueryProfile],
        tier_costs: TierCosts | None = None,
    ):
        """Initialize optimizer.

        Args:
            query_profiles: List of query profiles with signals and quality.
            tier_costs: Cost/latency model for each tier.
        """
        self.profiles = query_profiles
        self.tier_costs = tier_costs or TierCosts()
        self.n_queries = len(query_profiles)

        # Pre-compute arrays for vectorized evaluation
        self._mean_logits = np.array([p.mean_logit for p in query_profiles])
        self._logit_margins = np.array([p.logit_margin for p in query_profiles])
        self._quality_matrix = np.array([
            [p.quality_trivial, p.quality_discriminative, p.quality_uncertain]
            for p in query_profiles
        ])  # (N, 3)

        logger.info(
            "BudgetOptimizer initialized with %d queries, cost model: %s",
            self.n_queries,
            self.tier_costs,
        )

    @classmethod
    def from_calibration_data(
        cls,
        dataset: "OfflineDataset",
        tier_costs: TierCosts | None = None,
    ) -> "BudgetOptimizer":
        """Create optimizer from calibration dataset.

        Args:
            dataset: Offline calibration dataset with samples.
            tier_costs: Cost/latency model for each tier.

        Returns:
            Initialized BudgetOptimizer.
        """
        profiles = []
        for sample in dataset.samples:
            # Compute per-tier quality from ground truth
            effects = list(sample.ground_truth_effects.values())
            oracle_effect = sample.oracle_effect

            # TRIVIAL quality: use cheapest model's effect (assume index 0)
            # In practice, this should be calibrated to your cheapest model
            quality_trivial = effects[0] if effects else 0.5

            # DISCRIMINATIVE quality: use GNN's top pick
            top1_model = sample.top1_model
            quality_discriminative = sample.ground_truth_effects.get(
                top1_model, oracle_effect
            )

            # UNCERTAIN quality: use oracle (ensemble should match oracle)
            quality_uncertain = oracle_effect

            profiles.append(
                QueryProfile(
                    query_id=sample.query_id,
                    mean_logit=sample.mean_logit,
                    logit_margin=sample.logit_margin,
                    quality_trivial=quality_trivial,
                    quality_discriminative=quality_discriminative,
                    quality_uncertain=quality_uncertain,
                )
            )

        return cls(profiles, tier_costs)

    def _assign_tiers(
        self,
        t_trivial: float,
        t_uncertain_ml: float,
        t_uncertain_margin: float,
    ) -> np.ndarray:
        """Assign queries to tiers based on thresholds.

        Args:
            t_trivial: TRIVIAL tier threshold (mean_logit).
            t_uncertain_ml: UNCERTAIN tier threshold (mean_logit).
            t_uncertain_margin: UNCERTAIN tier threshold (logit_margin).

        Returns:
            Array of tier indices (0=TRIVIAL, 1=DISCRIMINATIVE, 2=UNCERTAIN).
        """
        n = self.n_queries
        tiers = np.ones(n, dtype=np.int32)  # Default to DISCRIMINATIVE (1)

        # TRIVIAL: mean_logit > threshold
        trivial_mask = self._mean_logits > t_trivial
        tiers[trivial_mask] = 0

        # UNCERTAIN: mean_logit < threshold OR logit_margin < threshold
        uncertain_mask = (
            (self._mean_logits < t_uncertain_ml)
            | (self._logit_margins < t_uncertain_margin)
        )
        # UNCERTAIN takes precedence for non-trivial queries
        tiers[uncertain_mask & ~trivial_mask] = 2

        return tiers

    def _evaluate_thresholds(
        self,
        t_trivial: float,
        t_uncertain_ml: float,
        t_uncertain_margin: float,
    ) -> tuple[float, float, float, dict[str, float]]:
        """Evaluate a threshold configuration.

        Args:
            t_trivial: TRIVIAL tier threshold.
            t_uncertain_ml: UNCERTAIN tier mean_logit threshold.
            t_uncertain_margin: UNCERTAIN tier margin threshold.

        Returns:
            Tuple of (quality, cost_per_query, latency_ms, tier_distribution).
        """
        tiers = self._assign_tiers(t_trivial, t_uncertain_ml, t_uncertain_margin)

        # Compute quality as weighted average
        quality = np.mean([
            self._quality_matrix[i, t] for i, t in enumerate(tiers)
        ])

        # Compute tier distribution
        tier_counts = np.bincount(tiers, minlength=3)
        tier_dist = {
            "TRIVIAL": float(tier_counts[0] / self.n_queries),
            "DISCRIMINATIVE": float(tier_counts[1] / self.n_queries),
            "UNCERTAIN": float(tier_counts[2] / self.n_queries),
        }

        # Compute expected cost and latency
        costs = np.array([
            self.tier_costs.trivial_cost,
            self.tier_costs.discriminative_cost,
            self.tier_costs.uncertain_cost,
        ])
        latencies = np.array([
            self.tier_costs.trivial_latency_ms,
            self.tier_costs.discriminative_latency_ms,
            self.tier_costs.uncertain_latency_ms,
        ])

        cost_per_query = float(np.mean(costs[tiers]))
        latency_ms = float(np.mean(latencies[tiers]))

        return quality, cost_per_query, latency_ms, tier_dist

    def optimize_for_budget(
        self,
        cost_budget_per_1k: float | None = None,
        latency_budget_ms: float | None = None,
        grid_resolution: int = 20,
    ) -> OptimizationResult:
        """Find optimal thresholds for given budget constraints.

        Uses grid search over threshold space to find the configuration
        that maximizes quality subject to cost and/or latency constraints.

        Args:
            cost_budget_per_1k: Maximum cost per 1K queries (USD).
            latency_budget_ms: Maximum average latency (ms).
            grid_resolution: Number of points per dimension in grid search.

        Returns:
            OptimizationResult with optimal thresholds and metrics.
        """
        # Define search grid - WIDE bounds to explore full cost-quality tradeoff
        # trivial_mean_logit: -1.5 (97% TRIVIAL) to +1.5 (6% TRIVIAL)
        # uncertain_mean_logit: -2.0 (no UNCERTAIN) to +1.0 (most UNCERTAIN)
        t_trivial_range = np.linspace(-1.5, 1.5, grid_resolution)
        t_uncertain_ml_range = np.linspace(-2.0, 1.0, grid_resolution)
        t_uncertain_margin_range = np.linspace(0.0, 0.5, grid_resolution)

        best_result = None
        best_quality = -1.0

        for t_triv in t_trivial_range:
            for t_unc_ml in t_uncertain_ml_range:
                for t_unc_margin in t_uncertain_margin_range:
                    quality, cost, latency, tier_dist = self._evaluate_thresholds(
                        t_triv, t_unc_ml, t_unc_margin
                    )
                    cost_per_1k = cost * 1000

                    # Check constraints
                    if cost_budget_per_1k is not None and cost_per_1k > cost_budget_per_1k:
                        continue
                    if latency_budget_ms is not None and latency > latency_budget_ms:
                        continue

                    if quality > best_quality:
                        best_quality = quality
                        best_result = OptimizationResult(
                            thresholds=TieredThresholds(
                                trivial_mean_logit=float(t_triv),
                                uncertain_mean_logit=float(t_unc_ml),
                                uncertain_logit_margin=float(t_unc_margin),
                            ),
                            quality=quality,
                            cost_per_query=cost,
                            cost_per_1k=cost_per_1k,
                            latency_ms=latency,
                            tier_distribution=tier_dist,
                            n_queries=self.n_queries,
                            solver_status="optimal",
                        )

        if best_result is None:
            # No feasible solution - return least-cost configuration
            logger.warning(
                "No feasible solution for budget constraints. "
                "Returning minimum cost configuration."
            )
            # Use maximum trivial threshold to minimize cost
            quality, cost, latency, tier_dist = self._evaluate_thresholds(
                0.3, 0.0, 0.25  # Most aggressive towards TRIVIAL
            )
            best_result = OptimizationResult(
                thresholds=TieredThresholds(
                    trivial_mean_logit=0.3,
                    uncertain_mean_logit=0.0,
                    uncertain_logit_margin=0.25,
                ),
                quality=quality,
                cost_per_query=cost,
                cost_per_1k=cost * 1000,
                latency_ms=latency,
                tier_distribution=tier_dist,
                n_queries=self.n_queries,
                solver_status="infeasible_relaxed",
            )

        return best_result

    def compute_pareto_frontier(
        self,
        n_points: int = 50,
        grid_resolution: int = 15,
    ) -> ParetoFrontier:
        """Compute Pareto frontier of quality-cost tradeoffs.

        Sweeps through cost budgets and finds optimal thresholds for each,
        then filters to non-dominated points.

        Args:
            n_points: Number of budget points to evaluate.
            grid_resolution: Resolution for threshold grid search.

        Returns:
            ParetoFrontier with optimal tradeoff points.
        """
        # First, find the cost range by evaluating extremes
        # Max TRIVIAL: low threshold → most queries routed cheap
        _, min_cost, _, _ = self._evaluate_thresholds(-1.5, -2.0, 0.5)
        # Max UNCERTAIN: high threshold → most queries routed to ensemble
        _, max_cost, _, _ = self._evaluate_thresholds(1.5, 1.0, 0.0)

        cost_budgets = np.linspace(min_cost * 1000, max_cost * 1000, n_points)

        all_points = []
        for budget in cost_budgets:
            result = self.optimize_for_budget(
                cost_budget_per_1k=budget,
                grid_resolution=grid_resolution,
            )
            all_points.append(
                ParetoPoint(
                    quality=result.quality,
                    cost_per_1k=result.cost_per_1k,
                    latency_ms=result.latency_ms,
                    thresholds=result.thresholds,
                    tier_distribution=result.tier_distribution,
                )
            )

        # Filter to Pareto-optimal points
        pareto_points = self._filter_pareto(all_points)

        # Compute ranges
        qualities = [p.quality for p in pareto_points]
        costs = [p.cost_per_1k for p in pareto_points]

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

    def _filter_pareto(self, points: list[ParetoPoint]) -> list[ParetoPoint]:
        """Filter to Pareto-optimal points (maximize quality, minimize cost)."""
        pareto = []
        for p1 in points:
            is_dominated = False
            for p2 in points:
                if p1 is p2:
                    continue
                # p2 dominates p1 if p2 is better on both objectives
                if p2.quality >= p1.quality and p2.cost_per_1k <= p1.cost_per_1k:
                    if p2.quality > p1.quality or p2.cost_per_1k < p1.cost_per_1k:
                        is_dominated = True
                        break
            if not is_dominated:
                pareto.append(p1)

        # Sort by cost for nice output
        pareto.sort(key=lambda p: p.cost_per_1k)
        return pareto

    def save_optimization_cache(self, path: str | Path) -> None:
        """Save query profiles and Pareto frontier for runtime use.

        This saves all data needed for budget-constrained inference:
        1. Query signal distributions (for validation)
        2. Pareto frontier (for threshold lookup)
        3. Tier cost model

        Args:
            path: Output path (JSON format).
        """
        frontier = self.compute_pareto_frontier()

        cache = {
            "n_queries": self.n_queries,
            "signal_stats": {
                "mean_logit": {
                    "mean": float(np.mean(self._mean_logits)),
                    "std": float(np.std(self._mean_logits)),
                    "min": float(np.min(self._mean_logits)),
                    "max": float(np.max(self._mean_logits)),
                },
                "logit_margin": {
                    "mean": float(np.mean(self._logit_margins)),
                    "std": float(np.std(self._logit_margins)),
                    "min": float(np.min(self._logit_margins)),
                    "max": float(np.max(self._logit_margins)),
                },
            },
            "tier_costs": {
                "trivial_cost": self.tier_costs.trivial_cost,
                "discriminative_cost": self.tier_costs.discriminative_cost,
                "uncertain_cost": self.tier_costs.uncertain_cost,
                "trivial_latency_ms": self.tier_costs.trivial_latency_ms,
                "discriminative_latency_ms": self.tier_costs.discriminative_latency_ms,
                "uncertain_latency_ms": self.tier_costs.uncertain_latency_ms,
            },
            "pareto_frontier": frontier.to_dict(),
        }

        with open(path, "w") as f:
            json.dump(cache, f, indent=2)

        logger.info("Saved optimization cache to %s", path)


def load_budget_thresholds(
    cache_path: str | Path,
    cost_budget_per_1k: float,
) -> TieredThresholds:
    """Load optimal thresholds for a given budget from cached Pareto frontier.

    This is the main runtime API for budget-constrained inference.

    Args:
        cache_path: Path to saved optimization cache.
        cost_budget_per_1k: Maximum cost per 1K queries (USD).

    Returns:
        Optimal TieredThresholds for the given budget.

    Raises:
        ValueError: If no feasible configuration exists for the budget.
    """
    with open(cache_path) as f:
        cache = json.load(f)

    frontier = ParetoFrontier.load(cache_path)
    thresholds = frontier.get_thresholds_for_budget(cost_budget_per_1k)

    if thresholds is None:
        min_cost = cache["pareto_frontier"]["cost_range"][0]
        raise ValueError(
            f"No feasible configuration for budget ${cost_budget_per_1k}/1K. "
            f"Minimum achievable cost is ${min_cost:.2f}/1K."
        )

    return thresholds
