from __future__ import annotations

import math

import pytest
import torch

from crypto_spot_bot.models import Candle
from crypto_spot_bot.time_series import _torch_head_outputs
from tools.train_torch_recurrent_forecaster import (
    OUTPUT_LAYOUT,
    RecurrentReturnModel,
    _barrier_outcome,
    _export_head_state,
    _prepare_data,
)


def _candle(index: int, *, open_: float, high: float, low: float, close: float) -> Candle:
    return Candle(index, open_, high, low, close, 100.0)


def test_barrier_target_uses_next_open_and_marks_take_profit_first() -> None:
    candles = [
        _candle(0, open_=90.0, high=101.0, low=89.0, close=100.0),
        _candle(1, open_=100.0, high=102.0, low=99.0, close=101.0),
        _candle(2, open_=101.0, high=104.0, low=100.0, close=103.0),
    ]

    net_return, event = _barrier_outcome(
        candles,
        end_index=0,
        horizon=2,
        stop_loss_percent=0.02,
        take_profit_percent=0.03,
        round_trip_cost=0.002,
    ) or (math.nan, math.nan)

    assert event == 1.0
    assert net_return == pytest.approx(math.log(1.03) - 0.002)


def test_barrier_target_resolves_same_candle_tie_as_stop_loss() -> None:
    candles = [
        _candle(0, open_=100.0, high=101.0, low=99.0, close=100.0),
        _candle(1, open_=100.0, high=104.0, low=97.0, close=101.0),
    ]

    net_return, event = _barrier_outcome(
        candles,
        end_index=0,
        horizon=1,
        stop_loss_percent=0.02,
        take_profit_percent=0.03,
        round_trip_cost=0.002,
    ) or (math.nan, math.nan)

    assert event == 0.0
    assert net_return == pytest.approx(math.log(0.98) - 0.002)


def test_multitask_head_export_matches_runtime_inference() -> None:
    torch.manual_seed(7)
    model = RecurrentReturnModel(
        architecture="gru",
        input_size=2,
        hidden_size=4,
        num_layers=1,
        dropout=0.0,
        output_size=2 * len(OUTPUT_LAYOUT),
        attention_pooling=False,
        context_norm=False,
        multitask_head=True,
        head_hidden_size=6,
    )
    model.eval()
    context = torch.tensor([[0.2, -0.1, 0.4, 0.3]], dtype=torch.float32)
    with torch.no_grad():
        shared = model.head_activation(model.head_hidden(context))
        returns = model.return_head(shared).view(1, 2, 4)
        events = model.event_head(shared).view(1, 2, 1)
        expected = torch.cat((returns, events), dim=2).reshape(-1).tolist()

    entry = {"multitask_head": True, **_export_head_state(model)}
    actual = _torch_head_outputs(context[0].tolist(), entry, hidden_size=4)

    assert actual == pytest.approx(expected, abs=2e-6)


def test_pooled_training_populates_symbol_identity_feature() -> None:
    candles = [
        _candle(
            index,
            open_=100.0 + index * 0.01,
            high=100.2 + index * 0.01,
            low=99.8 + index * 0.01,
            close=100.0 + index * 0.01,
        )
        for index in range(180)
    ]

    prepared = _prepare_data(
        symbol="BTCUSDT",
        candles=candles,
        feature_names=["return_1", "symbol_is_BTCUSDT", "symbol_is_ETHUSDT"],
        lookback=8,
        target_horizons=[3],
        decision_horizon=3,
        round_trip_cost=0.002,
        stop_loss_percent=0.04,
        take_profit_percent=0.035,
        market_candles={"BTCUSDT": candles},
        trend_candles=candles,
        validation_window=24,
        holdout_window=32,
        clip=8.0,
        device=torch.device("cpu"),
    )

    assert prepared is not None
    assert torch.all(prepared.train_x[:, :, 1] == 1.0)
    assert torch.all(prepared.train_x[:, :, 2] == 0.0)