"""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,
    component_var,
    labeled_component_risk,
    marginal_risk_contribution,
    risk_contribution,
)


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)