"""Triple Barrier Method for generating training labels.

Based on Marcos Lopez de Prado's "Advances in Financial Machine Learning".
Generates labels for entry signal classification and SL/TP prediction.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Tuple, List

import numpy as np
import pandas as pd
from loguru import logger

from ai.features import calculate_atr, true_range


@dataclass
class BarrierResult:
    """Result of triple barrier analysis for a single bar."""
    entry_idx: int
    entry_price: float
    entry_time: int
    side: str  # "LONG" or "SHORT"
    outcome: int  # 1 = TP hit, -1 = SL hit, 0 = timeout
    exit_price: float
    exit_idx: int
    bars_to_exit: int
    pnl_pct: float
    max_favorable_excursion: float  # MFE - best price reached
    max_adverse_excursion: float  # MAE - worst price reached
    atr_at_entry: float
    suggested_sl_distance: float  # in price units
    suggested_tp_distance: float  # in price units


@dataclass
class BarrierConfig:
    """Configuration for triple barrier labeling."""
    tp_atr_mult: float = 1.5  # TP = entry ± ATR * mult (breakeven WR = 40%)
    sl_atr_mult: float = 1.0  # SL = entry ∓ ATR * mult
    max_holding_bars: int = 48  # Maximum bars to hold position
    atr_period: int = 14
    min_atr_pct: float = 0.3  # Minimum ATR% to consider valid
    skip_bars_after_entry: int = 1  # Skip N bars after entry (execution delay)
    trend_period: int = 50  # SMA period for trend filter (0 = disabled)
    side_selection: str = "trend"  # "trend" | "momentum" — how to pick side per bar


def compute_triple_barrier_labels(
    df: pd.DataFrame,
    config: Optional[BarrierConfig] = None,
    sides: Optional[List[str]] = None,
) -> pd.DataFrame:
    """Compute triple barrier labels for all bars in DataFrame.
    
    Args:
        df: OHLCV DataFrame with columns [timestamp, Open, High, Low, Close, Volume]
        config: Barrier configuration parameters
        sides: List of sides to analyze ["LONG", "SHORT"]. If None, both are used.
        
    Returns:
        DataFrame with labels for training:
        - entry_idx, entry_price, entry_time
        - side (LONG/SHORT)
        - outcome (1=TP, -1=SL, 0=timeout)
        - exit_price, exit_idx, bars_to_exit
        - pnl_pct, mfe, mae
        - atr_at_entry, sl_distance, tp_distance
    """
    if config is None:
        config = BarrierConfig()
    
    if sides is None:
        sides = ["LONG", "SHORT"]
    
    if df.empty or len(df) < config.atr_period + config.max_holding_bars + 10:
        logger.warning(f"Недостаточно данных для labeling: {len(df)} баров")
        return pd.DataFrame()
    
    df = df.sort_values("timestamp").reset_index(drop=True)
    
    df["ATR"] = calculate_atr(df, config.atr_period)
    df["tr"] = true_range(df)
    
    # Trend filter: compute SMA for direction determination
    if config.trend_period > 0 and len(df) > config.trend_period + 10:
        df["SMA_trend"] = df["Close"].rolling(config.trend_period).mean()
    else:
        df["SMA_trend"] = float("nan")
    
    results: List[BarrierResult] = []
    
    start_idx = config.atr_period + 5
    end_idx = len(df) - config.max_holding_bars - config.skip_bars_after_entry
    
    for idx in range(start_idx, end_idx):
        atr_val = df["ATR"].iloc[idx]
        close_price = df["Close"].iloc[idx]
        
        if pd.isna(atr_val) or atr_val <= 0:
            continue
        
        atr_pct = (atr_val / close_price) * 100
        if atr_pct < config.min_atr_pct:
            continue
        
        # 🚫 КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: ровно ОДИН лейбл на бар.
        # Раньше генерация разрешала ОБЕ стороны (LONG+SHORT) для баров около SMA.
        # Модель получала идентичные фичи (отличается 1 бит side из ~400) с
        # противоположными таргетами → не могла обучиться → AUC=0.5.
        #
        # Теперь: сторона определяется строго по тренду. Если тренд неопределён
        # (или SMA недоступен), бар пропускается — никаких дуальных меток.
        side = _select_side_for_bar(df, idx, config, sides)
        if side is None:
            continue
        
        result = _analyze_single_barrier(
            df=df,
            entry_idx=idx,
            side=side,
            atr_val=atr_val,
            config=config,
        )
        if result is not None:
            results.append(result)
    
    if not results:
        logger.warning("Не сгенерировано ни одной метки")
        return pd.DataFrame()
    
    labels_df = pd.DataFrame([
        {
            "entry_idx": r.entry_idx,
            "entry_price": r.entry_price,
            "entry_time": r.entry_time,
            "side": r.side,
            "outcome": r.outcome,
            "exit_price": r.exit_price,
            "exit_idx": r.exit_idx,
            "bars_to_exit": r.bars_to_exit,
            "pnl_pct": r.pnl_pct,
            "mfe": r.max_favorable_excursion,
            "mae": r.max_adverse_excursion,
            "atr_at_entry": r.atr_at_entry,
            "sl_distance": r.suggested_sl_distance,
            "tp_distance": r.suggested_tp_distance,
        }
        for r in results
    ])
    
    tp_count = (labels_df["outcome"] == 1).sum()
    sl_count = (labels_df["outcome"] == -1).sum()
    timeout_count = (labels_df["outcome"] == 0).sum()
    
    logger.info(
        f"Triple Barrier labeling: {len(labels_df)} меток "
        f"(TP={tp_count}, SL={sl_count}, Timeout={timeout_count})"
    )
    
    return labels_df


def _select_side_for_bar(
    df: pd.DataFrame,
    idx: int,
    config: BarrierConfig,
    sides: List[str],
) -> Optional[str]:
    """Select ONE trading side for a given bar based on trend/momentum.
    
    CRITICAL: Returns exactly one side or None. NEVER returns multiple sides
    for the same bar — dual labels with nearly identical features but opposite
    targets is the #1 cause of AUC≈0.5 convergence failure.
    
    Args:
        df: OHLCV DataFrame with SMA_trend column
        idx: Bar index
        config: Barrier configuration
        sides: Allowed sides from caller ["LONG", "SHORT"]
        
    Returns:
        Selected side ("LONG" or "SHORT"), or None if no clear direction
    """
    if len(sides) == 1:
        return sides[0]
    
    if config.side_selection == "momentum":
        # Momentum-based: compare close to open of current bar
        close = df["Close"].iloc[idx]
        open_ = df["Open"].iloc[idx]
        return "LONG" if close >= open_ else "SHORT"
    
    # Default: trend-based selection
    sma_val = df["SMA_trend"].iloc[idx] if "SMA_trend" in df.columns else float("nan")
    
    if pd.notna(sma_val) and sma_val > 0:
        close_price = df["Close"].iloc[idx]
        trend_up = close_price > sma_val
        if trend_up:
            return "LONG" if "LONG" in sides else None
        else:
            return "SHORT" if "SHORT" in sides else None
    else:
        # No trend info available: use momentum as fallback
        close = df["Close"].iloc[idx]
        open_ = df["Open"].iloc[idx]
        if close >= open_ and "LONG" in sides:
            return "LONG"
        elif close < open_ and "SHORT" in sides:
            return "SHORT"
        return None


def _analyze_single_barrier(
    df: pd.DataFrame,
    entry_idx: int,
    side: str,
    atr_val: float,
    config: BarrierConfig,
) -> Optional[BarrierResult]:
    """Analyze triple barrier for a single entry point.
    
    Args:
        df: Full OHLCV DataFrame
        entry_idx: Index of entry bar
        side: "LONG" or "SHORT"
        atr_val: ATR value at entry
        config: Barrier configuration
        
    Returns:
        BarrierResult or None if analysis failed
    """
    entry_price = df["Close"].iloc[entry_idx]
    entry_time = int(df["timestamp"].iloc[entry_idx])
    
    sl_distance = atr_val * config.sl_atr_mult
    tp_distance = atr_val * config.tp_atr_mult
    
    if side == "LONG":
        sl_price = entry_price - sl_distance
        tp_price = entry_price + tp_distance
    else:
        sl_price = entry_price + sl_distance
        tp_price = entry_price - tp_distance
    
    scan_start = entry_idx + config.skip_bars_after_entry
    scan_end = min(entry_idx + config.max_holding_bars + 1, len(df))
    
    outcome = 0
    exit_price = entry_price
    exit_idx = scan_end - 1
    bars_to_exit = config.max_holding_bars
    
    mfe = 0.0
    mae = 0.0
    
    for i in range(scan_start, scan_end):
        high = df["High"].iloc[i]
        low = df["Low"].iloc[i]
        close = df["Close"].iloc[i]
        
        if side == "LONG":
            current_pnl_high = (high - entry_price) / entry_price
            current_pnl_low = (low - entry_price) / entry_price
            
            mfe = max(mfe, current_pnl_high)
            mae = min(mae, current_pnl_low)
            
            if low <= sl_price:
                outcome = -1
                exit_price = sl_price
                exit_idx = i
                bars_to_exit = i - entry_idx
                break
            elif high >= tp_price:
                outcome = 1
                exit_price = tp_price
                exit_idx = i
                bars_to_exit = i - entry_idx
                break
        else:
            current_pnl_high = (entry_price - low) / entry_price
            current_pnl_low = (entry_price - high) / entry_price
            
            mfe = max(mfe, current_pnl_high)
            mae = min(mae, current_pnl_low)
            
            if high >= sl_price:
                outcome = -1
                exit_price = sl_price
                exit_idx = i
                bars_to_exit = i - entry_idx
                break
            elif low <= tp_price:
                outcome = 1
                exit_price = tp_price
                exit_idx = i
                bars_to_exit = i - entry_idx
                break
    
    if outcome == 0:
        exit_price = df["Close"].iloc[exit_idx]
        bars_to_exit = exit_idx - entry_idx
    
    if side == "LONG":
        pnl_pct = (exit_price - entry_price) / entry_price
    else:
        pnl_pct = (entry_price - exit_price) / entry_price
    
    return BarrierResult(
        entry_idx=entry_idx,
        entry_price=entry_price,
        entry_time=entry_time,
        side=side,
        outcome=outcome,
        exit_price=exit_price,
        exit_idx=exit_idx,
        bars_to_exit=bars_to_exit,
        pnl_pct=pnl_pct,
        max_favorable_excursion=mfe,
        max_adverse_excursion=mae,
        atr_at_entry=atr_val,
        suggested_sl_distance=sl_distance,
        suggested_tp_distance=tp_distance,
    )


def filter_quality_labels(
    labels_df: pd.DataFrame,
    min_win_rate: float = 0.3,
    max_win_rate: float = 0.7,
    min_avg_pnl: float = -0.02,
) -> pd.DataFrame:
    """Filter labels to ensure balanced training data.
    
    Removes labels from periods with extreme win rates (too easy/too hard).
    
    Args:
        labels_df: DataFrame from compute_triple_barrier_labels
        min_win_rate: Minimum acceptable win rate
        max_win_rate: Maximum acceptable win rate
        min_avg_pnl: Minimum average PnL%
        
    Returns:
        Filtered DataFrame
    """
    if labels_df.empty:
        return labels_df
    
    win_rate = (labels_df["outcome"] == 1).mean()
    avg_pnl = labels_df["pnl_pct"].mean()
    
    logger.info(f"Label statistics: WR={win_rate:.2%}, Avg PnL={avg_pnl:.2%}")
    
    if win_rate < min_win_rate or win_rate > max_win_rate:
        logger.warning(
            f"Win rate {win_rate:.2%} вне диапазона [{min_win_rate:.2%}, {max_win_rate:.2%}]. "
            f"Рекомендуется скорректировать параметры барьеров."
        )
    
    filtered = labels_df[labels_df["pnl_pct"] >= min_avg_pnl].copy()
    
    if len(filtered) < len(labels_df):
        logger.info(f"Отфильтровано {len(labels_df) - len(filtered)} меток с низким PnL")
    
    return filtered


def create_training_targets(
    labels_df: pd.DataFrame,
    min_pnl_for_entry: float = 0.001,
) -> pd.DataFrame:
    """Create training targets from barrier labels.
    
    Transforms raw labels into targets suitable for multi-task learning:
    - entry_signal: binary (1 = profitable entry, 0 = unprofitable)
    - sl_distance_norm: SL distance normalized by ATR
    - tp_distance_norm: TP distance normalized by ATR
    - expected_pnl: expected PnL based on outcome
    
    Args:
        labels_df: DataFrame from compute_triple_barrier_labels
        min_pnl_for_entry: Minimum PnL% to consider entry profitable (default: 0.001)
            0.1% to account for commissions/slippage
        
    Returns:
        DataFrame with training targets
    """
    if labels_df.empty:
        return labels_df
    
    targets = labels_df.copy()
    
    # entry_signal = 1 для TP hits (outcome == 1)
    # + Timeouts с PnL > min_pnl_for_entry (учитывает комиссии)
    targets["entry_signal"] = (
        (targets["outcome"] == 1) |
        ((targets["outcome"] == 0) & (targets["pnl_pct"] > min_pnl_for_entry))
    ).astype(float)
    
    targets["sl_distance_norm"] = targets["sl_distance"] / (targets["atr_at_entry"] + 1e-10)
    targets["tp_distance_norm"] = targets["tp_distance"] / (targets["atr_at_entry"] + 1e-10)
    
    targets["expected_pnl"] = targets["pnl_pct"]
    
    # Regression targets: blend deterministic barrier distance with MAE/MFE.
    #
    #   barrier_sl = config.sl_atr_mult (deterministic, e.g. 1.0 ATR)
    #   barrier_tp = config.tp_atr_mult (deterministic, e.g. 1.5 ATR)
    #   mae_atr = |MAE| * entry_price / ATR (varied: 0.0-3.0 ATR)
    #   mfe_atr = MFE * entry_price / ATR (varied: 0.0-5.0 ATR)
    #
    # HYBRID_ALPHA = 0.8 (weight for barrier distance):
    #   - 80% barrier + 20% MAE/MFE ensures targets are CLOSE to deterministic
    #   - Model can learn a consistent mapping (not chasing noisy MAE/MFE)
    #   - Small MAE/MFE component adds market-condition variation to prevent
    #     the model from learning a pure constant
    #   - Previously HYBRID_ALPHA=0.5 created excessively noisy targets —
    #     MAE/MFE could be 5x the barrier distance, making targets unpredictable
    HYBRID_ALPHA = 0.8  # weight for barrier distance (vs 1-alpha for MAE/MFE)
    
    barrier_sl = targets.get("sl_distance_norm", pd.Series(1.0, index=targets.index))
    barrier_tp = targets.get("tp_distance_norm", pd.Series(1.5, index=targets.index))
    
    entry_prices = targets["entry_price"] if "entry_price" in targets.columns else 1.0
    mae_sl = targets["mae"].abs() * entry_prices / (targets["atr_at_entry"] + 1e-10)
    mfe_tp = targets["mfe"] * entry_prices / (targets["atr_at_entry"] + 1e-10)
    
    # Clip targets to exactly match model output range (see model_v2.py):
    #   sl_distance ∈ [0.5, 3.0], tp_distance ∈ [0.5, 5.0]
    targets["target_sl_atr"] = (HYBRID_ALPHA * barrier_sl + (1 - HYBRID_ALPHA) * mae_sl).clip(0.5, 3.0)
    targets["target_tp_atr"] = (HYBRID_ALPHA * barrier_tp + (1 - HYBRID_ALPHA) * mfe_tp).clip(0.5, 5.0)
    
    positive_rate = targets["entry_signal"].mean()
    logger.info(
        f"Training targets: {len(targets)} примеров, "
        f"entry_signal positive rate: {positive_rate:.2%}"
    )
    
    if positive_rate < 0.1:
        logger.warning(
            f"Очень низкий positive rate ({positive_rate:.2%}). "
            f"Рекомендуется уменьшить TP или увеличить SL."
        )
    elif positive_rate > 0.7:
        logger.warning(
            f"Очень высокий positive rate ({positive_rate:.2%}). "
            f"Возможно, барьеры слишком лёгкие для достижения."
        )
    
    return targets