import numpy as np
import pandas as pd
import pytest

from strategies.mlp_volatility_stop import event_study, replay_stop, volatility
from strategies.strategy_mlp_scores import _score_to_signals
from tools.run_mlp_state_policy_experiment import window_metrics


def frame(n=100):
    return pd.DataFrame({'time': pd.date_range('2020-01-01', periods=n, freq='4h'),
                         'open': np.full(n, 100.), 'high': np.full(n, 100.),
                         'low': np.full(n, 100.), 'close': np.full(n, 100.),
                         'activation_score': np.resize([200., 0., 0., 0.], n)})


def opportunities(f, **params):
    return _score_to_signals(f.copy(), params, np.ones(len(f)), f.close.to_numpy(),
                             f.high.to_numpy(), return_opportunities=True)


def run(f, o, widths=None):
    return replay_stop(f, o, start='2020-01-01', end='2020-01-31', widths=widths)


@pytest.mark.parametrize('window', [1, 4])
@pytest.mark.parametrize('trail', [0, 5])
def test_disabled_matches_engine_exactly(window, trail):
    f = frame(300)
    rng = np.random.default_rng(71)
    f.close = rng.uniform(80, 120, len(f))
    f.high = np.maximum(f.close, 110.)
    f.activation_score = rng.normal(0, 300, len(f))
    params = {'i_exit_score_window': window, 'i_trailing_stop_threshold': trail,
              '_pine_time_start': '2020-01-01', '_pine_time_end': '2020-01-31 23:59:59'}
    expected = _score_to_signals(f.copy(), params, np.ones(len(f)), f.close.to_numpy(), f.high.to_numpy())
    actual, _ = run(f, opportunities(f, **params))
    for c in ['position', 'execute_entry', 'execute_exit']:
        np.testing.assert_array_equal(expected[c], actual[c])
    assert window_metrics(expected, '2020-01-01', '2020-01-31') == window_metrics(actual, '2020-01-01', '2020-01-31')
    if trail:
        control, _ = run(f, opportunities(f, **params), np.full(len(f), .05))
        np.testing.assert_array_equal(control.position, actual.position)


def manual_ops(f):
    o = opportunities(f, i_trailing_stop_threshold=5)
    o['entry'][:] = False
    o['exit'][:] = False
    o['entry'][0] = True
    return o


def test_fill_bar_ratchet_floor_and_score_precedence():
    f = frame(5)
    f.loc[1, 'high'] = 200  # ignored on fill bar
    f.loc[2, ['high', 'close']] = [110, 108]
    f.loc[3, ['high', 'close']] = [109, 104]
    o = manual_ops(f)
    o['exit'][3] = o['entry'][3] = True
    out, events = run(f, o, np.array([.05, .05, .1, .1, .1]))
    assert out.stop_price.iloc[0] == 95
    assert out.stop_price.iloc[1] == 95
    assert out.stop_price.iloc[2] == pytest.approx(104.5)
    assert out.stop_price.iloc[3] == pytest.approx(104.5)  # width expansion cannot loosen
    assert out.execute_exit.iloc[3] and not out.execute_entry.iloc[3]
    assert events[0]['reason'] == 'trailing_stop' and events[0]['score_exit']


def test_entry_floor_close_overshoot_and_fees():
    f = frame(3)
    f.loc[1:, 'close'] = 90
    out, events = run(f, manual_ops(f), np.full(3, .1))
    assert out.stop_price.iloc[0] == 95  # max loss threshold bound remains 5%
    assert out.execute_exit.iloc[1]
    assert out.strategy_return.iloc[1] == pytest.approx(-.11)  # return - both fees
    assert events[0]['threshold_overshoot_pct'] == pytest.approx(5.)
    assert window_metrics(out, '2020-01-01', '2020-01-31')['pnl_pct'] == pytest.approx(-11.)


def test_equality_and_terminal_entry_and_time_gate():
    f = frame(4)
    f.loc[1, 'close'] = 95
    o = manual_ops(f)
    out, _ = run(f, o)
    assert out.execute_exit.iloc[1]
    o['time_ok'][1:] = False
    held, _ = run(f, o)
    assert held.position.iloc[-1] == 1 and not held.execute_exit.any()
    o = manual_ops(f)
    o['entry'][:] = False
    o['entry'][-1] = True
    terminal, _ = run(f, o)
    assert terminal.execute_entry.iloc[-1] and terminal.strategy_return.sum() == 0


def test_volatility_seed_and_prefix_causality():
    f = frame(100)
    f.high = 102.
    f.low = 98.
    v = volatility(f)
    assert v.atr_price.iloc[:13].isna().all()
    assert v.atr_price.iloc[13] == pytest.approx(.04)
    pd.testing.assert_frame_equal(v.iloc[:50], volatility(f.iloc[:50]))
    o = manual_ops(f)
    o['entry'][0] = False
    o['entry'][30] = True
    widths = v.atr_price.to_numpy()*2
    full, _ = run(f, o, widths)
    short_o = {k: value[:50] if isinstance(value, np.ndarray) else value for k, value in o.items()}
    prefix, _ = run(f.iloc[:50], short_o, widths[:50])
    for col in ['position', 'stop_price', 'strategy_return']:
        np.testing.assert_array_equal(full[col][:50], prefix[col])
    widths[50:] = .025
    changed, _ = run(f, o, widths)
    np.testing.assert_array_equal(full.stop_price[:50], changed.stop_price[:50])
    varying = frame(100)
    varying.close = 100+np.sin(np.arange(100))*3
    varying.high = np.maximum(varying.open, varying.close)+2
    varying.low = np.minimum(varying.open, varying.close)-2
    first = volatility(varying)
    varying.loc[50:, ['open', 'high', 'low', 'close']] *= 5
    pd.testing.assert_frame_equal(first.iloc[:50], volatility(varying).iloc[:50])


def test_bad_prices_and_active_missing_width_rejected():
    f = frame()
    with pytest.raises(ValueError, match='width'):
        run(f, manual_ops(f), np.full(len(f), np.nan))
    with pytest.raises(ValueError, match='length'):
        run(f, manual_ops(f), np.ones(3))
    f.loc[5, 'close'] = np.nan
    with pytest.raises(ValueError, match='finite'):
        volatility(f)


def test_event_censoring_gaps_and_overlap_across_reasons():
    f = frame(30)
    decisions = [{'bar': t, 'entry_bar': 0, 'action': 'exit', 'reason': reason, 'score_exit': False}
                 for t, reason in [(14, 'trailing_stop'), (16, 'score_exit'), (25, 'trailing_stop')]]
    e = event_study(f, decisions, volatility(f), '2020-01-31', (6,))
    assert e.valid.tolist() == [True, True, False]
    assert e.nonoverlap.tolist() == [True, False, False]
    f.loc[20:, 'time'] += pd.Timedelta(hours=4)
    e = event_study(f, decisions, volatility(f), '2020-01-31', (6,))
    assert not e.valid.any()


def test_event_hold_advantage_uses_canonical_additive_fee():
    f = frame(30)
    f.loc[15:, 'close'] = 110.
    f.loc[20:, 'close'] = 121.
    f.high = np.maximum(f.close, f.open)
    decisions = [{'bar': 14, 'entry_bar': 0, 'action': 'exit',
                  'reason': 'trailing_stop', 'score_exit': False}]
    event = event_study(f, decisions, volatility(f), '2020-01-31', (6,)).iloc[0]
    assert event.hold_advantage_pct == pytest.approx(100*(1.1*1.095-.995))


def test_calibration_and_selection_ignore_future_data():
    from tools.run_mlp_volatility_stop_experiment import calibrate, nominate, ATR_NAMES
    f = frame(100)
    f.time = pd.date_range('2020-12-20', periods=100, freq='4h')
    f.loc[:19, 'time'] = pd.date_range('2017-11-01', periods=20, freq='4h')
    f.high = 102.
    f.low = 98.
    o = opportunities(f, i_trailing_stop_threshold=5)
    _, first = calibrate(f, o, volatility(f))
    future = f.time >= pd.Timestamp('2021-01-01')
    f.loc[future, ['open', 'high', 'low', 'close']] *= 3
    _, second = calibrate(f, o, volatility(f))
    assert first == second
    rows = [{'stage': 'validation', 'window': str(year), 'variant': name,
             'pnl_pct': 10 if name in ATR_NAMES else 5, 'calmar': 2 if name in ATR_NAMES else 1,
             'dd_pct': -5, 'entries': 10} for year in range(3)
            for name in ['baseline', *ATR_NAMES, 'fixed_atr_5', 'fixed_atr_7_5']]
    original = pd.DataFrame(rows)
    assert nominate(original)['nominated'] == 'atr_5'  # tie breaks toward narrow scale
    leaked = pd.concat([original, original.assign(stage='historical_evaluation', pnl_pct=-999)])
    assert nominate(leaked) == nominate(original)