"""Time scaling of the mean-CVaR LP / MILP in the number of scenarios S (HiGHS locally, optionally cuOpt)."""
import sys, time, json, zlib
import numpy as np
from scipy.stats import norm, qmc
from quantcore.config import get_settings
from quantcore.data.market_data import MarketDataService
from quantcore.optimization.service import OptimizationService, OptimizationRequest
from quantcore.optimization.problem import MeanCVaRProblem, ProblemSpec, OptimizationConstraints, TransactionCosts
from quantcore.optimization.solvers import HighsSolver, CuOptSolver, to_cuopt_payload
from quantcore.models.covariance import CovarianceModel, symmetric_inv_sqrt

N_ASSETS = int(sys.argv[1]); SIZES = [int(x) for x in sys.argv[2].split(",")]; MODE = sys.argv[3]  # highs | cuopt
BETA = 0.95; W = 1_000_000.0; LAMBDA = 0.05 * W; H = 5
settings = get_settings(); md = MarketDataService(settings); svc = OptimizationService(settings, md)
req = OptimizationRequest(max_assets=N_ASSETS, solver="highs", cash=W)
symbols = svc.candidate_symbols(req)
mu_ann, _, _ = svc.expected_returns(req, symbols)
r = md.returns(symbols, currency="base", freq="daily").dropna()
cov_h = CovarianceModel.from_dict(req.covariance).estimate(r).values * H
mu_h = (mu_ann / 252).reindex(symbols).values * H
prices = md.latest_prices(symbols, currency="base").reindex(symbols).values
sectors = md.sectors(symbols).fillna("Unknown").tolist()
N = len(symbols)
vals, vecs = np.linalg.eigh(cov_h); o = np.argsort(vals)[::-1]; A = vecs[:, o] * np.sqrt(vals[o])

def scen_sobol_mm(S, seed=0):
    z = norm.ppf(qmc.Sobol(d=N, scramble=True, seed=seed).random(S))
    z = z - z.mean(0); z = z @ symmetric_inv_sqrt(z.T @ z / S)
    return mu_h + z @ A.T

print(f"N={N}, beta={BETA}, mode={MODE}")
print(f"{'S':>8} {'nnz':>10} {'gen s':>7} {'build s':>8} {'LP s':>8} {'MILP s':>8} {'status':>10} {'E[P&L]':>9} {'json MB':>8} {'cuopt s':>8} {'cu solve':>9}")
for S in SIZES:
    t0 = time.time(); scen = scen_sobol_mm(S); t_gen = time.time() - t0
    spec = ProblemSpec(symbols=symbols, prices=prices, scenarios=scen, sectors=sectors, cash=W, holdings=np.zeros(N),
                       confidence=BETA, cvar_limit=LAMBDA, objective="max_return", constraints=OptimizationConstraints(),
                       costs=TransactionCosts())
    prob = MeanCVaRProblem(spec)
    t0 = time.time(); model = prob.build(); t_build = time.time() - t0
    nnz = model.A.nnz
    row = f"{S:>8} {nnz:>10} {t_gen:>7.1f} {t_build:>8.1f}"
    if MODE == "highs":
        lp = MeanCVaRProblem(spec).build(); lp.var_types = ["C"] * len(lp.var_types)
        t0 = time.time(); sol_lp = HighsSolver().solve(lp, time_limit=300); t_lp = time.time() - t0
        t0 = time.time(); sol = HighsSolver().solve(model, time_limit=300, mip_gap=1e-4); t_mip = time.time() - t0
        obj = prob.decode(sol.x)["expected_pnl"] if sol.x is not None else float("nan")
        row += f" {t_lp:>8.1f} {t_mip:>8.1f} {sol.status:>10} {obj:>9.1f}"
    else:
        t0 = time.time(); payload = to_cuopt_payload(model, 300.0, {"mip_relative_gap": 1e-4, "log_to_console": False})
        body = json.dumps(payload); mb = len(body) / 1e6; t_json = time.time() - t0
        t0 = time.time(); sol = CuOptSolver(settings).solve(model, time_limit=300.0, extra_config={"mip_relative_gap": 1e-4, "log_to_console": False}); t_cu = time.time() - t0
        obj = prob.decode(sol.x)["expected_pnl"] if sol.x is not None else float("nan")
        row += f" {'':>8} {'':>8} {sol.status:>10} {obj:>9.1f} {mb:>8.1f} {t_cu:>8.1f} {sol.solve_time:>9.1f}"
    print(row, flush=True)