from typing import Dict

import numpy as np
import talib as ta
from crypto_data_downloader.utils import load_pkl
from trading_models.utils import D_TYPE

from fast_trading_simulator.sim import map_trades, simulate
from fast_trading_simulator.utils import pymoo_minimize

path = "./futures_data_2025-08-01_2025-11-20.pkl"
data: D_TYPE = load_pkl(path, gz=True)
ref_sym = "BTCUSDT"
price_idx = 1
T = len(data[ref_sym])
data = {k: v for k, v in data.items() if len(v) == T}
market = np.array(list(data.values()))
tot_fee = 1e-3
liq_fee = np.full(len(data), 0.02)


def loss_fn(p: Dict, plot=False):
    MA = [ta.SMA, ta.KAMA][int(p["ma_idx"])]
    all_act = []
    for sym, v in data.items():
        price = v[:, price_idx]
        obs = price / MA(price, int(p["ma_period"])) - 1
        pos = np.where(obs < -p["delta"], 1, 0)
        pos = np.where(obs > p["delta"], -1, pos)
        lev = np.full(T, int(p["lev"]))
        timeout = np.full(T, p["timeout"])
        take_profit = np.full(T, p["take_profit"])
        stop_loss = np.full(T, p["stop_loss"])
        act = np.array([pos, lev, timeout, take_profit, stop_loss]).T
        all_act.append(act)
    action = np.array(all_act)
    trades = simulate(market, action, tot_fee, liq_fee, use_ratio=0.2, alloc_ratio=0.01)
    res = map_trades(trades, plot=plot)
    final_worth = res["worth"][-1]
    return -final_worth


if 0:
    conf = {
        "ma_idx": [0, 1, 1],
        "ma_period": [10, 100, 1],
        "delta": [0.01, 0.1, 0.005],
        "lev": [1, 3, 1],
        "timeout": [10, 100, 1],
        "take_profit": [0.01, 0.1, 0.01],
        "stop_loss": [-0.5, -0.1, 0.01],
    }
    pymoo_minimize(loss_fn, conf)
else:
    p_best = {
        "ma_idx": 0,
        "ma_period": 26,
        "delta": 0.1,
        "lev": 3,
        "timeout": 74,
        "take_profit": 0.01,
        "stop_loss": -0.47,
    }
    loss_fn(p_best, plot=True)
