import datetime as dt

import polars as pl

from wp39.features import daily_block


def test_split_adjusted_overnight_return():
    day = pl.DataFrame({"session_date": [dt.date(2020, 8, 28), dt.date(2020, 8, 31)], "d_open": [500.0, 127.0],
                        "d_close": [499.0, 129.0], "d_high": [505.0, 131.0], "d_low": [495.0, 126.0], "d_volume": [1e6, 4e6]})
    acts = pl.DataFrame({"symbol": ["AAPL"], "ex_date": [dt.date(2020, 8, 31)], "split_ratio": [4.0], "dividend": [0.0]})
    d = daily_block(day, acts, "AAPL")
    assert abs(d["d_ret_cc"][1] - (129.0 / (499.0 / 4.0) - 1)) < 1e-9
    assert abs(d["d_ret_co"][1] - (127.0 / (499.0 / 4.0) - 1)) < 1e-9


def test_dividend_adjusted_return():
    day = pl.DataFrame({"session_date": [dt.date(2021, 2, 4), dt.date(2021, 2, 5)], "d_open": [100.0, 99.0],
                        "d_close": [100.0, 99.5], "d_high": [101.0, 100.0], "d_low": [99.0, 98.0], "d_volume": [1e6, 1e6]})
    acts = pl.DataFrame({"symbol": ["X"], "ex_date": [dt.date(2021, 2, 5)], "split_ratio": [1.0], "dividend": [0.5]})
    d = daily_block(day, acts, "X")
    assert abs(d["d_ret_cc"][1] - ((99.5 + 0.5) / 100.0 - 1)) < 1e-12