import numpy as np
import pandas as pd

from tools import train_mlp


def test_purge_removes_exact_labels_crossing_validation():
    times = pd.Series(pd.date_range('2020-01-01', periods=40, freq='4h'))
    valid = np.ones(40, dtype=bool)
    mask = train_mlp.purged_training_mask(valid, times, times.iloc[25], 10)
    assert np.flatnonzero(mask).tolist() == list(range(15))
    assert (times.shift(-10)[mask] < times.iloc[25]).all()


def test_training_loader_epoch_and_full_final_day(tmp_path, monkeypatch):
    times = pd.date_range('2020-01-01', periods=13, freq='4h')
    path = tmp_path / 'input.csv'
    pd.DataFrame({'time': [int(t.timestamp()) for t in times], 'close': 1}).to_csv(path, index=False)
    monkeypatch.setattr(train_mlp.config, 'TRAIN_START', '2020-01-01')
    monkeypatch.setattr(train_mlp.config, 'TRAIN_END', '2020-01-02')
    got = train_mlp.load_data(path)
    assert len(got) == 12
    assert got.time.iloc[-1] == pd.Timestamp('2020-01-02 20:00')


def test_fitness_preserves_preboundary_history_and_full_end_day(monkeypatch):
    times = pd.Series(pd.date_range('2020-01-01', periods=24, freq='4h'))
    frame = pd.DataFrame({'time': times, 'close': 100., 'high': 101.})
    layers = [(np.zeros((1, 10)), np.zeros(1))]
    theta = np.r_[train_mlp.flatten_layers(layers), [0., 0., 0., 0.]]
    context = train_mlp.make_fitness_context(frame, np.zeros((24, 10)), [10, 1], theta[:11],
        [(None, '2020-01-02', '2020-01-03')], 1, 1, 0.)
    monkeypatch.setattr(train_mlp, '_score_to_signals', lambda d, *args: d)
    seen = {}
    def metrics(sl, **kwargs):
        seen['first'] = sl.time.iloc[0]
        seen['last'] = sl.time.iloc[-1]
        seen.update(kwargs)
        return {'Calmar Ratio': 1., 'Total Trades': 1}
    monkeypatch.setattr(train_mlp, 'calculate_metrics', metrics)
    train_mlp.evaluate_theta(theta, context)
    assert seen['first'] == times.iloc[0]
    assert seen['last'] == pd.Timestamp('2020-01-03 20:00')
    assert seen['score_start'] == '2020-01-02'