import numpy as np
import pandas as pd
from src.live.basket_trend_worker import BasketTrendWorker


def _ramp_df(n=300, slope=1.0):
    c = np.linspace(100, 100 + slope * n, n)
    ts = (pd.date_range("2024-01-01", periods=n, freq="4h", tz="UTC").astype("int64") // 10**6)
    return pd.DataFrame({"timestamp": ts, "open": c, "high": c, "low": c, "close": c, "volume": 1.0})


def test_basket_goes_long_in_uptrend(tmp_path):
    w = BasketTrendWorker(universe=["AAA", "BBB"], tf="4h", capital=1000.0, data_dir=tmp_path)
    data = {"AAA": _ramp_df(slope=1.0), "BBB": _ramp_df(slope=1.0)}
    w.tick(data)
    assert w.positions["AAA"] == 1.0 and w.positions["BBB"] == 1.0


def test_basket_flat_in_downtrend(tmp_path):
    w = BasketTrendWorker(universe=["AAA"], tf="4h", capital=1000.0, data_dir=tmp_path)
    data = {"AAA": _ramp_df(slope=-1.0)}
    w.tick(data)
    assert w.positions["AAA"] == 0.0


def test_basket_persists_and_resumes(tmp_path):
    w = BasketTrendWorker(universe=["AAA"], tf="4h", capital=1000.0, data_dir=tmp_path)
    w.tick({"AAA": _ramp_df(slope=1.0)})
    w2 = BasketTrendWorker(universe=["AAA"], tf="4h", capital=1000.0, data_dir=tmp_path)
    assert w2.positions["AAA"] == 1.0          # stato ripreso da status.json


def _ramp_df_now(n=300, slope=1.0, forming_close=None):
    """Serie 4h che termina ADESSO: l'ultima riga e' la candela IN FORMAZIONE."""
    import time
    c = np.linspace(100, 100 + slope * n, n)
    if forming_close is not None:
        c = c.copy()
        c[-1] = forming_close
    bar = 4 * 3_600_000
    now_ms = int(time.time() * 1000)
    ts = now_ms - bar * np.arange(n - 1, -1, -1)
    return pd.DataFrame({"timestamp": ts, "open": c, "high": c, "low": c, "close": c, "volume": 1.0})


def test_basket_ignores_forming_bar(tmp_path):
    # downtrend su barre COMPLETE (target flat); la candela in formazione e' un
    # spike estremo che flipperebbe EMA20>EMA100 se contata -> va ignorata.
    w = BasketTrendWorker(universe=["AAA"], tf="4h", capital=1000.0, data_dir=tmp_path)
    w.tick({"AAA": _ramp_df_now(slope=-0.3, forming_close=1e6)})
    assert w.positions["AAA"] == 0.0


def test_basket_fee_includes_leverage(tmp_path):
    # flip 0->1: fee = cap * POS * LEV * fee_rt/2 (reference honest_improve2:150)
    w = BasketTrendWorker(universe=["AAA"], tf="4h", capital=1000.0, data_dir=tmp_path)
    w.tick({"AAA": _ramp_df(slope=1.0)})
    expected = 1000.0 - 1000.0 * w.position_size * w.leverage * (w.fee_rt / 2)
    assert np.isclose(w.capital, expected)