import numpy as np
import pytest

from chronos_x.data.synthetic import synthetic_bars
from chronos_x.data.features import build_grid, N_DAY
from chronos_x.data.quality import qa_report


@pytest.fixture(scope="module")
def grid():
    df, actions = synthetic_bars(start="2024-01-02", end="2024-02-15", seed=1, split=("2024-01-22", 4.0))
    arrays, meta = build_grid(df, actions, "2024-01-02", "2024-02-15", "SYN", "stock")
    return df, actions, arrays, meta


def test_qa_clean(grid):
    df, *_ = grid
    q = qa_report(df)
    assert q["ohlc_inconsistent"] == 0 and q["duplicates"] == 0 and q["off_interval"] == 0


def test_extended_hours_excluded_and_coverage(grid):
    df, actions, arrays, meta = grid
    assert meta.counts["rows_extended_hours"] > 0
    assert 0.95 < meta.counts["coverage"] <= 1.0
    assert arrays["mask"].shape[1] == 391


def test_split_adjusted_prices_continuous(grid):
    df, actions, arrays, meta = grid
    dates = meta.session_dates
    i = dates.index("2024-01-22")
    base, mask = arrays["base"], arrays["mask"]
    # last close of day before ex-date vs first open on ex-date: after explicit split adjustment the jump is small
    k1 = np.max(np.flatnonzero(mask[i - 1])); k0 = np.argmax(mask[i] > 0)
    jump = base[i, k0, 0] - base[i - 1, k1, 3]
    assert abs(jump) < 0.05
    assert meta.split_factors[i - 1] == 4.0 and meta.split_factors[i] == 1.0


def test_buckets_consistent_with_minutes(grid):
    df, actions, arrays, meta = grid
    base, mask, b5, m5, b30, m30 = (arrays[k] for k in ("base", "mask", "b5", "m5", "b30", "m30"))
    s = 5
    for j in range(78):
        sl = slice(5 * j, 5 * j + 5)
        pres = mask[s, sl] > 0
        if not pres.any():
            assert m5[s, j] == 0
            continue
        assert m5[s, j] == 1
        assert np.isclose(b5[s, j, 3], base[s, sl, 3][pres][-1])
        assert np.isclose(b5[s, j, 0], base[s, sl, 0][pres][0])
        assert np.isclose(b5[s, j, 1], base[s, sl, 1][pres].max())
        assert np.isclose(b5[s, j, 4], np.log1p(np.expm1(base[s, sl, 4][pres].astype(np.float64)).sum()), atol=1e-4)
    # closing print (slot 390) is never inside a bucket
    assert b30.shape[1] == 13 and b5.shape[1] == 78


def test_seasonal_profile_is_past_only(grid):
    df, actions, arrays, meta = grid
    seas, base, mask = arrays["seas"], arrays["base"], arrays["mask"]
    assert np.isnan(seas[0]).all()
    s, k = 10, 100
    past = [base[t, k, 4] for t in range(max(0, s - 20), s) if mask[t, k]]
    assert np.isclose(seas[s, k], np.mean(past), atol=1e-5)