"""Weighted scenarios, tail-row generation and the exact Gaussian (outer-approximation) mode."""

import itertools

import numpy as np
import pytest
from scipy.stats import norm

from quantcore.config import Settings
from quantcore.models import sampling as smp
from quantcore.optimization.gaussian_exact import solve_gaussian_exact
from quantcore.optimization.problem import (
    MeanCVaRProblem,
    OptimizationConstraints,
    ProblemSpec,
    TransactionCosts,
    cvar_multiplier,
)
from quantcore.optimization.rowgen import restricted_spec, solve_by_row_generation
from quantcore.optimization.solvers import HighsSolver

SETTINGS = Settings(cuopt_url="http://127.0.0.1:9", mip_rel_gap=1e-6)


def highs(model):
    return HighsSolver(SETTINGS).solve(model, time_limit=120)


def gaussian_inputs(seed=0, n=6):
    rng = np.random.default_rng(seed)
    symbols = [f"S{i}" for i in range(n)]
    prices = rng.uniform(20, 400, n)
    vol = rng.uniform(0.02, 0.06, n)
    corr = 0.4 * np.ones((n, n)) + 0.6 * np.eye(n)
    cov = corr * np.outer(vol, vol)
    mean = rng.uniform(0.0, 0.006, n)
    sectors = ["Tech", "Fin", "Health"] * (n // 3) + ["Tech"] * (n % 3)
    return symbols, prices, mean, cov, sectors


def gaussian_scenarios(mean, cov, S, seed=1):
    z = smp.moment_match(smp.standard_normal_sample(S, len(mean), "sobol", seed))
    return mean + z @ smp.pca_root(cov).T


def test_weighted_scenarios_equal_duplicated_scenarios():
    symbols, prices, mean, cov, sectors = gaussian_inputs()
    scen = gaussian_scenarios(mean, cov, 256)
    dup = np.vstack([scen, scen[:64]])  # the first 64 scenarios counted twice
    p = np.ones(256)
    p[:64] = 2
    p /= p.sum()
    common = dict(symbols=symbols, prices=prices, sectors=sectors, cash=100_000.0, cvar_limit=3_000.0)
    a = MeanCVaRProblem(ProblemSpec(scenarios=dup, **common))
    b = MeanCVaRProblem(ProblemSpec(scenarios=scen, probabilities=p, **common))
    ra, rb = highs(a.build()), highs(b.build())
    assert ra.ok and rb.ok
    da, db = a.decode(ra.x), b.decode(rb.x)
    assert abs(ra.objective - rb.objective) < 1e-6 * abs(ra.objective)
    np.testing.assert_array_equal(da["shares"], db["shares"])
    assert abs(da["cvar"] - db["cvar"]) < 1e-6 and abs(da["expected_pnl"] - db["expected_pnl"]) < 1e-6


def test_expected_returns_drive_the_objective():
    symbols, prices, mean, cov, sectors = gaussian_inputs()
    scen = gaussian_scenarios(mean, cov, 256)
    common = dict(symbols=symbols, prices=prices, scenarios=scen, sectors=sectors, cash=100_000.0, cvar_limit=3_000.0,
                  costs=TransactionCosts(proportional_buy=0.0))
    shifted = mean + np.array([0.01, 0, 0, 0, 0, 0])  # the model says S0 earns 1 % more than the scenarios show
    prob = MeanCVaRProblem(ProblemSpec(expected_returns=shifted, **common))
    res = highs(prob.build())
    d = prob.decode(res.x)
    assert abs(d["expected_pnl"] - (shifted @ d["value"])) < 1e-6  # objective uses the model mean ...
    assert d["expected_pnl"] > d["expected_pnl_scenarios"] + 0.005 * d["value"][0] - 1e-6  # ... not the sample mean
    assert d["shares"][0] > 0
    with pytest.raises(ValueError):
        ProblemSpec(expected_returns=np.zeros(3), **common)


def test_row_generation_matches_full_milp():
    symbols, prices, mean, cov, sectors = gaussian_inputs(seed=2)
    scen = gaussian_scenarios(mean, cov, 2048)
    spec = ProblemSpec(symbols=symbols, prices=prices, scenarios=scen, sectors=sectors, cash=100_000.0,
                       cvar_limit=3_000.0, constraints=OptimizationConstraints(max_weight=0.3, max_sector_weight=0.6))
    full = MeanCVaRProblem(spec)
    res_full = highs(full.build())
    res_rg, problem_rg, info = solve_by_row_generation(spec, highs)
    assert res_rg.ok and info.converged and info.iterations >= 1
    assert max(info.working_set_sizes) < 0.5 * spec.n_scenarios  # solved on a fraction of the rows
    d_full, d_rg = full.decode(res_full.x), problem_rg.decode(res_rg.x)
    assert abs(d_rg["expected_pnl"] - d_full["expected_pnl"]) <= 1e-5 * abs(d_full["expected_pnl"]) + 1e-6
    assert d_rg["cvar"] <= 3_000.0 * (1 + 1e-7)
    # the expanded solution vector is feasible for the full model: u_s >= L_s - alpha for every scenario
    m = full.build()
    Ax = m.A @ res_rg.x
    assert (Ax >= m.row_lb - 1e-6).all() and (Ax <= m.row_ub + 1e-6).all()
    # a restricted spec keeps the full-set probabilities and expected returns
    sub = restricted_spec(spec, np.arange(10))
    assert sub.n_scenarios == 10 and abs(sub.probabilities.sum() - 10 / 2048) < 1e-12
    np.testing.assert_allclose(sub.expected_returns, spec.mean_returns)


def test_row_generation_rebalancing_mode():
    symbols, prices, mean, cov, sectors = gaussian_inputs(seed=3)
    scen = gaussian_scenarios(mean, cov, 1024)
    holdings = np.array([40, 0, 10, 0, 5, 0], dtype=float)
    spec = ProblemSpec(symbols=symbols, prices=prices, scenarios=scen, sectors=sectors, cash=20_000.0,
                       holdings=holdings, cvar_limit=2_500.0, constraints=OptimizationConstraints(max_weight=0.3, max_sector_weight=0.6))
    res, problem, info = solve_by_row_generation(spec, highs)
    assert res.ok and info.converged
    d = problem.decode(res.x)
    np.testing.assert_array_equal(d["shares"], holdings + d["buys"] - d["sells"])
    assert d["cvar"] <= 2_500.0 * (1 + 1e-7) and d["cash_end"] >= -1e-6


def brute_force_gaussian(prices, mean, cov, sectors, cash, limit, beta, max_weight, costs):
    """Enumerate all integer share vectors within the caps (small instance)."""
    ub = np.floor(max_weight * cash / prices).astype(int)
    k = cvar_multiplier(beta)
    best, best_n = -np.inf, None
    for shares in itertools.product(*[range(u + 1) for u in ub]):
        s = np.array(shares, dtype=float)
        v = s * prices
        tc = costs * v.sum()
        if v.sum() + tc > cash + 1e-9:
            continue
        pnl = mean @ v - tc
        cvar = -pnl + k * np.sqrt(v @ cov @ v)
        if cvar <= limit + 1e-9 and pnl > best:
            best, best_n = pnl, s
    return best, best_n


def test_gaussian_exact_matches_brute_force():
    prices = np.array([400.0, 650.0, 900.0])
    vol = np.array([0.03, 0.05, 0.04])
    corr = np.array([[1, 0.5, 0.2], [0.5, 1, 0.3], [0.2, 0.3, 1]])
    cov = corr * np.outer(vol, vol)
    mean = np.array([0.003, 0.006, 0.004])
    sectors = ["A", "B", "C"]
    cash, limit, beta = 10_000.0, 500.0, 0.95
    costs = TransactionCosts(proportional_buy=0.002)
    spec = ProblemSpec(symbols=list("XYZ"), prices=prices, scenarios=None, sectors=sectors, cash=cash, cvar_limit=limit,
                       confidence=beta, gaussian_mean=mean, gaussian_cov=cov, costs=costs,
                       constraints=OptimizationConstraints(max_weight=0.6, max_sector_weight=1.0))
    res, problem, info = solve_gaussian_exact(spec, highs, tol=1e-9)
    assert res.ok and info.converged and info.cuts >= 2
    d = problem.decode(res.x)
    best, best_n = brute_force_gaussian(prices, mean, cov, sectors, cash, limit, beta, 0.6, 0.002)
    np.testing.assert_array_equal(d["shares"], best_n)
    assert abs(d["expected_pnl"] - best) < 1e-6
    assert d["cvar"] <= limit * (1 + 1e-9)
    # closed-form risk figures of the decoded portfolio
    v = d["value"]
    sigma = np.sqrt(v @ cov @ v)
    assert abs(d["sigma"] - sigma) < 1e-9
    assert abs(d["var"] - (-(mean @ v - d["total_cost"]) + norm.ppf(beta) * sigma)) < 1e-9
    assert res.reproducible and "outer approximation" in res.message


def test_gaussian_exact_other_objectives_and_validation():
    symbols, prices, mean, cov, sectors = gaussian_inputs(seed=4)
    base = dict(symbols=symbols, prices=prices, scenarios=None, sectors=sectors, cash=100_000.0,
                gaussian_mean=mean, gaussian_cov=cov)
    # min CVaR at a target expected P&L: converges, meets the target, epigraph equals the true CVaR
    spec = ProblemSpec(objective="min_cvar", target_pnl=150.0, **base)
    res, problem, info = solve_gaussian_exact(spec, highs)
    d = problem.decode(res.x)
    assert info.converged and d["expected_pnl"] >= 150.0 - 1e-6
    assert abs(d["cvar"] - d["cvar_model"]) <= 1e-6 * max(d["cvar"], 1.0)
    # utility objective
    spec_u = ProblemSpec(objective="utility", risk_aversion=0.05, **base)  # at 0.5 staying in cash is optimal
    res_u, problem_u, info_u = solve_gaussian_exact(spec_u, highs)
    du = problem_u.decode(res_u.x)
    assert info_u.converged and du["invested"] > 0
    # the CVaR-limited solution is on the mean-variance frontier: no feasible portfolio with the same
    # expected P&L has a smaller sigma (checked on the continuous relaxation via the cut model)
    spec_m = ProblemSpec(cvar_limit=3_000.0, **base)
    res_m, problem_m, _ = solve_gaussian_exact(spec_m, highs)
    dm = problem_m.decode(res_m.x)
    assert dm["cvar"] <= 3_000.0 * (1 + 1e-6)
    with pytest.raises(ValueError):
        ProblemSpec(symbols=symbols, prices=prices, scenarios=None, sectors=sectors, cash=1.0, cvar_limit=1.0)
    with pytest.raises(ValueError):
        MeanCVaRProblem(spec_m).build()  # no cuts yet
    with pytest.raises(ValueError):
        solve_by_row_generation(spec_m, highs)


def test_moment_matched_sobol_scenarios_reproduce_the_exact_optimum():
    """The sample-average approximation with Gaussian scenarios (used linearly) converges to the exact
    mean-variance-equivalent optimum; plain Monte Carlo at the same size does not."""
    symbols, prices, mean, cov, sectors = gaussian_inputs(seed=6, n=8)
    common = dict(symbols=symbols, prices=prices, sectors=sectors, cash=1_000_000.0, cvar_limit=30_000.0,
                  costs=TransactionCosts(proportional_buy=0.0),
                  constraints=OptimizationConstraints(max_weight=0.25, max_sector_weight=0.6))
    exact_spec = ProblemSpec(scenarios=None, gaussian_mean=mean, gaussian_cov=cov, **common)
    res_e, prob_e, _info = solve_gaussian_exact(exact_spec, highs, tol=1e-8)
    d_e = prob_e.decode(res_e.x)
    k = cvar_multiplier(0.95)

    def true_cvar(v):
        return -(mean @ v) + k * np.sqrt(v @ cov @ v)

    # moment-matched Sobol' scenarios, model mean in the objective
    scen = gaussian_scenarios(mean, cov, 4096, seed=3)
    prob_s = MeanCVaRProblem(ProblemSpec(scenarios=scen, expected_returns=mean, **common))
    d_s = prob_s.decode(highs(prob_s.build()).x)
    assert abs(d_s["expected_pnl"] - d_e["expected_pnl"]) <= 0.01 * d_e["expected_pnl"]
    assert true_cvar(d_s["value"]) <= 30_000.0 * 1.02
    # plain Monte Carlo of the same size with the sample mean in the objective (the old default)
    rng = np.random.default_rng(3)
    mc = mean + rng.standard_normal((4096, 8)) @ np.linalg.cholesky(cov).T
    prob_m = MeanCVaRProblem(ProblemSpec(scenarios=mc, **common))
    d_m = prob_m.decode(highs(prob_m.build()).x)
    true_pnl_mc = mean @ d_m["value"]
    assert abs(true_pnl_mc - d_e["expected_pnl"]) > abs(d_s["expected_pnl"] - d_e["expected_pnl"])
    assert d_m["expected_pnl"] > true_pnl_mc  # in-sample optimism of the plain MC solution


def test_row_generation_stops_at_the_time_budget():
    """The time limit is the budget of the whole request, not of each pass.

    Every pass is a MILP of its own, so a deadline that has already passed must end the loop with
    the incumbent (feasible, not proven optimal) instead of starting another solve.
    """
    import time

    symbols, prices, mean, cov, sectors = gaussian_inputs(seed=2)
    scen = gaussian_scenarios(mean, cov, 2048)
    spec = ProblemSpec(symbols=symbols, prices=prices, scenarios=scen, sectors=sectors, cash=100_000.0,
                       cvar_limit=3_000.0, constraints=OptimizationConstraints(max_weight=0.3, max_sector_weight=0.6))
    calls = []

    def counted(model):
        calls.append(1)
        return highs(model)

    res, problem, info = solve_by_row_generation(spec, counted, deadline=time.monotonic() - 1.0)
    assert len(calls) == 1 and info.iterations == 1 and not info.converged
    assert res.status == "FeasibleFound" and "time limit" in res.message
    d = problem.decode(res.x)  # the incumbent is a real portfolio, measured on all scenarios
    assert d["shares"].sum() > 0 and np.isfinite(d["cvar"])
    # with a budget in the future the loop runs to convergence as before
    res2, _, info2 = solve_by_row_generation(spec, highs, deadline=time.monotonic() + 600)
    assert info2.converged and res2.status == "Optimal"


def test_the_solve_function_hands_on_the_remaining_budget():
    from quantcore.optimization.service import OptimizationRequest, OptimizationService

    svc = OptimizationService.__new__(OptimizationService)  # no data needed for this
    svc.settings = SETTINGS
    req = OptimizationRequest(time_limit=30.0, solver="highs", cvar_limit=1.0)
    limits = []
    deadline = __import__("time").monotonic() + 30.0

    class Recorder:
        def __init__(self, settings):
            pass

        def solve(self, model, time_limit=None):
            limits.append(time_limit)
            return None

    import quantcore.optimization.service as service_module
    original = service_module.HighsSolver
    service_module.HighsSolver = Recorder
    try:
        fn = svc._solve_fn(req, prefer_local=True, deadline=deadline)
        fn(None)
        fn(None)
    finally:
        service_module.HighsSolver = original
    assert 0 < limits[1] <= limits[0] <= 30.0  # each call gets what is left, never the full limit again