# tests/test_labels.py
import numpy as np
import pandas as pd
from src.data.features import generate_labels_dual_direction
from config import BASE_HORIZON, MIN_HORIZON, MAX_HORIZON


def make_df_long():
    """DataFrame足够长以通过valid_start=20"""
    n = 35
    closes = list(range(100, 100 + n))
    highs = [c + 1 for c in closes]
    lows = [c - 1 for c in closes]
    df = pd.DataFrame({'Close': closes, 'High': highs, 'Low': lows})
    atr = pd.Series([1.0] * n)
    return df, atr


def make_df_short():
    """DataFrame for short scenario"""
    n = 35
    closes = list(range(100, 100 - n, -1))
    highs = [c + 1 for c in closes]
    lows = [c - 1 for c in closes]
    df = pd.DataFrame({'Close': closes, 'High': highs, 'Low': lows})
    atr = pd.Series([1.0] * n)
    return df, atr


def test_long_hit_tp_before_sl():
    df, atr = make_df_long()
    labels_long, labels_short = generate_labels_dual_direction(df, atr, rr_ratio=2.0)
    assert labels_long[20] == 1.0


def test_short_hit_tp_before_sl():
    df, atr = make_df_short()
    labels_long, labels_short = generate_labels_dual_direction(df, atr, rr_ratio=2.0)
    assert labels_short[20] == 1.0
    assert labels_long[20] == 0.0


def test_long_hit_sl_before_tp():
    """LONG: цена падает, SL hit раньше TP"""
    n = 35
    # Price dropping: closes go 120, 119, 118...
    closes = np.array(list(range(120, 120 - n, -1)))
    highs = closes + 0.3  # Not reaching TP
    lows = closes - 0.5   # Goes below SL quickly
    df = pd.DataFrame({'Close': closes, 'High': highs, 'Low': lows})
    atr = pd.Series(np.ones(n))

    labels_long, _ = generate_labels_dual_direction(df, atr, rr_ratio=2.0)
    # LONG at i=20: entry=100, TP=102, SL=99
    # Future lows go below 99, TP at 102 never reached
    # SL hit first -> label = 0
    assert labels_long[20] == 0.0


def test_short_hit_sl_before_tp():
    """SHORT: цена растёт, SL hit раньше TP"""
    n = 35
    # Price rising: closes go 100, 101, 102...
    closes = np.array(list(range(100, 100 + n)))
    highs = closes + 0.5  # Goes above SL quickly
    lows = closes - 0.3    # Never reaches TP
    df = pd.DataFrame({'Close': closes, 'High': highs, 'Low': lows})
    atr = pd.Series(np.ones(n))

    _, labels_short = generate_labels_dual_direction(df, atr, rr_ratio=2.0)
    # SHORT at i=20: entry=120, TP=118, SL=121
    # Future highs exceed 121, TP at 118 never reached
    # SL hit first -> label = 0
    assert labels_short[20] == 0.0


def test_rr_ratio_3():
    """Проверка RR=3"""
    n = 35
    closes = np.array(list(range(100, 100 + n)))
    highs = closes + 1
    lows = closes - 1
    df = pd.DataFrame({'Close': closes, 'High': highs, 'Low': lows})
    atr = pd.Series(np.ones(n))

    labels_long, _ = generate_labels_dual_direction(df, atr, rr_ratio=3.0)
    # LONG at i=20: entry=120, TP=123, SL=119
    # Future highs: first >= 123 is at index 3 (low=122, high=123)
    assert labels_long[20] == 1.0