import numpy as np
import pandas as pd
import pytest

from strategies.mlp_state_policy import StatePolicy, replay_state_policy
from strategies.strategy_mlp_scores import _score_to_signals


def replay(entry, exits, policy=StatePolicy(), risk=None, close=None, trail=0):
    n = len(entry)
    return replay_state_policy(np.array(entry, bool), np.array(exits, bool),
        np.array(exits, bool), np.zeros(n), 0, np.ones(n) * 100 if close is None else np.array(close),
        np.ones(n) * 100, trail, policy=policy, risk=risk)


def test_macro_blocks_trigger_and_exactly_n_following_bars():
    result = replay([1]*5, [0]*5, StatePolicy(macro_memory=2), [1, 0, 0, 0, 0])
    assert result['execute_entry'].tolist() == [False, False, False, True, False]
    assert result['policy_action'].tolist() == ['stay_flat']*3 + ['enter_long', 'hold_long']


def test_exit_cooldown_and_no_same_bar_reentry():
    result = replay([1]*6, [0, 1, 0, 0, 0, 0], StatePolicy(exit_cooldown=2))
    assert result['position'].tolist() == [1, 0, 0, 0, 1, 1]


def test_hold_latches_exit_until_minimum_age():
    result = replay([1, 0, 0, 0, 0], [0, 1, 0, 0, 0], StatePolicy(min_hold=3))
    assert result['position'].tolist() == [1, 1, 1, 0, 0]


def test_trailing_stop_overrides_hold():
    result = replay([1, 0, 0], [0, 0, 0], StatePolicy(min_hold=12), close=[100, 80, 80], trail=.1)
    assert result['execute_exit'].tolist() == [False, True, False]


@pytest.mark.parametrize('window', [1, 3])
@pytest.mark.parametrize('trail', [0, 10])
def test_null_policy_exactly_matches_existing_replay(window, trail):
    rng = np.random.default_rng(22)
    frame = pd.DataFrame({'activation_score': rng.normal(0, 200, 400),
                          'time': pd.date_range('2020-01-01', periods=400),
                          'close': rng.uniform(80, 100, 400), 'high': np.ones(400)*101})
    params = {'i_exit_score_window': window, 'i_trailing_stop_threshold': trail,
              '_pine_time_start': '2020-02-01', '_pine_time_end': '2020-11-01'}
    args = (params, np.ones(400), frame.close.to_numpy(), frame.high.to_numpy())
    base = _score_to_signals(frame.copy(), *args)
    experimental = _score_to_signals(frame.copy(), *args, state_policy=StatePolicy())
    for column in ('position', 'execute_entry', 'execute_exit'):
        np.testing.assert_array_equal(base[column], experimental[column])


def test_policy_is_prefix_causal():
    kwargs = dict(policy=StatePolicy(macro_memory=2, exit_cooldown=2, min_hold=3))
    full = replay([1]*8, [0, 1, 0, 0, 0, 1, 0, 0], risk=[0, 0, 1, 0, 0, 0, 0, 0], **kwargs)
    prefix = replay([1]*4, [0, 1, 0, 0], risk=[0, 0, 1, 0], **kwargs)
    for column in full:
        np.testing.assert_array_equal(full[column][:4], prefix[column])


def test_negative_memory_rejected():
    with pytest.raises(ValueError):
        StatePolicy(macro_memory=-1)


def test_repeated_macro_risk_refreshes_block():
    result = replay([1]*6, [0]*6, StatePolicy(macro_memory=2), [1, 0, 1, 0, 0, 0])
    assert np.flatnonzero(result['execute_entry']).tolist() == [5]


def test_deferred_exit_respects_time_window():
    result = replay_state_policy(np.array([1, 0, 0, 0], bool), np.array([0, 1, 0, 0], bool),
        np.zeros(4, bool), np.zeros(4), 0, np.ones(4)*100, np.ones(4)*100, 0,
        policy=StatePolicy(min_hold=2), time_ok=[True, True, False, False])
    assert result['position'].tolist() == [1, 1, 1, 1]


def test_policy_changes_real_score_signal_chain():
    frame = pd.DataFrame({'activation_score': [200, 0, 0, 0, 0, 0, 0],
                          'close': [100.]*7, 'high': [100.]*7})
    args = ({}, np.ones(7), frame.close.to_numpy(), frame.high.to_numpy())
    baseline = _score_to_signals(frame.copy(), *args)
    held = _score_to_signals(frame.copy(), *args, state_policy=StatePolicy(min_hold=4))
    assert np.flatnonzero(baseline.execute_exit).tolist() == [2]
    assert np.flatnonzero(held.execute_exit).tolist() == [5]