import numpy as np
import pandas as pd
import pytest

from quantcore.models.black_litterman import (
    BlackLittermanModel, View, idzorek_omega, idzorek_omega_numeric, view_from_sentiment,
)


@pytest.fixture
def model():
    symbols = ["A", "B", "C", "D"]
    vols = np.array([0.20, 0.25, 0.15, 0.30]) / np.sqrt(52)
    corr = np.array([[1, .5, .3, .2], [.5, 1, .4, .1], [.3, .4, 1, .3], [.2, .1, .3, 1]])
    cov = pd.DataFrame(np.outer(vols, vols) * corr, index=symbols, columns=symbols)
    w = pd.Series([0.4, 0.3, 0.2, 0.1], index=symbols)
    return BlackLittermanModel(cov, w, market_return=0.07 / 52, risk_free=0.01 / 52, tau=0.05, periods_per_year=52)


def test_equilibrium_is_consistent(model):
    # the market portfolio earns exactly the market excess return under the implied returns
    assert model.w @ model.pi == pytest.approx(model.market_return - model.risk_free)
    assert model.delta > 0
    res = model.posterior([])
    pd.testing.assert_series_equal(res.posterior, res.prior)
    # market weights are recovered as the unconstrained MV weights of the prior (tau-adjusted cov ignored here)
    w_mv = np.linalg.solve(model.delta * model.sigma, model.pi)
    np.testing.assert_allclose(w_mv, model.w)


def test_idzorek_closed_form_equals_numeric(model):
    p = np.array([0.0, 1.0, 0.0, 0.0])
    q = (0.12 - 0.01) / 52
    for c in (0.1, 0.35, 0.5, 0.8, 0.99):
        omega_cf = idzorek_omega(p, model.tau * model.sigma, c)
        omega_num = idzorek_omega_numeric(p, q, model.pi, model.sigma, model.tau, model.delta, c)
        assert omega_cf == pytest.approx(omega_num, rel=1e-4)


def test_confidence_moves_weights_by_fraction(model):
    view = View({"B": 1.0}, q=0.15, confidence=0.4)
    res = model.posterior([view])
    w_post = np.linalg.solve(model.delta * model.sigma, res.posterior.values)
    full = model.posterior([View({"B": 1.0}, q=0.15, confidence=1.0)])
    w_full = np.linalg.solve(model.delta * model.sigma, full.posterior.values)
    tilt_partial = w_post - model.w
    tilt_full = w_full - model.w
    np.testing.assert_allclose(tilt_partial, 0.4 * tilt_full, atol=1e-10)


def test_full_confidence_reproduces_view(model):
    res = model.posterior([View({"B": 1.0}, q=0.15, confidence=1.0)])
    assert res.posterior_ann["B"] == pytest.approx(0.15, abs=1e-9)


def test_positive_view_raises_asset_and_relative_view(model):
    res = model.posterior([View({"A": 1.0}, q=0.20, confidence=0.6)])
    assert res.posterior_ann["A"] > res.prior_ann["A"]
    rel = model.posterior([View({"C": 1.0, "D": -1.0}, q=0.05, confidence=0.5, label="C beats D")])
    assert (rel.posterior_ann["C"] - rel.posterior_ann["D"]) > (rel.prior_ann["C"] - rel.prior_ann["D"])
    assert rel.views[0]["label"] == "C beats D"
    assert rel.posterior_cov.shape == (4, 4)


def test_view_from_sentiment(model):
    v = view_from_sentiment(model, "C", sentiment=1.0, confidence=0.7, strength=0.25)
    assert v.q == pytest.approx(model.prior_ann["C"] + 0.25 * 0.15, rel=1e-6)
    neg = view_from_sentiment(model, "C", sentiment=-1.0, confidence=0.7)
    assert neg.q < model.prior_ann["C"]
    with pytest.raises(ValueError):
        View({"A": 1.0}, q=0.1, confidence=0.0)


def test_documentation_example_numbers(model):
    """docs/measurement_and_models.tex, Example 7.2 (the fixture is that four-asset model)."""
    assert model.market_var * 52 == pytest.approx(0.02502, abs=1e-5)
    assert model.delta == pytest.approx(2.3986, abs=1e-4)
    np.testing.assert_allclose(model.prior_ann.values, [0.0736, 0.0880, 0.0435, 0.0550], atol=5e-5)
    assert model.w @ model.prior_ann.values == pytest.approx(0.07, abs=1e-9)
    res = model.posterior([View({"B": 1.0}, q=0.15, confidence=0.4)])
    p = np.array([0.0, 1.0, 0.0, 0.0])
    base = p @ (model.tau * model.sigma) @ p
    assert base == pytest.approx(6.0096e-5, rel=1e-4)
    assert res.omega[0] == pytest.approx(1.5 * base, rel=1e-12)
    assert res.omega[0] == pytest.approx(9.0144e-5, rel=1e-4)
    np.testing.assert_allclose(res.posterior_ann.values, [0.0835, 0.1128, 0.0494, 0.0580], atol=5e-5)
    w_post = np.linalg.solve(model.delta * model.sigma, res.posterior.values)
    full = model.posterior([View({"B": 1.0}, q=0.15, confidence=1.0)])
    w_100 = np.linalg.solve(model.delta * model.sigma, full.posterior.values)
    assert w_100[1] == pytest.approx(0.7139, abs=5e-5) and w_post[1] == pytest.approx(0.4656, abs=5e-5)
    np.testing.assert_allclose(w_post[[0, 2, 3]], model.w[[0, 2, 3]], atol=1e-10)