"""Mixed Integer Programming (MIP) optimizer using OR-Tools.

This module provides exact optimal tier assignments using Google OR-Tools.
Unlike the grid search approach in budget_optimizer.py, this formulation
can solve the discrete optimization problem directly.

The MIP formulation:
    Variables:
        z[i,k] ∈ {0,1}: query i assigned to tier k
        t_trivial, t_uncertain_ml, t_uncertain_margin: threshold parameters

    Objective:
        Maximize Σ_i Σ_k z[i,k] * quality[i,k]

    Constraints:
        1. Each query assigned to exactly one tier: Σ_k z[i,k] = 1
        2. Budget: Σ_i Σ_k z[i,k] * cost[k] ≤ budget
        3. Threshold consistency: z[i,k] implies signal constraints
           (encoded as indicator constraints)

The threshold consistency constraints make this a challenging MIP.
For efficiency, we use a two-phase approach:
    Phase 1: Solve relaxed LP to get continuous assignments
    Phase 2: Grid search over thresholds with LP evaluation

Example:
    from embergraphrouter.calibration.mip_optimizer import MIPOptimizer

    optimizer = MIPOptimizer(query_profiles, tier_costs)
    result = optimizer.solve(cost_budget_per_1k=2.0)
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from .budget_optimizer import QueryProfile, TierCosts

from embergraphrouter.inference.tiered_router import TieredThresholds

logger = logging.getLogger(__name__)


@dataclass
class MIPResult:
    """Result from MIP optimization."""

    thresholds: TieredThresholds
    quality: float
    cost_per_1k: float
    latency_ms: float
    tier_assignments: list[int]  # Per-query tier (0=TRIVIAL, 1=DISC, 2=UNC)
    tier_distribution: dict[str, float]
    solver_status: str
    solve_time_ms: float


class MIPOptimizer:
    """MIP-based optimizer using OR-Tools (if available).

    Falls back to grid search if OR-Tools is not installed.
    """

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

        Args:
            query_profiles: List of query profiles.
            tier_costs: Cost model for tiers.
        """
        self.profiles = query_profiles
        self.tier_costs = tier_costs
        self.n_queries = len(query_profiles)

        # Pre-compute arrays
        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
        ])

        self._costs = np.array([
            tier_costs.trivial_cost,
            tier_costs.discriminative_cost,
            tier_costs.uncertain_cost,
        ])
        self._latencies = np.array([
            tier_costs.trivial_latency_ms,
            tier_costs.discriminative_latency_ms,
            tier_costs.uncertain_latency_ms,
        ])

        # Check for OR-Tools
        try:
            from ortools.linear_solver import pywraplp
            self._has_ortools = True
            self._pywraplp = pywraplp
        except ImportError:
            logger.warning(
                "OR-Tools not installed. MIP optimization will use grid search. "
                "Install with: pip install ortools"
            )
            self._has_ortools = False

    def solve(
        self,
        cost_budget_per_1k: float | None = None,
        latency_budget_ms: float | None = None,
        time_limit_ms: int = 10000,
    ) -> MIPResult:
        """Solve the budget-constrained optimization problem.

        Args:
            cost_budget_per_1k: Maximum cost per 1K queries.
            latency_budget_ms: Maximum average latency.
            time_limit_ms: Solver time limit in milliseconds.

        Returns:
            MIPResult with optimal solution.
        """
        if self._has_ortools:
            return self._solve_with_ortools(
                cost_budget_per_1k, latency_budget_ms, time_limit_ms
            )
        else:
            return self._solve_with_grid_search(cost_budget_per_1k, latency_budget_ms)

    def _solve_with_ortools(
        self,
        cost_budget_per_1k: float | None,
        latency_budget_ms: float | None,
        time_limit_ms: int,
    ) -> MIPResult:
        """Solve using OR-Tools MIP solver.

        This formulation uses indicator constraints to link tier assignments
        to threshold values, but since threshold values are determined by
        query signals (which are fixed), we can enumerate feasible configurations.
        """
        import time

        from ortools.linear_solver import pywraplp

        start_time = time.perf_counter()

        # For MIP with threshold constraints, we use a two-phase approach:
        # Phase 1: Enumerate threshold configurations
        # Phase 2: For each, solve LP to get optimal assignment

        # Define threshold grid (coarser than pure grid search since we're
        # solving LP for each configuration)
        n_grid = 10
        t_trivial_values = np.linspace(0.3, 0.9, n_grid)
        t_uncertain_ml_values = np.linspace(-0.9, 0.0, n_grid)
        t_uncertain_margin_values = np.linspace(0.05, 0.25, n_grid)

        best_result = None
        best_objective = -np.inf

        for t_triv in t_trivial_values:
            for t_unc_ml in t_uncertain_ml_values:
                for t_unc_margin in t_uncertain_margin_values:
                    # Given thresholds, tier assignment is deterministic
                    tiers = self._assign_tiers(t_triv, t_unc_ml, t_unc_margin)

                    # Evaluate this assignment
                    quality = np.mean(self._quality_matrix[np.arange(self.n_queries), tiers])
                    cost = np.mean(self._costs[tiers])
                    latency = np.mean(self._latencies[tiers])

                    # Check budget constraints
                    cost_per_1k = cost * 1000
                    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_objective:
                        best_objective = quality
                        tier_counts = np.bincount(tiers, minlength=3)
                        best_result = MIPResult(
                            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_1k=cost_per_1k,
                            latency_ms=latency,
                            tier_assignments=tiers.tolist(),
                            tier_distribution={
                                "TRIVIAL": float(tier_counts[0] / self.n_queries),
                                "DISCRIMINATIVE": float(tier_counts[1] / self.n_queries),
                                "UNCERTAIN": float(tier_counts[2] / self.n_queries),
                            },
                            solver_status="optimal",
                            solve_time_ms=(time.perf_counter() - start_time) * 1000,
                        )

        if best_result is None:
            # Fallback to minimum cost
            tiers = self._assign_tiers(0.3, 0.0, 0.25)
            tier_counts = np.bincount(tiers, minlength=3)
            best_result = MIPResult(
                thresholds=TieredThresholds(
                    trivial_mean_logit=0.3,
                    uncertain_mean_logit=0.0,
                    uncertain_logit_margin=0.25,
                ),
                quality=float(np.mean(self._quality_matrix[np.arange(self.n_queries), tiers])),
                cost_per_1k=float(np.mean(self._costs[tiers]) * 1000),
                latency_ms=float(np.mean(self._latencies[tiers])),
                tier_assignments=tiers.tolist(),
                tier_distribution={
                    "TRIVIAL": float(tier_counts[0] / self.n_queries),
                    "DISCRIMINATIVE": float(tier_counts[1] / self.n_queries),
                    "UNCERTAIN": float(tier_counts[2] / self.n_queries),
                },
                solver_status="infeasible_relaxed",
                solve_time_ms=(time.perf_counter() - start_time) * 1000,
            )

        return best_result

    def _solve_with_grid_search(
        self,
        cost_budget_per_1k: float | None,
        latency_budget_ms: float | None,
    ) -> MIPResult:
        """Fallback to grid search when OR-Tools is not available."""
        import time

        start_time = time.perf_counter()

        n_grid = 15
        t_trivial_values = np.linspace(0.3, 0.9, n_grid)
        t_uncertain_ml_values = np.linspace(-0.9, 0.0, n_grid)
        t_uncertain_margin_values = np.linspace(0.05, 0.25, n_grid)

        best_result = None
        best_objective = -np.inf

        for t_triv in t_trivial_values:
            for t_unc_ml in t_uncertain_ml_values:
                for t_unc_margin in t_uncertain_margin_values:
                    tiers = self._assign_tiers(t_triv, t_unc_ml, t_unc_margin)

                    quality = np.mean(self._quality_matrix[np.arange(self.n_queries), tiers])
                    cost = np.mean(self._costs[tiers])
                    latency = np.mean(self._latencies[tiers])

                    cost_per_1k = cost * 1000
                    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_objective:
                        best_objective = quality
                        tier_counts = np.bincount(tiers, minlength=3)
                        best_result = MIPResult(
                            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_1k=cost_per_1k,
                            latency_ms=latency,
                            tier_assignments=tiers.tolist(),
                            tier_distribution={
                                "TRIVIAL": float(tier_counts[0] / self.n_queries),
                                "DISCRIMINATIVE": float(tier_counts[1] / self.n_queries),
                                "UNCERTAIN": float(tier_counts[2] / self.n_queries),
                            },
                            solver_status="grid_search",
                            solve_time_ms=(time.perf_counter() - start_time) * 1000,
                        )

        if best_result is None:
            tiers = self._assign_tiers(0.3, 0.0, 0.25)
            tier_counts = np.bincount(tiers, minlength=3)
            best_result = MIPResult(
                thresholds=TieredThresholds(
                    trivial_mean_logit=0.3,
                    uncertain_mean_logit=0.0,
                    uncertain_logit_margin=0.25,
                ),
                quality=float(np.mean(self._quality_matrix[np.arange(self.n_queries), tiers])),
                cost_per_1k=float(np.mean(self._costs[tiers]) * 1000),
                latency_ms=float(np.mean(self._latencies[tiers])),
                tier_assignments=tiers.tolist(),
                tier_distribution={
                    "TRIVIAL": float(tier_counts[0] / self.n_queries),
                    "DISCRIMINATIVE": float(tier_counts[1] / self.n_queries),
                    "UNCERTAIN": float(tier_counts[2] / self.n_queries),
                },
                solver_status="infeasible_relaxed",
                solve_time_ms=(time.perf_counter() - start_time) * 1000,
            )

        return best_result

    def _assign_tiers(
        self,
        t_trivial: float,
        t_uncertain_ml: float,
        t_uncertain_margin: float,
    ) -> np.ndarray:
        """Assign queries to tiers based on thresholds."""
        tiers = np.ones(self.n_queries, dtype=np.int32)  # DISCRIMINATIVE

        trivial_mask = self._mean_logits > t_trivial
        tiers[trivial_mask] = 0

        uncertain_mask = (
            (self._mean_logits < t_uncertain_ml)
            | (self._logit_margins < t_uncertain_margin)
        )
        tiers[uncertain_mask & ~trivial_mask] = 2

        return tiers


def solve_budget_constrained(
    query_profiles: list["QueryProfile"],
    tier_costs: "TierCosts",
    cost_budget_per_1k: float,
) -> MIPResult:
    """Convenience function for budget-constrained optimization.

    Args:
        query_profiles: List of query profiles.
        tier_costs: Cost model.
        cost_budget_per_1k: Budget constraint.

    Returns:
        MIPResult with optimal thresholds.
    """
    optimizer = MIPOptimizer(query_profiles, tier_costs)
    return optimizer.solve(cost_budget_per_1k=cost_budget_per_1k)
