"""The large-scale driver against the integer MILP on a synthetic universe (no market data needed).

    python scripts/experiments/large_scale.py 300 16384            # N stocks, S scenarios
    python scripts/experiments/large_scale.py 750 16384,65536 --cardinality 40 --cuopt http://spark-6d4a:5000

Prints, per size: the full-MILP row generation (quantcore's existing path) and the driver, with wall
time, objective, the driver's bound and gap, working sets and solver calls. ``--cuopt URL`` adds the
cuOpt server as the large-model solver of the driver (and as the MILP solver of the reference run).
"""

from __future__ import annotations

import argparse
import time

import numpy as np

from quantcore.config import Settings
from quantcore.optimization.large_scale import solve_large_scale
from quantcore.optimization.problem import OptimizationConstraints, ProblemSpec
from quantcore.optimization.rowgen import solve_by_row_generation
from quantcore.optimization.solvers import HighsSolver, solve_with_fallback


def market(N: int, S: int, seed: int = 0, K: int = 6, horizon_days: int = 5):
    """A K-factor universe with Student-t(4) scale-mixture scenarios (the benchmark's generator)."""
    rng = np.random.default_rng(seed)
    prices = np.exp(rng.uniform(np.log(10.0), np.log(600.0), N))
    sectors = [f"sector_{s}" for s in rng.integers(0, 10, N)]
    f_vol = np.concatenate([[0.18], rng.uniform(0.08, 0.20, K - 1)])
    F = np.diag((f_vol ** 2) * horizon_days / 252.0)
    B = np.column_stack([rng.normal(1.0, 0.3, N)] + [rng.normal(0.0, 0.5, N) for _ in range(K - 1)])
    D = (rng.uniform(0.15, 0.40, N) ** 2) * horizon_days / 252.0
    mu = rng.normal(0.06, 0.08, N) * horizon_days / 252.0
    rng2 = np.random.default_rng(seed + 7919)
    z = rng2.standard_normal((S, K)) @ np.linalg.cholesky(F).T
    e = rng2.standard_normal((S, N)) * np.sqrt(D)
    r = mu + z @ B.T + e
    w = np.sqrt(4.0 / rng2.chisquare(4.0, size=(S, 1)))
    r = mu + (r - mu) * w * np.sqrt(0.5)
    return prices, sectors, mu, r


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("N", type=int)
    ap.add_argument("S", help="comma-separated scenario counts")
    ap.add_argument("--cardinality", type=int, default=None, help="max_positions (adds binaries)")
    ap.add_argument("--cuopt", default=None, help="cuOpt URL for the large models (default: HiGHS only)")
    ap.add_argument("--no-reference", action="store_true", help="skip the full-MILP reference run")
    ap.add_argument("--time-limit", type=float, default=600.0)
    ap.add_argument("--seed", type=int, default=0)
    a = ap.parse_args()
    settings = Settings(cuopt_url=a.cuopt or "http://127.0.0.1:9", mip_rel_gap=1e-4, solver_fallback=True)
    highs = lambda m: HighsSolver(settings).solve(m, time_limit=a.time_limit)
    cuopt = (lambda m: solve_with_fallback(m, settings, solver="cuopt", time_limit=a.time_limit)) if a.cuopt else None
    W = 1_000_000.0
    print(f"{'N':>4} {'S':>6} {'K':>4} | {'MILP s':>8} {'objective':>12} | {'driver s':>8} {'objective':>12} "
          f"{'bound':>12} {'gap':>8} {'stage':>7} {'rows':>6} {'calls':>5} {'escal':>5}")
    for S in (int(x) for x in a.S.split(",")):
        prices, sectors, mu, scen = market(a.N, S, a.seed)
        cons = OptimizationConstraints(max_weight=0.10, max_sector_weight=0.25, max_positions=a.cardinality)
        spec = ProblemSpec(symbols=[f"S{i}" for i in range(a.N)], prices=prices, scenarios=scen, sectors=sectors,
                           cash=W, cvar_limit=0.05 * W, expected_returns=mu, constraints=cons)
        ref = "-"
        ref_obj = float("nan")
        if not a.no_reference:
            t0 = time.perf_counter()
            res, problem, _ = solve_by_row_generation(spec, cuopt or highs)
            ref = f"{time.perf_counter() - t0:8.1f}"
            ref_obj = problem.decode(res.x)["expected_pnl"] if res.ok else float("nan")
        t0 = time.perf_counter()
        res, problem, info = solve_large_scale(spec, highs, cuopt, mip_gap=1e-4)
        t_ls = time.perf_counter() - t0
        obj = problem.decode(res.x)["expected_pnl"] if res.ok else float("nan")
        print(f"{a.N:>4} {S:>6} {a.cardinality or 0:>4} | {ref:>8} {ref_obj:12.2f} | {t_ls:8.1f} {obj:12.2f} "
              f"{(info.bound or float('nan')):12.2f} {(info.gap if info.gap is not None else float('nan')):8.1e} "
              f"{info.integer_stage:>7} {info.working_set:>6} {len(info.calls):>5} {info.escalated!s:>5}", flush=True)
        for c in info.calls:
            print(f"      {c['stage']:>10} rows={c['rows']:>6} vars={c['vars']:>7} {c['solver']:>6} {c['status']:>13} "
                  f"{c['time']:7.2f}s")


if __name__ == "__main__":
    main()