"""Sampling noise of the mean-CVaR LP (weights) under the empirical-innovation model with different resampling schemes.

'truth' = a common evaluation set of 500k plain bootstrap draws (sampling error of CVaR for a fixed portfolio ~0.1%).
"""
import sys, time
import numpy as np, pandas as pd
from scipy.stats import 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_sqrt, symmetric_inv_sqrt
from quantcore.models.simulation import EmpiricalInnovationSimulator
from quantcore.measurement.risk import scenario_cvar

N_ASSETS = int(sys.argv[1]) if len(sys.argv) > 1 else 30
S = int(sys.argv[2]) if len(sys.argv) > 2 else 2048
N_SEEDS = int(sys.argv[3]) if len(sys.argv) > 3 else 8
BETA = float(sys.argv[4]) if len(sys.argv) > 4 else 0.95
W, H = 1_000_000.0, 5
LAMBDA = 0.05 * W
S_EVAL = 500_000

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_model = CovarianceModel.from_dict(req.covariance)
sim = EmpiricalInnovationSimulator(cov_model, max_history=req.innovation_history, standardize=req.innovation_standardize)
t0 = time.time()
eps = sim.innovations(r)
e = eps.values - eps.values.mean(axis=0)
e, info = sim._restandardize(e)
T, N = e.shape
cov = cov_model.estimate(r).values
A = symmetric_sqrt(cov)
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()
print(f"N={N} assets, T={T} innovations ({time.time()-t0:.0f}s), S={S}, h={H}, beta={BETA}, {info}")
# market-factor ordering of the innovation dates (projection on the first eigenvector of Sigma_T)
w1 = np.linalg.eigh(cov)[1][:, -1]
factor = e @ (A @ w1)                      # innovation of the dominant direction
rank_order = np.argsort(factor)            # dates sorted from worst to best market innovation


def moment_match(x, target_var):
    x = x - x.mean(axis=0)
    c = x.T @ x / len(x)
    return x @ symmetric_inv_sqrt(c) * np.sqrt(target_var)


def summed_innovations(kind, seed):
    rng = np.random.default_rng(seed)
    if kind == "bootstrap":                                   # current implementation
        idx = rng.integers(0, T, size=(S, H))
        return e[idx].sum(axis=1)
    if kind == "bootstrap+mm":
        idx = rng.integers(0, T, size=(S, H))
        return moment_match(e[idx].sum(axis=1), H)
    if kind == "balanced+mm":                                 # every date used equally often
        reps = int(np.ceil(S * H / T))
        idx = rng.permutation(np.tile(np.arange(T), reps))[: S * H].reshape(S, H)
        return moment_match(e[idx].sum(axis=1), H)
    if kind == "lhs-ranks+mm":                                # stratified on the market-factor rank, per day
        u = qmc.LatinHypercube(d=H, scramble=True, seed=seed).random(S)
        idx = rank_order[np.minimum((u * T).astype(int), T - 1)]
        return moment_match(e[idx].sum(axis=1), H)
    if kind == "sobol-ranks+mm":
        u = qmc.Sobol(d=H, scramble=True, seed=seed).random(S)
        idx = rank_order[np.minimum((u * T).astype(int), T - 1)]
        return moment_match(e[idx].sum(axis=1), H)
    if kind == "windows+mm":                                  # deterministic: all overlapping h-day windows
        idx = np.arange(T - H + 1)[:, None] + np.arange(H)[None, :]
        return moment_match(e[idx].sum(axis=1), H)
    if kind == "windows":
        idx = np.arange(T - H + 1)[:, None] + np.arange(H)[None, :]
        return e[idx].sum(axis=1)
    raise ValueError(kind)


def scenarios(kind, seed):
    return mu_h + summed_innovations(kind, seed) @ A.T


def solve_lp(scen):
    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(proportional_buy=0.0, proportional_sell=0.0))
    prob = MeanCVaRProblem(spec)
    model = prob.build()
    model.var_types = ["C"] * len(model.var_types)
    sol = HighsSolver().solve(model, time_limit=300)
    n = sol.x[prob.i_n]
    pnl = scen @ (n * prices)
    return n, float(pnl.mean()), scenario_cvar(-pnl, BETA)[1]


# common evaluation set (plain bootstrap, independent seed)
rng_eval = np.random.default_rng(987654321)
eval_scen = mu_h + e[rng_eval.integers(0, T, size=(S_EVAL, H))].sum(axis=1) @ A.T
print(f"evaluation set: {S_EVAL} draws, mean of column means {eval_scen.mean():.2e}")


def true_stats(n):
    pnl = eval_scen @ (n * prices)
    return float(pnl.mean()), scenario_cvar(-pnl, BETA)[1]


results = {}
KINDS = sys.argv[5].split(",") if len(sys.argv) > 5 else ["bootstrap", "bootstrap+mm", "balanced+mm", "lhs-ranks+mm", "sobol-ranks+mm"]
for kind in KINDS:
    rows, t0 = [], time.time()
    for seed in range(1 if kind.startswith("windows") else N_SEEDS):
        n, epnl_in, cvar_in = solve_lp(scenarios(kind, seed))
        epnl_t, cvar_t = true_stats(n)
        rows.append({"epnl_in": epnl_in, "cvar_in": cvar_in, "epnl_true": epnl_t, "cvar_true": cvar_t, "n_pos": int((n > 1e-6).sum()), "w": n * prices / W})
    df = pd.DataFrame(rows)
    ws = np.vstack(df["w"].values)
    pair = [np.abs(ws[i] - ws[j]).sum() for i in range(len(ws)) for j in range(i + 1, len(ws))]
    results[kind] = {
        "E[P&L] in-sample mean": df.epnl_in.mean(), "E[P&L] in-sample sd": df.epnl_in.std(),
        "E[P&L] true mean": df.epnl_true.mean(), "E[P&L] true sd": df.epnl_true.std(),
        "true CVaR mean": df.cvar_true.mean(), "true CVaR sd": df.cvar_true.std(),
        "CVaR excess over limit (%)": 100 * (df.cvar_true.mean() / LAMBDA - 1),
        "positions": df.n_pos.mean(), "turnover between seeds (%)": 100 * np.mean(pair) if pair else float("nan"), "seconds": time.time() - t0,
    }
    print(f"{kind:15s} done in {time.time()-t0:.0f}s")

out = pd.DataFrame(results).T
pd.set_option("display.width", 250); pd.set_option("display.float_format", lambda v: f"{v:10.2f}")
print(f"\nEMPIRICAL model, continuous (weights) problem, CVaR limit {LAMBDA:.0f} CHF, beta={BETA}")
print(out.T.to_string())