import numpy as np
import pandas as pd
import pytest
from src.portfolio import weighting as W


def test_family_of():
    assert W.family_of("PR_ETHBTC") == "PAIRS"
    assert W.family_of("SH_BTC") == "SHAPE"
    assert W.family_of("TSM01") == "TSM"
    assert W.family_of("MR01_BTC") == "FADE"
    assert W.family_of("DIP01_BTC") == "HONEST"


def test_equal_sums_to_one():
    w = W.equal(["a", "b", "c", "d"])
    assert pytest.approx(sum(w.values())) == 1.0
    assert all(abs(v - 0.25) < 1e-9 for v in w.values())


def test_manual_normalizes():
    w = W.manual(["a", "b"], {"a": 3, "b": 1})
    assert pytest.approx(w["a"]) == 0.75 and pytest.approx(w["b"]) == 0.25


def test_cap_limits_family_and_redistributes():
    ids = ["PR_ETHBTC", "PR_LTCETH", "MR01_BTC", "MR02_BTC"]
    w = W.cap(ids, caps={"PAIRS": 0.30})
    pairs_w = w["PR_ETHBTC"] + w["PR_LTCETH"]
    assert pytest.approx(pairs_w, abs=1e-9) == 0.30
    assert pytest.approx(sum(w.values())) == 1.0
    assert w["MR01_BTC"] > 0.25


def test_inverse_vol_prefers_low_vol():
    idx = pd.date_range("2024-01-01", periods=100, freq="D", tz="UTC")
    rng = np.random.default_rng(0)
    df = pd.DataFrame({"lo": rng.normal(0, 0.01, 100), "hi": rng.normal(0, 0.05, 100)}, index=idx)
    w = W.inverse_vol(["lo", "hi"], df, lookback=90)
    assert w["lo"] > w["hi"]
    assert pytest.approx(sum(w.values())) == 1.0