import numpy as np
import pandas as pd
import pytest
from scipy import stats

from quantcore.measurement import risk as rk


def test_normal_var_cvar_reference():
    # sigma = 1 %, mu = 0: VaR95 = 1.6449 %, CVaR95 = 2.0627 %, VaR99 = 2.3263 %, CVaR99 = 2.6652 %
    assert rk.normal_var(0.0, 0.01, 0.95) == pytest.approx(0.016449, abs=1e-6)
    assert rk.normal_cvar(0.0, 0.01, 0.95) == pytest.approx(0.020627, abs=1e-6)
    assert rk.normal_var(0.0, 0.01, 0.99) == pytest.approx(0.023263, abs=1e-6)
    assert rk.normal_cvar(0.0, 0.01, 0.99) == pytest.approx(0.026652, abs=1e-6)
    # a positive mean lowers the loss quantile one-for-one
    assert rk.normal_var(0.001, 0.01, 0.95) == pytest.approx(0.015449, abs=1e-6)


def test_historical_var_cvar_reference():
    r = pd.Series(np.arange(-0.10, 0.10, 0.001))  # 200 equally spaced returns
    var95 = rk.historical_var(r, 0.95)
    assert var95 == pytest.approx(-np.quantile(r, 0.05))
    cvar95 = rk.historical_cvar(r, 0.95)
    tail = r[r <= np.quantile(r, 0.05)]
    assert cvar95 == pytest.approx(-tail.mean())
    assert cvar95 > var95


def test_scenario_cvar_is_tail_mean():
    losses = np.array([1.0, 5.0, 3.0, 9.0, 7.0, 2.0, 8.0, 4.0, 6.0, 10.0])  # S = 10
    var, cvar = rk.scenario_cvar(losses, 0.80)  # k = 2 -> mean of the two largest
    assert cvar == pytest.approx((10 + 9) / 2)
    assert var == pytest.approx(9.0)
    var, cvar = rk.scenario_cvar(losses, 0.95)  # k = 0.5 -> fractional tail: 0.5 * 10 / 0.5
    assert cvar == pytest.approx(10.0)
    var, cvar = rk.scenario_cvar(losses, 0.85)  # k = 1.5 -> (10 + 0.5 * 9) / 1.5
    assert cvar == pytest.approx((10 + 0.5 * 9) / 1.5)
    assert var == pytest.approx(9.0)


def test_scenario_cvar_matches_ru_minimisation():
    rng = np.random.default_rng(0)
    losses = rng.standard_normal(500)
    beta = 0.95
    _, cvar = rk.scenario_cvar(losses, beta)
    grid = np.linspace(-1, 4, 20001)
    ru = np.min([a + np.maximum(losses - a, 0).mean() / (1 - beta) for a in grid])
    assert cvar == pytest.approx(ru, abs=1e-3)


def test_loss_to_money():
    assert rk.loss_to_money(0.0, 1000.0) == 0.0
    assert rk.loss_to_money(np.log(1 / 0.9), 1000.0) == pytest.approx(100.0)


def test_asset_risk_table_columns_and_scaling():
    rng = np.random.default_rng(1)
    r = pd.DataFrame(rng.standard_normal((1500, 2)) * 0.01, index=pd.bdate_range("2019-01-01", periods=1500), columns=["A", "B"])
    day = rk.asset_risk_table(r, "daily", 200, 0.94, (0.95,), "day")
    week = rk.asset_risk_table(r, "daily", 200, 0.94, (0.95,), "week")
    assert {"vol_rolling_ann", "vol_ewm_ann", "var95_normal_rolling", "cvar95_hist"} <= set(day.columns)
    assert week["horizon_periods"].iloc[0] == 5
    # sqrt(5) scaling of the Gaussian VaR for a zero-mean window (approximately, mean is small)
    ratio = week["var95_normal_rolling"] / day["var95_normal_rolling"]
    assert ratio.between(2.0, 2.5).all()
    assert (day["cvar95_normal_ewm"] > day["var95_normal_ewm"]).all()


def test_portfolio_risk_normal_vs_scenarios(md):
    svc = rk.RiskService(md)
    holdings = {"S0": 100, "S1": 50, "S2": 200}
    pr = svc.portfolio_risk(holdings, cash=1000.0, horizon="week", method="normal")
    assert pr.value > 0 and set(pr.confidence) == {"0.95", "0.99"}
    assert pr.confidence["0.99"]["cvar_money"] > pr.confidence["0.95"]["cvar_money"] > 0
    # historical estimate is of the same order of magnitude
    ph = svc.portfolio_risk(holdings, cash=1000.0, horizon="week", method="historical")
    assert 0.3 < ph.confidence["0.95"]["cvar_money"] / pr.confidence["0.95"]["cvar_money"] < 3.0
    # scenarios: Gaussian scenarios with the same covariance reproduce the normal CVaR
    rng = np.random.default_rng(0)
    r = md.returns(list(holdings), currency="base", freq="daily")
    from quantcore.measurement import estimators as est
    cov = est.rolling_cov(r, 200)
    mu_h = r.iloc[-200:].mean().values * 5
    scen = rng.multivariate_normal(mu_h, cov.values * 5, size=200_000)
    ps = svc.portfolio_risk(holdings, cash=1000.0, horizon="week", method="scenarios", scenarios=scen)
    # differences: sampling noise and the exact money mapping (exp) vs. the log-return closed form
    assert ps.confidence["0.95"]["cvar_money"] == pytest.approx(pr.confidence["0.95"]["cvar_money"], rel=0.06)


def test_a_one_asset_portfolio_is_a_valid_request(md):
    """``np.cov`` of a single column returns a 0-d array; the rolling estimator has to keep it 2-d.

    The simplest call anyone makes -- the risk of a single position -- went through this path.
    """
    svc = rk.RiskService(md)
    from quantcore.measurement import estimators as est

    r = md.returns(["S0"], currency="base", freq="daily")
    assert est.rolling_cov(r, 200).shape == (1, 1)
    assert est.ewm_cov(r, 0.94).shape == (1, 1)
    for method in ("normal", "historical"):
        for cov_method in ("rolling", "ewm"):
            pr = svc.portfolio_risk({"S0": 100}, cash=1000.0, horizon="week", method=method,
                                    cov_method=cov_method)
            assert pr.value > 0 and pr.confidence["0.95"]["cvar_money"] > 0
    # a one-asset portfolio is the asset: the parametric VaR is z * sd_h - mean_h exactly
    one = svc.portfolio_risk({"S0": 100}, cash=0.0, horizon="week", method="normal", cov_method="rolling")
    sd_week = float(np.sqrt(est.rolling_cov(r, 200).iloc[0, 0] * 5))
    mean_week = float(r.iloc[-200:].mean().iloc[0] * 5)
    assert one.confidence["0.95"]["var"] == pytest.approx(1.6448536269514722 * sd_week - mean_week, rel=1e-9)


def test_weekly_reference_values_from_documentation():
    """docs/measurement_and_models.tex, Example 4.1: sigma = 1 % daily, one week, V0 = 100 000."""
    s_h = 0.01 * np.sqrt(5)
    assert rk.normal_var(0.0, s_h, 0.95) == pytest.approx(0.036780, abs=1e-6)
    assert rk.normal_cvar(0.0, s_h, 0.95) == pytest.approx(0.046124, abs=1e-6)
    assert rk.normal_var(0.0, s_h, 0.99) == pytest.approx(0.052019, abs=1e-6)
    assert rk.normal_cvar(0.0, s_h, 0.99) == pytest.approx(0.059596, abs=1e-6)
    v0 = 100_000.0
    assert rk.loss_to_money(rk.normal_var(0.0, s_h, 0.95), v0) == pytest.approx(3611.2, abs=0.1)
    assert rk.loss_to_money(rk.normal_cvar(0.0, s_h, 0.95), v0) == pytest.approx(4507.6, abs=0.1)
    assert rk.loss_to_money(rk.normal_var(0.0, s_h, 0.99), v0) == pytest.approx(5068.9, abs=0.1)
    assert rk.loss_to_money(rk.normal_cvar(0.0, s_h, 0.99), v0) == pytest.approx(5785.5, abs=0.1)


def test_weighted_quantile_matches_numpy_and_handles_weights():
    rng = np.random.default_rng(3)
    x = rng.standard_normal(5000)
    qs = (0.01, 0.05, 0.5, 0.95)
    np.testing.assert_allclose(rk.weighted_quantile(x, qs), np.quantile(x, qs), rtol=1e-12)
    # equal weights reproduce the unweighted result (the weighted convention places the
    # observations at the midpoints of their probability steps, so the tails differ by half a rank)
    w = np.full(len(x), 1 / len(x))
    np.testing.assert_allclose(rk.weighted_quantile(x, qs, w), np.quantile(x, qs), atol=0.01)
    # column-wise on a matrix
    m = rng.standard_normal((2000, 3))
    np.testing.assert_allclose(rk.weighted_quantile(m, (0.1, 0.9)), np.quantile(m, (0.1, 0.9), axis=0), rtol=1e-12)
    # a sample that over-represents the left tail, corrected by its weights, recovers the standard normal
    tilted = np.concatenate([rng.standard_normal(4000), rng.standard_normal(4000) - 2.0])
    p = np.concatenate([np.full(4000, 1.0), stats.norm.pdf(tilted[4000:]) / stats.norm.pdf(tilted[4000:] + 2.0)])
    p = p / p.sum()
    for q in (0.01, 0.05, 0.5):
        assert rk.weighted_quantile(tilted, q, p)[0] == pytest.approx(stats.norm.ppf(q), abs=0.12)


def test_value_loss_maps_a_log_loss_to_a_share_of_the_position():
    assert rk.value_loss(0.0) == pytest.approx(0.0)
    assert rk.value_loss(0.30) == pytest.approx(1 - np.exp(-0.30))      # 25.9 %, not 30 %
    assert rk.value_loss(5.0) < 1.0                                     # can never reach 100 %
    np.testing.assert_allclose(rk.value_loss([0.01, 0.1]), [1 - np.exp(-0.01), 1 - np.exp(-0.1)])
    # a money loss is the value-space loss times the position
    assert rk.loss_to_money(0.30, 1_000.0) == pytest.approx(1_000.0 * rk.value_loss(0.30))


def test_value_space_var_transforms_exactly_but_cvar_needs_the_sigma_squared_term():
    """``1 - e^{-x}`` is monotone, so a VaR carries over unchanged; a CVaR is a tail *mean* and
    the concavity costs a ``sigma^2 / 2`` term -- the same correction as the Ito drift."""
    m, s, beta = 0.01, 0.30, 0.95
    r = np.random.default_rng(1).normal(m, s, 2_000_000)
    loss = 1 - np.exp(r)
    q = np.quantile(loss, beta)
    assert rk.normal_value_var(m, s, beta) == pytest.approx(q, abs=2e-3)
    assert rk.normal_value_var(m, s, beta) == pytest.approx(rk.value_loss(rk.normal_var(m, s, beta)), rel=1e-12)
    exact = rk.normal_value_cvar(m, s, beta)
    assert exact == pytest.approx(loss[loss >= q].mean(), abs=2e-3)
    naive = rk.value_loss(rk.normal_cvar(m, s, beta))       # what converting afterwards would give
    assert naive > exact                                    # concavity: the naive figure overstates
    assert 0.002 < naive - exact < 0.005
    # the correction vanishes as sigma -> 0 and grows with sigma
    assert rk.value_loss(rk.normal_cvar(m, 0.01, beta)) - rk.normal_value_cvar(m, 0.01, beta) < 1e-5
    assert rk.value_loss(rk.normal_cvar(m, 0.6, beta)) - rk.normal_value_cvar(m, 0.6, beta) > 0.005
    assert rk.normal_value_cvar(m, 3.0, beta) < 1.0         # never a total loss, whatever the volatility


def test_asset_risk_table_in_value_space():
    rng = np.random.default_rng(5)
    r = pd.DataFrame(rng.standard_normal((1500, 2)) * np.array([0.012, 0.03]),
                     index=pd.bdate_range("2019-01-01", periods=1500), columns=["A", "B"])
    log = rk.asset_risk_table(r, "daily", horizon="month", confidence=(0.95,))
    val = rk.asset_risk_table(r, "daily", horizon="month", confidence=(0.95,), space="value")
    for col in ("var95_normal_ewm", "var95_hist"):
        np.testing.assert_allclose(val[col], rk.value_loss(log[col]), rtol=1e-7)   # quantiles: exact
    for col in ("cvar95_normal_ewm", "cvar95_hist"):
        assert (val[col] < rk.value_loss(log[col])).all()                          # tail means: corrected
        assert (val[col] < 1.0).all() and (val[col] > 0).all()
    np.testing.assert_allclose(val["vol_ewm_ann"], log["vol_ewm_ann"])             # only the losses move
    with pytest.raises(ValueError):
        rk.asset_risk_table(r, "daily", space="nonsense")


def test_zero_drift_removes_the_extrapolated_mean_from_the_parametric_risk():
    """At a quarterly horizon the estimated mean dominates the VaR: a stock with a strong recent
    rally would otherwise look *less* risky than a quiet one."""
    rng = np.random.default_rng(11)
    r = pd.DataFrame({"trend": rng.standard_normal(1500) * 0.02 + 0.004,   # strong drift
                      "quiet": rng.standard_normal(1500) * 0.02},
                     index=pd.bdate_range("2019-01-01", periods=1500))
    est_ = rk.asset_risk_table(r, "daily", horizon="quarter", confidence=(0.95,))
    zero = rk.asset_risk_table(r, "daily", horizon="quarter", confidence=(0.95,), drift="zero")
    assert est_.loc["trend", "var95_normal_rolling"] < est_.loc["quiet", "var95_normal_rolling"]
    assert zero.loc["trend", "var95_normal_rolling"] == pytest.approx(
        zero.loc["quiet", "var95_normal_rolling"], rel=0.05)      # same volatility -> same risk
    assert zero.loc["trend", "var95_normal_rolling"] > est_.loc["trend", "var95_normal_rolling"]
    # the historical columns keep the drift of their window either way
    assert zero.loc["trend", "var95_hist"] == est_.loc["trend", "var95_hist"]
    with pytest.raises(ValueError):
        rk.asset_risk_table(r, "daily", drift="wishful")