import datetime as dt
import math

import numpy as np
import polars as pl

from wp39.features import aggregate_buckets, bucket_returns, session_features


def _reg_day(n=391, seed=0):
    rng = np.random.default_rng(seed)
    r = rng.normal(0, 1e-3, n)
    px = 100 * np.exp(np.cumsum(r))
    return pl.DataFrame({"session_date": [dt.date(2024, 3, 5)] * n, "m": list(range(n)), "open": np.r_[100.0, px[:-1]],
                         "high": px * 1.0005, "low": px * 0.9995, "close": px, "volume": rng.integers(100, 1000, n).astype(float),
                         "close_minute": [390] * n})


def test_bucket_aggregation_counts_and_close_print():
    reg = _reg_day()
    b5 = aggregate_buckets(reg, 5); b30 = aggregate_buckets(reg, 30)
    assert b5.height == 78 and b30.height == 13
    assert b5["close"][-1] == reg["close"][-1] and b30["close"][-1] == reg["close"][-1]
    assert abs(b5["volume"].sum() - reg["volume"].sum()) < 1e-9


def test_rv_identities_and_shares():
    reg = _reg_day()
    for f in (1, 5, 30):
        out = session_features(aggregate_buckets(reg, f), f, None)
        r = out.row(0, named=True); p = f"f{f}m_"
        assert 0 <= r[p + "rsneg_share"] <= 1 and 0 <= r[p + "rj"] <= 1 and 0 < r[p + "present_share"] <= 1.0 + 1e-9
        assert 0 <= r[p + "vol_first30"] <= 1 and 0 <= r[p + "down_vol_share"] <= 1 and 0 <= r[p + "vol_gini"] <= 1
    o1 = session_features(aggregate_buckets(reg, 1), 1, None).row(0, named=True)
    rr = bucket_returns(aggregate_buckets(reg, 1))["r"].to_numpy()
    assert abs(o1["f1m_rv"] - np.sum(rr ** 2)) < 1e-12
    assert abs(o1["f1m_drift"] - math.log(reg["close"][-1] / reg["open"][0])) < 1e-12