"""Risk contribution contracts and validation tests."""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from quant_engine.risk import (
    ComponentRiskResult,
    CovarianceSnapshot,
    component_var,
    estimate_covariance_snapshot,
    labeled_component_risk,
    marginal_risk_contribution,
    risk_contribution,
)


def test_estimate_covariance_snapshot_is_complete_case_and_reproducible() -> None:
    dates = pd.date_range("2026-01-05", periods=6, freq="B")
    returns = pd.DataFrame(
        {
            "A": [0.01, 0.02, 0.03, 0.04, 0.05, 99.0],
            "B": [0.02, 0.01, np.nan, 0.03, 0.04, -99.0],
        },
        index=dates,
    )
    as_of = dates[4]

    snapshot = estimate_covariance_snapshot(
        returns,
        as_of_date=as_of,
        lookback_sessions=4,
        min_observations=3,
        data_snapshot_id="market-returns-20260109-v1",
        return_frequency="1d",
        periods_per_year=252,
    )

    expected_window = returns.loc[:as_of].tail(4)
    expected = expected_window.dropna(how="any").cov()
    pd.testing.assert_frame_equal(snapshot.covariance, expected)
    assert snapshot.snapshot_id.startswith("sample-cov-v1:")
    assert snapshot.as_of_date == as_of.date()
    assert snapshot.method == "sample"
    assert snapshot.window_start_date == expected_window.index[0].date()
    assert snapshot.window_end_date == as_of.date()
    assert snapshot.observations == 3
    assert snapshot.lookback_sessions == 4
    assert snapshot.missing_policy == "complete_case"
    assert snapshot.data_snapshot_id == "market-returns-20260109-v1"
    assert len(snapshot.input_sha256) == 64

    future_changed = returns.copy()
    future_changed.loc[dates[-1], :] = [1_000_000.0, -1_000_000.0]
    repeated = estimate_covariance_snapshot(
        future_changed,
        as_of_date=as_of,
        lookback_sessions=4,
        min_observations=3,
        data_snapshot_id="market-returns-20260109-v1",
        return_frequency="1d",
        periods_per_year=252,
    )
    assert repeated.snapshot_id == snapshot.snapshot_id
    pd.testing.assert_frame_equal(repeated.covariance, snapshot.covariance)


def test_covariance_snapshot_identity_captures_data_and_estimator_contract() -> None:
    dates = pd.date_range("2026-01-05", periods=4, freq="B")
    returns = pd.DataFrame(
        {"A": [0.01, 0.02, -0.01, 0.03], "B": [0.02, -0.01, 0.01, 0.04]},
        index=dates,
    )
    base = estimate_covariance_snapshot(
        returns,
        as_of_date=dates[-1],
        lookback_sessions=4,
        min_observations=3,
        data_snapshot_id="snapshot-a",
    )
    different_source = estimate_covariance_snapshot(
        returns,
        as_of_date=dates[-1],
        lookback_sessions=4,
        min_observations=3,
        data_snapshot_id="snapshot-b",
    )

    assert base.snapshot_id != different_source.snapshot_id
    assert base.covariance.equals(different_source.covariance)


def test_estimate_covariance_snapshot_rejects_ambiguous_or_insufficient_history() -> None:
    dates = pd.date_range("2026-01-05", periods=4, freq="B")
    returns = pd.DataFrame(
        {"A": [0.01, np.nan, 0.03, 0.04], "B": [0.02, 0.01, np.nan, 0.03]},
        index=dates,
    )

    with pytest.raises(ValueError, match="complete observations"):
        estimate_covariance_snapshot(
            returns,
            as_of_date=dates[-1],
            lookback_sessions=4,
            min_observations=3,
            data_snapshot_id="snapshot-a",
        )

    with pytest.raises(ValueError, match="strictly increasing"):
        estimate_covariance_snapshot(
            returns.iloc[::-1],
            as_of_date=dates[-1],
            lookback_sessions=4,
            min_observations=2,
            data_snapshot_id="snapshot-a",
        )


def test_covariance_snapshot_is_validated_and_immutable_by_interface() -> None:
    covariance = pd.DataFrame(
        [[0.04, 0.01], [0.01, 0.09]],
        index=["A", "B"],
        columns=["A", "B"],
    )
    snapshot = CovarianceSnapshot(
        snapshot_id="cov-20260107-v1",
        as_of_date="2026-01-07",
        covariance=covariance,
        return_frequency="1d",
        periods_per_year=252,
    )

    covariance.loc["A", "A"] = 999.0
    leaked_copy = snapshot.covariance
    leaked_copy.loc["B", "B"] = 999.0

    assert snapshot.as_of_date == pd.Timestamp("2026-01-07").date()
    assert snapshot.covariance.loc["A", "A"] == pytest.approx(0.04)
    assert snapshot.covariance.loc["B", "B"] == pytest.approx(0.09)


@pytest.mark.parametrize(
    ("kwargs", "message"),
    [
        ({"snapshot_id": ""}, "snapshot_id"),
        ({"return_frequency": ""}, "return_frequency"),
        ({"periods_per_year": 0}, "periods_per_year"),
    ],
)
def test_covariance_snapshot_rejects_incomplete_identity(
    kwargs: dict[str, object],
    message: str,
) -> None:
    values: dict[str, object] = {
        "snapshot_id": "cov-20260107-v1",
        "as_of_date": "2026-01-07",
        "covariance": pd.DataFrame([[0.04]], index=["A"], columns=["A"]),
        "return_frequency": "1d",
        "periods_per_year": 252,
    }
    values.update(kwargs)

    with pytest.raises((TypeError, ValueError), match=message):
        CovarianceSnapshot(**values)


def test_risk_contribution_sums_to_one_for_positive_portfolio_variance() -> None:
    weights = np.array([0.5, 0.5])
    covariance = np.diag([1.0, 4.0])

    result = risk_contribution(weights, covariance)

    np.testing.assert_allclose(result, [0.2, 0.8])
    assert result.sum() == pytest.approx(1.0)


def test_zero_variance_portfolio_falls_back_to_equal_contribution() -> None:
    result = risk_contribution(np.array([0.2, 0.3, 0.5]), np.zeros((3, 3)))

    np.testing.assert_allclose(result, np.full(3, 1 / 3))


def test_marginal_and_component_risk_follow_matrix_identities() -> None:
    weights = np.array([0.25, 0.75])
    covariance = np.array([[0.04, 0.01], [0.01, 0.09]])

    marginal = marginal_risk_contribution(weights, covariance)
    component = component_var(weights, covariance)

    np.testing.assert_allclose(marginal, covariance @ weights)
    np.testing.assert_allclose(component, weights * marginal)
    assert component.sum() == pytest.approx(weights @ covariance @ weights)


@pytest.mark.parametrize(
    "function",
    [risk_contribution, marginal_risk_contribution, component_var],
)
def test_risk_functions_reject_covariance_shape_mismatch(function) -> None:
    with pytest.raises(ValueError, match="does not match weights length"):
        function(np.array([0.5, 0.5]), np.eye(3))


@pytest.mark.parametrize(
    "function",
    [risk_contribution, marginal_risk_contribution, component_var],
)
def test_risk_functions_reject_empty_portfolio(function) -> None:
    with pytest.raises(ValueError, match="at least one asset"):
        function(np.array([]), np.empty((0, 0)))


def test_labeled_component_risk_aligns_covariance_and_closes_to_volatility() -> None:
    weights = pd.Series({"A": 0.25, "B": 0.75}, name="weight")
    covariance = pd.DataFrame(
        [[0.09, 0.01], [0.01, 0.04]],
        index=["B", "A"],
        columns=["B", "A"],
    )

    result = labeled_component_risk(weights, covariance)

    aligned = covariance.reindex(index=weights.index, columns=weights.index)
    expected_volatility = float(np.sqrt(weights @ aligned @ weights))
    assert isinstance(result, ComponentRiskResult)
    assert result.component.index.tolist() == ["A", "B"]
    assert result.portfolio_volatility == pytest.approx(expected_volatility)
    assert result.component.sum() == pytest.approx(expected_volatility)
    assert result.percentage.sum() == pytest.approx(1.0)


def test_component_risk_groups_actual_asset_contributions_by_label() -> None:
    weights = pd.Series({"A": 0.2, "B": 0.3, "C": 0.5})
    covariance = pd.DataFrame(np.diag([0.04, 0.09, 0.16]), index=weights.index, columns=weights.index)
    groups = pd.Series({"C": "growth", "A": "value", "B": "value"})

    result = labeled_component_risk(weights, covariance)
    grouped = result.grouped_component(groups)

    assert grouped.index.tolist() == ["growth", "value"]
    assert grouped.loc["value"] == pytest.approx(
        result.component.loc["A"] + result.component.loc["B"]
    )
    assert grouped.sum() == pytest.approx(result.portfolio_volatility)


def test_labeled_component_risk_rejects_asset_label_mismatch() -> None:
    weights = pd.Series({"A": 0.5, "B": 0.5})
    covariance = pd.DataFrame(np.eye(2), index=["A", "C"], columns=["A", "C"])

    with pytest.raises(ValueError, match="same asset labels"):
        labeled_component_risk(weights, covariance)


def test_labeled_component_risk_rejects_invalid_covariance() -> None:
    weights = pd.Series({"A": 0.5, "B": 0.5})
    asymmetric = pd.DataFrame([[1.0, 0.2], [0.1, 1.0]], index=weights.index, columns=weights.index)

    with pytest.raises(ValueError, match="symmetric"):
        labeled_component_risk(weights, asymmetric)


def test_labeled_component_risk_rejects_zero_variance_portfolio() -> None:
    weights = pd.Series({"A": 0.5, "B": 0.5})
    covariance = pd.DataFrame(np.zeros((2, 2)), index=weights.index, columns=weights.index)

    with pytest.raises(ValueError, match="positive portfolio variance"):
        labeled_component_risk(weights, covariance)