import numpy as np
import pandas as pd

from strategies.mlp_slow_regime_router import (
    SlowRouterConfig,
    regime_head_scores,
    slow_regime_states,
)
from tools.run_mlp_slow_regime_head_experiment import _parse_chart_time, fit_residual_heads


def _frame(n=12):
    return pd.DataFrame({
        "mvrv_zscore_cont": np.linspace(-0.8, 0.8, n),
        "rvol_norm": np.zeros(n),
        "fed_net_liq_sign": np.ones(n),
        "close": np.linspace(100.0, 130.0, n),
    })


def test_slow_router_is_causal_and_honors_minimum_dwell():
    config = SlowRouterConfig(smoothing_span=1, trend_window=1, enter_threshold=0.2,
                              exit_threshold=0.05, min_dwell_bars=3)
    frame = _frame()
    states = slow_regime_states(frame, config)
    changed_future = frame.copy()
    changed_future.loc[8:, "mvrv_zscore_cont"] = -1.0

    np.testing.assert_array_equal(states[:8], slow_regime_states(changed_future, config)[:8])
    changes = np.flatnonzero(np.diff(states) != 0) + 1
    assert all(current - previous >= config.min_dwell_bars for previous, current in zip(changes, changes[1:]))


def test_tiny_heads_reduce_to_shared_output_when_residuals_are_zero():
    trunk = np.array([[0.2, -0.3], [0.1, 0.4], [-0.2, 0.1]])
    base_head = (np.array([[0.7, -0.2]]), np.array([0.05]))
    heads = (np.zeros((3, 2)), np.zeros(3))
    states = np.array([-1, 0, 1])

    scores = regime_head_scores(trunk, states, base_head, heads)

    expected = np.tanh(trunk @ base_head[0].T + base_head[1]).ravel() * 1000.0
    np.testing.assert_allclose(scores, expected)


def test_residual_head_fit_keeps_the_shared_trunk_out_of_the_trainable_shape():
    rng = np.random.default_rng(7)
    trunk = rng.normal(size=(48, 8))
    states = np.repeat([-1, 0, 1], 16)
    y = np.tanh(trunk[:, 0] * 0.2)
    valid = np.ones(48, dtype=bool)
    times = pd.Series(pd.date_range("2025-01-01", periods=48, freq="6h"))
    base_head = (np.zeros((1, 8)), np.zeros(1))

    heads, metadata = fit_residual_heads(
        trunk, states, base_head, y, valid, times, seed=3,
        val_start=times.iloc[36], epochs=2,
    )

    assert heads[0].shape == (3, 8)
    assert heads[1].shape == (3,)
    assert metadata["train_bars"] == 36
    assert metadata["val_bars"] == 12


def test_experiment_accepts_tradingview_unix_second_timestamps():
    parsed = _parse_chart_time(pd.Series([1508068800]))

    assert parsed.iloc[0] == pd.Timestamp("2017-10-15 12:00:00")