import numpy as np
import pandas as pd

from wp46.calendars.sessions import assign_sessions, xnys_early_closes
from wp46.models.baselines import HARRV, _har_X
from wp46.realized_measures import session_measures
from wp46.scoring.losses import mse, qlike, qlike_zero_safe


def test_qlike_properties():
    y = np.array([1.0, 2.0, 0.5]); h = y.copy()
    assert np.allclose(qlike(y, h), 0.0)
    assert np.all(qlike(y, h * 1.5) > 0) and np.all(qlike(y, h * 0.5) > 0)
    # zero-safe form equals qlike when y > 0 and is finite when y == 0
    assert np.allclose(qlike_zero_safe(y, h * 1.3), qlike(y, h * 1.3))
    assert np.isfinite(qlike_zero_safe(np.array([0.0]), np.array([2.0])))[0]
    assert mse(y, h)[0] == 0.0


def test_session_measures_units_and_grid():
    # 391 one-minute closes with constant 0.1% log return -> rv_1m = 390 * 0.01 = 3.9 (percent^2)
    minutes = np.arange(391)
    closes = 100.0 * np.exp(0.001 * minutes)
    m = session_measures(minutes, closes)
    assert abs(m["rv_1m"] - 3.9) < 1e-9
    assert abs(m["rv_5m"] - 78 * 0.25) < 1e-9          # 78 five-minute returns of 0.5%
    assert m["n_5m"] == 78 and m["n_bars"] == 391
    assert abs(m["rv_5m_ss"] - m["rv_5m"]) < 0.3       # offset grids have 77 returns instead of 78 (constant drift)
    assert m["rsv_pos_5m"] > 0 and m["rsv_neg_5m"] == 0


def test_xnys_early_closes_known_dates():
    ec = xnys_early_closes("2024-01-01", "2024-12-31")
    assert "2024-07-03" in ec and "2024-11-29" in ec and "2024-12-24" in ec
    assert "2024-03-08" not in ec


def test_assign_sessions_futures_window():
    ts = pd.Series(pd.to_datetime(["2024-07-01T21:30:00Z", "2024-07-01T22:30:00Z", "2024-07-02T13:30:00Z", "2024-07-02T21:30:00Z"]))
    s = assign_sessions(ts, "futures")
    # 17:30 ET Jul 1 is in the maintenance break (no session); 18:30 ET Jul 1 and 09:30 ET Jul 2 belong to session Jul 2; 17:30 ET Jul 2 is a break
    assert not s["in_window"][0] and s["in_window"][1] and s["in_window"][2] and not s["in_window"][3]
    assert str(s["session"][1].date()) == "2024-07-02" and str(s["session"][2].date()) == "2024-07-02"


def test_har_direct_targets_and_regressors():
    y = np.arange(1.0, 101.0)
    X = _har_X(y)
    assert np.isnan(X[20]).all() and X[21, 1] == 22.0 and X[21, 2] == np.mean(y[17:22]) and X[21, 3] == np.mean(y[0:22])
    f = HARRV(horizons=(1, 5)).fit(np.zeros(100), y)
    fc = HARRV(horizons=(1, 5)).forecast(f, np.zeros(100), y, 5)
    # linear series: h=1 forecast = 101, h=5 cumulative = 101+...+105 = 515
    assert abs(fc.cumulative(1) - 101) < 1e-6 and abs(fc.cumulative(5) - 515) < 1e-5