"""Tests for optimization module — expression parsing, LP/MIP solving, provider registry.

Tests the documented requirements:
    1. Expression parser handles all supported term forms
    2. LP product-mix problem solves to known-correct optimum
    3. Infeasible problems are detected
    4. MIP integer/binary variables work correctly
    5. Provider registry discovers optimization domain
"""

import pytest

from platoon.learning.optimization import (
    Constraint,
    OptimizationProblem,
    ProblemSense,
    Variable,
    VariableType,
    parse_constraint,
    parse_linear_expression,
)
from platoon.learning.providers import get_provider, list_providers

# Check if OR-Tools is installed
try:
    from ortools.linear_solver import pywraplp  # noqa: F401

    ortools_available = True
except ImportError:
    ortools_available = False

requires_ortools = pytest.mark.skipif(
    not ortools_available, reason="OR-Tools not installed"
)


# ===========================================================================
# Expression parser
# ===========================================================================


class TestExpressionParser:
    def test_simple_terms(self):
        """Parse '2*x + 3*y' into two terms."""
        terms = parse_linear_expression("2*x + 3*y")
        assert terms == [(2.0, "x"), (3.0, "y")]

    def test_implicit_coefficient(self):
        """Variables without a coefficient have coefficient 1."""
        terms = parse_linear_expression("x + y")
        assert terms == [(1.0, "x"), (1.0, "y")]

    def test_negative_coefficient(self):
        """Leading negative sign is captured."""
        terms = parse_linear_expression("-3*z")
        assert terms == [(-3.0, "z")]

    def test_subtraction(self):
        """Subtraction produces negative coefficient."""
        terms = parse_linear_expression("x - 2*y")
        assert terms == [(1.0, "x"), (-2.0, "y")]

    def test_underscore_variable(self):
        """Variable names with underscores are valid."""
        terms = parse_linear_expression("2*x_1 + 3*y_2")
        assert terms == [(2.0, "x_1"), (3.0, "y_2")]

    def test_decimal_coefficient(self):
        """Decimal coefficients are parsed correctly."""
        terms = parse_linear_expression("2.5*x + 0.5*y")
        assert terms == [(2.5, "x"), (0.5, "y")]

    def test_single_variable(self):
        """A single variable with no coefficient."""
        terms = parse_linear_expression("x")
        assert terms == [(1.0, "x")]

    def test_three_terms(self):
        """Multiple terms with mixed signs."""
        terms = parse_linear_expression("3*a + 2*b - c")
        assert terms == [(3.0, "a"), (2.0, "b"), (-1.0, "c")]

    def test_empty_expression_raises(self):
        """Empty string raises ValueError."""
        with pytest.raises(ValueError):
            parse_linear_expression("")

    def test_no_valid_terms_raises(self):
        """Numeric-only string raises ValueError."""
        with pytest.raises(ValueError):
            parse_linear_expression("123")

    def test_constraint_lte(self):
        """Parse less-than-or-equal constraint."""
        terms, op, rhs = parse_constraint("2*x + 3*y <= 10")
        assert terms == [(2.0, "x"), (3.0, "y")]
        assert op == "<="
        assert rhs == 10.0

    def test_constraint_gte(self):
        """Parse greater-than-or-equal constraint."""
        terms, op, rhs = parse_constraint("x + y >= 5")
        assert terms == [(1.0, "x"), (1.0, "y")]
        assert op == ">="
        assert rhs == 5.0

    def test_constraint_eq(self):
        """Parse equality constraint."""
        terms, op, rhs = parse_constraint("x - y == 0")
        assert terms == [(1.0, "x"), (-1.0, "y")]
        assert op == "=="
        assert rhs == 0.0

    def test_constraint_no_operator_raises(self):
        """Missing operator raises ValueError."""
        with pytest.raises(ValueError):
            parse_constraint("2*x + 3*y")


# ===========================================================================
# LP solver — product-mix problem
# ===========================================================================


@requires_ortools
class TestProductMixLP:
    """Classic product-mix LP with known solution.

    Maximize: 5*x + 4*y
    Subject to:
        6*x + 4*y <= 24   (resource A)
        x + 2*y <= 6      (resource B)
        x, y >= 0

    Optimal: x=3, y=1.5, obj=21.0
    """

    def setup_method(self):
        self.problem = OptimizationProblem(
            name="product_mix",
            sense=ProblemSense.MAXIMIZE,
            variables=[
                Variable(name="x"),
                Variable(name="y"),
            ],
            objective="5*x + 4*y",
            constraints=[
                Constraint(expression="6*x + 4*y <= 24", name="resource_a"),
                Constraint(expression="x + 2*y <= 6", name="resource_b"),
            ],
        )
        provider = get_provider("optimization", backend="ortools")
        self.result = provider.solve(self.problem)

    def test_optimal_status(self):
        assert self.result.status == "optimal"

    def test_objective_value(self):
        assert self.result.objective_value == pytest.approx(21.0, abs=1e-6)

    def test_variable_x(self):
        var_x = next(v for v in self.result.variables if v.name == "x")
        assert var_x.value == pytest.approx(3.0, abs=1e-6)

    def test_variable_y(self):
        var_y = next(v for v in self.result.variables if v.name == "y")
        assert var_y.value == pytest.approx(1.5, abs=1e-6)

    def test_positive_objective(self):
        assert self.result.objective_value > 0

    def test_non_negative_vars(self):
        for v in self.result.variables:
            assert v.value >= -1e-9

    def test_constraint_a_satisfied(self):
        vals = {v.name: v.value for v in self.result.variables}
        assert 6 * vals["x"] + 4 * vals["y"] <= 24 + 1e-6

    def test_constraint_b_satisfied(self):
        vals = {v.name: v.value for v in self.result.variables}
        assert vals["x"] + 2 * vals["y"] <= 6 + 1e-6

    def test_solver_is_glop(self):
        assert self.result.solver == "GLOP"

    def test_metadata_has_iterations(self):
        assert "iterations" in self.result.metadata


# ===========================================================================
# Infeasible problem
# ===========================================================================


@requires_ortools
class TestInfeasibleProblem:
    def test_contradictory_constraints(self):
        """Contradictory constraints return infeasible status."""
        problem = OptimizationProblem(
            name="infeasible",
            sense=ProblemSense.MAXIMIZE,
            variables=[Variable(name="x")],
            objective="x",
            constraints=[
                Constraint(expression="x <= 1"),
                Constraint(expression="x >= 5"),
            ],
        )
        provider = get_provider("optimization", backend="ortools")
        result = provider.solve(problem)
        assert result.status == "infeasible"


# ===========================================================================
# MIP — integer and binary variables
# ===========================================================================


@requires_ortools
class TestMIPIntegerVariables:
    def test_integer_knapsack(self):
        """Integer knapsack: values [60,100,120], weights [10,20,30], capacity 50."""
        problem = OptimizationProblem(
            name="knapsack",
            sense=ProblemSense.MAXIMIZE,
            variables=[
                Variable(name="x0", var_type=VariableType.INTEGER, lb=0, ub=1),
                Variable(name="x1", var_type=VariableType.INTEGER, lb=0, ub=1),
                Variable(name="x2", var_type=VariableType.INTEGER, lb=0, ub=1),
            ],
            objective="60*x0 + 100*x1 + 120*x2",
            constraints=[
                Constraint(expression="10*x0 + 20*x1 + 30*x2 <= 50"),
            ],
        )
        provider = get_provider("optimization", backend="ortools")
        result = provider.solve(problem)

        assert result.status == "optimal"
        assert result.objective_value == pytest.approx(220.0, abs=1e-6)
        assert result.solver == "CBC"

    def test_binary_selection(self):
        """Binary selection: pick exactly 2 out of 3 items to maximize value."""
        problem = OptimizationProblem(
            name="selection",
            sense=ProblemSense.MAXIMIZE,
            variables=[
                Variable(name="a", var_type=VariableType.BINARY),
                Variable(name="b", var_type=VariableType.BINARY),
                Variable(name="c", var_type=VariableType.BINARY),
            ],
            objective="10*a + 20*b + 15*c",
            constraints=[
                Constraint(expression="a + b + c == 2"),
            ],
        )
        provider = get_provider("optimization", backend="ortools")
        result = provider.solve(problem)

        assert result.status == "optimal"
        assert result.objective_value == pytest.approx(35.0, abs=1e-6)

        vals = {v.name: v.value for v in result.variables}
        assert vals["b"] == pytest.approx(1.0, abs=1e-6)
        assert vals["c"] == pytest.approx(1.0, abs=1e-6)
        assert vals["a"] == pytest.approx(0.0, abs=1e-6)


# ===========================================================================
# Provider registration
# ===========================================================================


class TestProviderRegistration:
    def test_optimization_domain_exists(self):
        """Optimization domain shows up in provider listing."""
        providers = list_providers("optimization")
        assert len(providers) >= 1
        domains = {p["domain"] for p in providers}
        assert "optimization" in domains

    def test_explicit_ortools_backend(self):
        """Can request ortools backend by name."""
        providers = list_providers("optimization")
        backends = {p["backend"] for p in providers}
        assert "ortools" in backends

    @requires_ortools
    def test_auto_select_returns_ortools(self):
        """Auto-select picks ortools when available."""
        provider = get_provider("optimization")
        assert provider.backend == "ortools"

    def test_pulp_stub_registered(self):
        """PuLP stub is registered even if not installed."""
        providers = list_providers("optimization")
        backends = {p["backend"] for p in providers}
        assert "pulp" in backends
