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