from beequant.api.context import UserContext
from beequant.api.client import Client
import pandas as pd


def init(context: UserContext) -> None:
    context.position_size = 0.5  # Target 50% of equity per instrument
    context.log.info("Strategy initialized", source="Custom trading logic")


def handle_data(context: UserContext, datas: pd.DataFrame) -> None:
    client: Client = context.exchange[0]
    if datas is None or datas.empty:
        return

    for _, row in datas.iterrows():
        symbol = row["symbol"]
        pred = float(row.get("prediction", 0.0))
        if pred > 0:
            order_id = client.target_percent(
                symbol, context.position_size, threshold=0.03, auto_partial=True
            )
            context.log.info(f"Opened long {symbol}; order_id={order_id}", source="Custom trading logic")
        elif pred < 0:
            order_id = client.target_percent(symbol, 0, auto_partial=True)
            context.log.info(f"Closed {symbol}; order_id={order_id}", source="Custom trading logic")


def handle_tick(context: UserContext) -> None:
    pass

from beequant.api.context import UserContext
from beequant.api.client import Client
import pandas as pd


def init(context: UserContext) -> None:
    context.max_hold = 2          # Hold no more than two assets
    context.per_pct = 0.5         # Allocate 50% per asset
    context.take_profit = 0.05    # Take profit at 5%
    context.stop_loss = 0.02      # Stop loss at 2%
    context.entry_price = {}      # Record each symbol's entry price
    context.log.info("Ranked multi-asset rebalancing initialized", source="Custom trading logic")


def handle_data(context: UserContext, datas: pd.DataFrame) -> None:
    client: Client = context.exchange[0]
    if datas is None or datas.empty:
        return

    # 1) Sort the current slice by prediction score, descending
    latest = datas.sort_values("prediction", ascending=False)
    target_symbols = latest.head(context.max_hold)["symbol"].tolist()

    # 2) Close positions that are no longer in the target set
    for pos in client.get_positions():
        if pos.symbol not in target_symbols and pos.amount != 0:
            client.target_percent(pos.symbol, 0, auto_partial=True)
            context.entry_price.pop(pos.symbol, None)
            context.log.info(f"Rotated out of {pos.symbol}", source="Custom trading logic")

    # 3) Rebalance targets and record their entry prices
    for symbol in target_symbols:
        order_id = client.target_percent(
            symbol, context.per_pct, threshold=0.03, auto_partial=True
        )
        if order_id and symbol not in context.entry_price:
            context.entry_price[symbol] = client.get_price(symbol)
            context.log.info(
                f"Opened long {symbol}; entry={context.entry_price[symbol]}",
                source="Custom trading logic",
            )


def handle_tick(context: UserContext) -> None:
    """Monitor take-profit and stop-loss conditions on every tick."""
    client: Client = context.exchange[0]
    for symbol, entry in list(context.entry_price.items()):
        price = client.get_price(symbol)
        if not price or not entry:
            continue
        ret = price / entry - 1
        if ret >= context.take_profit:
            client.target_percent(symbol, 0, auto_partial=True)
            context.entry_price.pop(symbol, None)
            context.log.info(f"Took profit on {symbol}; return={ret:.2%}", source="Custom trading logic")
        elif ret <= -context.stop_loss:
            client.target_percent(symbol, 0, auto_partial=True)
            context.entry_price.pop(symbol, None)
            context.log.info(f"Stopped out of {symbol}; return={ret:.2%}", source="Custom trading logic")

from beequant.api.context import UserContext
from beequant.api.client import Client
import pandas as pd

# One threshold set and prediction column per instrument, matching upstream predictionColumn
BTC = dict(symbol="BTC_USDT", pred_col="pred_btc",
           open_long=0.07, close_long=0.02, open_short=-0.03, max_pos=0.30)
ETH = dict(symbol="ETH_USDT", pred_col="pred_eth",
           open_long=0.09, close_long=-0.02, open_short=-0.03, max_pos=0.30)


def init(context: UserContext) -> None:
    context.log.info("BTC+ETH strategy initialized", source="Custom trading logic")


def _handle_symbol(context: UserContext, client: Client, cfg: dict, pred: float) -> None:
    """Reusable per-instrument entry and exit logic; long and short branches are independent."""
    symbol, max_pos = cfg["symbol"], cfg["max_pos"]
    # Long: move to +max_pos on a strong score; close only the long when it weakens
    if pred > cfg["open_long"]:
        client.target_percent(symbol, max_pos, threshold=0.01, auto_partial=True)
    elif pred < cfg["close_long"]:
        pos = client.get_position(symbol, "long")
        if pos and abs(pos.amount) > 0:
            client.close_long(symbol)
    # Short: move to -max_pos on a low score; close only the short when it turns positive
    if pred < cfg["open_short"]:
        client.target_percent(symbol, -max_pos, threshold=0.01, auto_partial=True)
    elif pred > 0:
        pos = client.get_position(symbol, "short")
        if pos and abs(pos.amount) > 0:
            client.close_short(symbol)


def handle_data(context: UserContext, datas: pd.DataFrame) -> None:
    client: Client = context.exchange[0]
    if datas is None or datas.empty:
        context.log.warn("Upstream data is empty; skipping", source="Custom trading logic")
        return
    # This upstream input is a wide table with one row per signal timestep
    row = datas.iloc[0]
    for cfg in (BTC, ETH):
        col = cfg["pred_col"]
        # An outer merge can produce NaN for one asset; skip it to avoid an invalid trade
        if col in row and pd.notna(row[col]):
            _handle_symbol(context, client, cfg, float(row[col]))
        else:
            context.log.warn(f"Missing {col}; skipping {cfg['symbol']}", source="Custom trading logic")


def handle_tick(context: UserContext) -> None:
    pass