"""Tail-row generation for the RU mean-CVaR MILP: solve on a working set of scenarios, add violated ones, repeat.

Full problem (new mode, proportional buy costs c):  max (g_full - c p)'n   s.t.
   g_s'n - c p'n + alpha + u_s >= 0           (all s)            [u_s >= L_s - alpha, L_s = -(g_s - c p)'n]
   alpha + sum_s u_s / ((1-beta) S) <= Lambda
   (1 + c) p'n <= W,  p_i n_i <= w_max W,  sector caps,  n integer >= 0
Restricted problem: rows only for s in the working set (u_s = 0 elsewhere) -> a relaxation; optimal for the
full problem as soon as no omitted scenario has L_s > alpha*.
"""
import sys, time
import numpy as np
import scipy.sparse as sps
from scipy.optimize import milp, LinearConstraint, Bounds
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
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(",")]
FULL = len(sys.argv) > 3 and sys.argv[3] == "full"
BETA = 0.95; W = 1_000_000.0; LAMBDA = 0.05 * W; H = 5; C = 0.001; WMAX = 0.10; SMAX = 0.25; GAP = 1e-4
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 np.expm1(mu_h + z @ A.T)  # simple returns, as in production


def solve_restricted(gain, work, S_full, mean_gain):
    """MILP on the working set `work` (indices into gain rows). Returns n, alpha, objective, seconds."""
    K = len(work)
    p = prices
    # variables: n (N), alpha (1), u (K)
    nv = N + 1 + K
    c = np.zeros(nv); c[:N] = -(mean_gain - C * p)   # minimise negative expected P&L (net of costs)
    rows, cols, vals_ = [], [], []
    lb, ub = [], []
    ri = 0
    G = gain[work] - C * p[None, :]                    # net gain per share in the working scenarios
    # scenario rows: G_s' n + alpha + u_s >= 0
    rows.extend(np.repeat(np.arange(K), N)); cols.extend(np.tile(np.arange(N), K)); vals_.extend(G.ravel())
    rows.extend(np.arange(K)); cols.extend([N] * K); vals_.extend([1.0] * K)
    rows.extend(np.arange(K)); cols.extend(N + 1 + np.arange(K)); vals_.extend([1.0] * K)
    lb.extend([0.0] * K); ub.extend([np.inf] * K); ri = K
    # cvar row: alpha + sum u / ((1-beta) S_full) <= Lambda
    rows.append(ri); cols.append(N); vals_.append(1.0)
    rows.extend([ri] * K); cols.extend(N + 1 + np.arange(K)); vals_.extend([1.0 / ((1 - BETA) * S_full)] * K)
    lb.append(-np.inf); ub.append(LAMBDA); ri += 1
    # budget (1+c) p'n <= W
    rows.extend([ri] * N); cols.extend(range(N)); vals_.extend((1 + C) * p); lb.append(-np.inf); ub.append(W); ri += 1
    # sector caps
    for sec in sorted(set(sectors)):
        idx = [i for i, s in enumerate(sectors) if s == sec]
        rows.extend([ri] * len(idx)); cols.extend(idx); vals_.extend(p[idx]); lb.append(-np.inf); ub.append(SMAX * W); ri += 1
    Am = sps.csr_matrix((vals_, (rows, cols)), shape=(ri, nv))
    vlb = np.zeros(nv); vub = np.full(nv, np.inf); vub[:N] = np.floor(WMAX * W / p + 1e-9); vlb[N] = -np.inf
    integrality = np.zeros(nv); integrality[:N] = 1
    t0 = time.time()
    res = milp(c, constraints=LinearConstraint(Am, np.array(lb), np.array(ub)), integrality=integrality,
               bounds=Bounds(vlb, vub), options={"mip_rel_gap": GAP, "time_limit": 1800})
    n = np.rint(res.x[:N]); alpha = res.x[N]
    return n, alpha, -res.fun, time.time() - t0


def row_generation(gain, verbose=True):
    S = len(gain)
    mean_gain = gain.mean(0)
    p = prices
    # initial working set: the worst 2 (1-beta) S scenarios of the equal-weight (cap-respecting) portfolio
    n0 = np.floor(W / N / p)
    L0 = -(gain - C * p) @ n0
    K0 = int(2 * (1 - BETA) * S)
    work = np.argsort(L0)[-K0:]
    t_solve = 0.0
    for it in range(1, 30):
        n, alpha, obj, dt = solve_restricted(gain, work, S, mean_gain)
        t_solve += dt
        L = -(gain - C * p) @ n
        viol = np.setdiff1d(np.where(L > alpha + 1e-9)[0], work)
        if verbose:
            print(f"    iter {it}: |work|={len(work)}, MILP {dt:.1f}s, obj {obj:.2f}, alpha {alpha:.1f}, violated {len(viol)}", flush=True)
        if len(viol) == 0:
            u = np.maximum(L - alpha, 0)
            cvar = alpha + u.sum() / ((1 - BETA) * S)
            return n, obj, cvar, it, t_solve
        work = np.union1d(work, viol)
    raise RuntimeError("no convergence")


print(f"N={N}, beta={BETA}, gap={GAP}")
for S in SIZES:
    gain = scen_sobol_mm(S) * prices[None, :]
    t0 = time.time()
    n, obj, cvar, iters, t_solve = row_generation(gain)
    t_rg = time.time() - t0
    print(f"S={S:>7}: row generation {iters} iterations, {t_rg:.1f}s total ({t_solve:.1f}s in MILPs), E[P&L]={obj:.2f}, CVaR={cvar:.1f}, positions={(n>0).sum()}", flush=True)
    if FULL:
        t0 = time.time()
        n_f, alpha_f, obj_f, dt = solve_restricted(gain, np.arange(S), S, gain.mean(0))
        print(f"S={S:>7}: full MILP {dt:.1f}s, E[P&L]={obj_f:.2f}, same shares: {np.array_equal(n, n_f)}, max |dn| = {np.abs(n - n_f).max():.0f}", flush=True)