"""
Phase 2 Target Builder
Creates targets for range prediction, ATR bins, and TP/SL classification
"""

import pandas as pd
import numpy as np
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Any
from enum import Enum
from loguru import logger
import yaml
from pathlib import Path


class RRConfig:
    """Risk:Reward configuration"""

    def __init__(self, sl: float, tp: float, name: str = None):
        self.sl = sl
        self.tp = tp
        self.rr_ratio = tp / sl
        self.name = name or f"rr_{int(self.rr_ratio)}_1"

    def __repr__(self):
        return f"RRConfig(sl={self.sl}, tp={self.tp}, rr={self.rr_ratio:.1f})"


@dataclass
class HorizonConfig:
    """Configuration for a prediction horizon"""
    name: str           # e.g., "15m", "1h"
    bars: int          # Number of 5m bars
    minutes: int       # Total minutes
    weight: float = 1.0
    enabled: bool = True


@dataclass
class TargetConfig:
    """Complete target configuration"""
    horizons: List[HorizonConfig]
    rr_configs: List[RRConfig]
    atr_bins: List[float] = field(default_factory=lambda: [0.25, 0.5, 1.0])
    start_offset: int = 1  # Start from t+1 (NOT t)


class Phase2TargetBuilder:
    """
    Builder for Phase 2 targets

    Creates:
    1. Delta targets (ΔHigh, ΔLow) - regression targets
    2. ATR-based bins - classification targets
    3. TP vs SL labels - binary classification targets
    """

    def __init__(self, config: Optional[TargetConfig] = None, config_path: str = None):
        """
        Initialize target builder

        Args:
            config: TargetConfig object
            config_path: Path to config file (alternative to config object)
        """
        if config is not None:
            self.config = config
        elif config_path:
            self.config = self._load_config(config_path)
        else:
            # Default configuration for XAUUSD
            self.config = TargetConfig(
                horizons=[
                    HorizonConfig(name="15m", bars=3, minutes=15, weight=0.6),
                    HorizonConfig(name="1h", bars=12, minutes=60, weight=0.4)
                ],
                rr_configs=[
                    RRConfig(sl=5.0, tp=10.0, name="rr_2_1"),
                    RRConfig(sl=5.0, tp=15.0, name="rr_3_1")
                ],
                atr_bins=[0.25, 0.5, 1.0],
                start_offset=1
            )

        logger.info(f"Initialized Phase2TargetBuilder with {len(self.config.horizons)} horizons")

    def _load_config(self, config_path: str) -> TargetConfig:
        """Load configuration from YAML file"""
        with open(config_path, 'r') as f:
            cfg = yaml.safe_load(f)

        horizons = [
            HorizonConfig(**h) for h in cfg.get('horizons', [])
        ]

        rr_configs = [
            RRConfig(**r) for r in cfg.get('targets', {}).get('tp_sl', {}).get('rr_configs', [])
        ]

        atr_thresholds = cfg.get('targets', {}).get('atr_bins', {}).get('thresholds', [0.25, 0.5, 1.0])

        return TargetConfig(
            horizons=horizons,
            rr_configs=rr_configs,
            atr_bins=atr_thresholds,
            start_offset=cfg.get('targets', {}).get('delta', {}).get('start_offset', 1)
        )

    def build_all_targets(
        self,
        df: pd.DataFrame,
        include_delta: bool = True,
        include_bins: bool = True,
        include_tp_sl: bool = True
    ) -> pd.DataFrame:
        """
        Build all Phase 2 targets

        Args:
            df: DataFrame with OHLCV data (must have 'high', 'low', 'close', 'ATR')
            include_delta: Include delta (range) targets
            include_bins: Include ATR-based bins
            include_tp_sl: Include TP vs SL labels

        Returns:
            DataFrame with all targets added
        """
        df = df.copy()

        # Verify required columns
        required = ['high', 'low', 'close']
        missing = [col for col in required if col not in df.columns]
        if missing:
            raise ValueError(f"Missing required columns: {missing}")

        # Build targets for each horizon
        for horizon in self.config.horizons:
            if not horizon.enabled:
                continue

            logger.info(f"Building targets for horizon: {horizon.name}")

            # 1. Delta targets (ΔHigh, ΔLow)
            if include_delta:
                df = self.calculate_delta_targets(df, horizon)

            # 2. ATR-based bins
            if include_bins and 'ATR' in df.columns:
                df = self.calculate_atr_bins(df, horizon)

            # 3. TP vs SL labels
            if include_tp_sl:
                for rr_config in self.config.rr_configs:
                    df = self.calculate_tp_sl_labels(df, horizon, rr_config)

        # Drop rows with NaN targets
        target_cols = [col for col in df.columns if col.startswith(('delta_', 'bin_', 'tp_first_'))]
        initial_len = len(df)
        df = df.dropna(subset=target_cols)
        dropped = initial_len - len(df)

        logger.info(f"Built {len(target_cols)} target columns, dropped {dropped} rows with NaN")
        return df

    def calculate_delta_targets(
        self,
        df: pd.DataFrame,
        horizon: HorizonConfig
    ) -> pd.DataFrame:
        """
        Calculate delta (range) targets

        CRITICAL: Start from t+1, NOT t (avoid data leakage)

        Δhigh = max(high[t+1 : t+horizon]) - close[t]
        Δlow  = close[t] - min(low[t+1 : t+horizon])

        Args:
            df: DataFrame with OHLCV
            horizon: Horizon configuration

        Returns:
            DataFrame with delta targets added
        """
        df = df.copy()
        start = self.config.start_offset  # Should be 1
        end = horizon.bars

        # Calculate future high (max of high from t+1 to t+horizon)
        future_highs = []
        for i in range(start, end + 1):
            future_highs.append(df['high'].shift(-i))

        future_high = pd.concat(future_highs, axis=1).max(axis=1)
        df[f'future_high_{horizon.name}'] = future_high

        # Calculate future low (min of low from t+1 to t+horizon)
        future_lows = []
        for i in range(start, end + 1):
            future_lows.append(df['low'].shift(-i))

        future_low = pd.concat(future_lows, axis=1).min(axis=1)
        df[f'future_low_{horizon.name}'] = future_low

        # Calculate deltas
        df[f'delta_high_{horizon.name}'] = future_high - df['close']
        df[f'delta_low_{horizon.name}'] = df['close'] - future_low

        # Also calculate normalized deltas (by ATR) if ATR available
        if 'ATR' in df.columns:
            df[f'delta_high_{horizon.name}_norm'] = df[f'delta_high_{horizon.name}'] / df['ATR']
            df[f'delta_low_{horizon.name}_norm'] = df[f'delta_low_{horizon.name}'] / df['ATR']

        logger.debug(f"Created delta targets for {horizon.name}")
        return df

    def calculate_atr_bins(
        self,
        df: pd.DataFrame,
        horizon: HorizonConfig,
        atr_column: str = 'ATR'
    ) -> pd.DataFrame:
        """
        Create ATR-based bins for classification

        Bins:
        - Bin 0: Δ < 0.25 * ATR (very small movement)
        - Bin 1: 0.25 * ATR ≤ Δ < 0.5 * ATR (small)
        - Bin 2: 0.5 * ATR ≤ Δ < 1.0 * ATR (medium)
        - Bin 3: Δ ≥ 1.0 * ATR (large)

        Args:
            df: DataFrame with delta targets and ATR
            horizon: Horizon configuration
            atr_column: Name of ATR column

        Returns:
            DataFrame with bin targets added
        """
        df = df.copy()

        if atr_column not in df.columns:
            logger.warning(f"ATR column '{atr_column}' not found, skipping bins")
            return df

        # Get delta columns
        delta_high_col = f'delta_high_{horizon.name}'
        delta_low_col = f'delta_low_{horizon.name}'

        if delta_high_col not in df.columns or delta_low_col not in df.columns:
            logger.warning(f"Delta columns not found for {horizon.name}, calculating first")
            df = self.calculate_delta_targets(df, horizon)

        # Calculate bins for delta_high
        delta_high_norm = df[delta_high_col] / df[atr_column]
        df[f'bin_high_{horizon.name}'] = self._assign_bins(delta_high_norm)

        # Calculate bins for delta_low
        delta_low_norm = df[delta_low_col] / df[atr_column]
        df[f'bin_low_{horizon.name}'] = self._assign_bins(delta_low_norm)

        logger.debug(f"Created ATR bins for {horizon.name}")
        return df

    def _assign_bins(self, normalized_delta: pd.Series) -> pd.Series:
        """
        Assign bins based on normalized delta values

        Args:
            normalized_delta: Delta values normalized by ATR

        Returns:
            Series with bin labels (0-3)
        """
        bins = pd.Series(index=normalized_delta.index, dtype='Int64')

        thresholds = self.config.atr_bins

        # Bin 0: < threshold[0]
        bins[normalized_delta < thresholds[0]] = 0

        # Bin 1: threshold[0] <= x < threshold[1]
        bins[(normalized_delta >= thresholds[0]) & (normalized_delta < thresholds[1])] = 1

        # Bin 2: threshold[1] <= x < threshold[2]
        bins[(normalized_delta >= thresholds[1]) & (normalized_delta < thresholds[2])] = 2

        # Bin 3: >= threshold[2]
        bins[normalized_delta >= thresholds[2]] = 3

        return bins

    def calculate_tp_sl_labels(
        self,
        df: pd.DataFrame,
        horizon: HorizonConfig,
        rr_config: RRConfig,
        direction: str = 'long'
    ) -> pd.DataFrame:
        """
        Calculate TP vs SL labels (binary classification)

        For each bar t, simulate a trade entry and check if TP or SL is hit first
        within the horizon window.

        For LONG trades:
        - Entry: close[t]
        - SL: entry - sl_value
        - TP: entry + tp_value
        - Label = 1 if price hits TP first, 0 if hits SL first or neither

        Args:
            df: DataFrame with OHLCV data
            horizon: Horizon configuration
            rr_config: R:R configuration (SL/TP values)
            direction: 'long' or 'short'

        Returns:
            DataFrame with TP/SL labels added
        """
        df = df.copy()
        start = self.config.start_offset
        end = horizon.bars

        # Column name
        col_name = f'tp_first_{horizon.name}_{rr_config.name}'

        if direction == 'long':
            labels = self._simulate_long_trades(
                df, start, end, rr_config.sl, rr_config.tp
            )
        else:
            labels = self._simulate_short_trades(
                df, start, end, rr_config.sl, rr_config.tp
            )

        df[col_name] = labels

        # Calculate some statistics
        valid_labels = labels.dropna()
        if len(valid_labels) > 0:
            tp_rate = valid_labels.mean()
            logger.info(f"TP/SL labels for {horizon.name} {rr_config.name}: "
                       f"TP rate = {tp_rate:.2%} ({valid_labels.sum():.0f}/{len(valid_labels)})")

        return df

    def _simulate_long_trades(
        self,
        df: pd.DataFrame,
        start_bar: int,
        end_bar: int,
        sl_value: float,
        tp_value: float
    ) -> pd.Series:
        """
        Simulate long trades and determine if TP or SL hits first

        Args:
            df: DataFrame with OHLCV
            start_bar: First bar to check (usually 1)
            end_bar: Last bar to check
            sl_value: Stop loss distance in price units
            tp_value: Take profit distance in price units

        Returns:
            Series with labels (1=TP first, 0=SL first or neither)
        """
        n = len(df)
        labels = pd.Series(index=df.index, dtype='float64')

        entry_prices = df['close'].values
        highs = df['high'].values
        lows = df['low'].values

        for i in range(n - end_bar):
            entry = entry_prices[i]
            sl_price = entry - sl_value
            tp_price = entry + tp_value

            tp_hit = False
            sl_hit = False
            tp_bar = end_bar + 1
            sl_bar = end_bar + 1

            # Check each bar in the horizon
            for j in range(start_bar, end_bar + 1):
                idx = i + j

                # Check if SL hit (low <= sl_price)
                if lows[idx] <= sl_price and not sl_hit:
                    sl_hit = True
                    sl_bar = j

                # Check if TP hit (high >= tp_price)
                if highs[idx] >= tp_price and not tp_hit:
                    tp_hit = True
                    tp_bar = j

            # Determine which hit first
            if tp_hit and sl_hit:
                # Both hit - which was first?
                labels.iloc[i] = 1 if tp_bar <= sl_bar else 0
            elif tp_hit:
                labels.iloc[i] = 1
            elif sl_hit:
                labels.iloc[i] = 0
            else:
                # Neither hit within horizon - count as loss
                labels.iloc[i] = 0

        return labels

    def _simulate_short_trades(
        self,
        df: pd.DataFrame,
        start_bar: int,
        end_bar: int,
        sl_value: float,
        tp_value: float
    ) -> pd.Series:
        """
        Simulate short trades and determine if TP or SL hits first

        Args:
            df: DataFrame with OHLCV
            start_bar: First bar to check (usually 1)
            end_bar: Last bar to check
            sl_value: Stop loss distance in price units
            tp_value: Take profit distance in price units

        Returns:
            Series with labels (1=TP first, 0=SL first or neither)
        """
        n = len(df)
        labels = pd.Series(index=df.index, dtype='float64')

        entry_prices = df['close'].values
        highs = df['high'].values
        lows = df['low'].values

        for i in range(n - end_bar):
            entry = entry_prices[i]
            sl_price = entry + sl_value  # SL is above for shorts
            tp_price = entry - tp_value  # TP is below for shorts

            tp_hit = False
            sl_hit = False
            tp_bar = end_bar + 1
            sl_bar = end_bar + 1

            # Check each bar in the horizon
            for j in range(start_bar, end_bar + 1):
                idx = i + j

                # Check if SL hit (high >= sl_price)
                if highs[idx] >= sl_price and not sl_hit:
                    sl_hit = True
                    sl_bar = j

                # Check if TP hit (low <= tp_price)
                if lows[idx] <= tp_price and not tp_hit:
                    tp_hit = True
                    tp_bar = j

            # Determine which hit first
            if tp_hit and sl_hit:
                labels.iloc[i] = 1 if tp_bar <= sl_bar else 0
            elif tp_hit:
                labels.iloc[i] = 1
            elif sl_hit:
                labels.iloc[i] = 0
            else:
                labels.iloc[i] = 0

        return labels

    def get_target_columns(self) -> Dict[str, List[str]]:
        """
        Get lists of target column names by type

        Returns:
            Dictionary with target column names grouped by type
        """
        targets = {
            'delta_regression': [],
            'delta_normalized': [],
            'bin_classification': [],
            'tp_sl_classification': []
        }

        for horizon in self.config.horizons:
            if not horizon.enabled:
                continue

            # Delta targets
            targets['delta_regression'].append(f'delta_high_{horizon.name}')
            targets['delta_regression'].append(f'delta_low_{horizon.name}')
            targets['delta_normalized'].append(f'delta_high_{horizon.name}_norm')
            targets['delta_normalized'].append(f'delta_low_{horizon.name}_norm')

            # Bin targets
            targets['bin_classification'].append(f'bin_high_{horizon.name}')
            targets['bin_classification'].append(f'bin_low_{horizon.name}')

            # TP/SL targets
            for rr in self.config.rr_configs:
                targets['tp_sl_classification'].append(f'tp_first_{horizon.name}_{rr.name}')

        return targets

    def get_target_statistics(self, df: pd.DataFrame) -> Dict[str, Any]:
        """
        Get statistics about target distributions

        Args:
            df: DataFrame with targets

        Returns:
            Dictionary with statistics
        """
        stats = {}
        target_cols = self.get_target_columns()

        # Delta statistics
        for col in target_cols['delta_regression']:
            if col in df.columns:
                stats[col] = {
                    'mean': df[col].mean(),
                    'std': df[col].std(),
                    'min': df[col].min(),
                    'max': df[col].max(),
                    'median': df[col].median()
                }

        # Bin distributions
        for col in target_cols['bin_classification']:
            if col in df.columns:
                dist = df[col].value_counts(normalize=True).sort_index()
                if len(dist) > 0:
                    stats[col] = {
                        'distribution': dist.to_dict(),
                        'majority_class': dist.idxmax(),
                        'majority_pct': dist.max()
                    }
                else:
                    stats[col] = {
                        'distribution': {},
                        'majority_class': None,
                        'majority_pct': 0.0
                    }

        # TP/SL distributions
        for col in target_cols['tp_sl_classification']:
            if col in df.columns:
                tp_rate = df[col].mean()
                stats[col] = {
                    'tp_rate': tp_rate,
                    'sl_rate': 1 - tp_rate,
                    'total_samples': df[col].notna().sum()
                }

        return stats


if __name__ == "__main__":
    # Test target builder
    import numpy as np

    # Create sample OHLCV data
    np.random.seed(42)
    n_samples = 1000

    # Generate realistic gold prices around $2000
    base_price = 2000
    returns = np.random.randn(n_samples) * 0.001  # 0.1% volatility per bar
    prices = base_price * np.cumprod(1 + returns)

    dates = pd.date_range(start='2024-01-01', periods=n_samples, freq='5min')

    df = pd.DataFrame({
        'open': prices,
        'high': prices * (1 + abs(np.random.randn(n_samples) * 0.001)),
        'low': prices * (1 - abs(np.random.randn(n_samples) * 0.001)),
        'close': prices * (1 + np.random.randn(n_samples) * 0.0005),
        'volume': np.random.randint(1000, 10000, n_samples),
        'ATR': np.full(n_samples, 5.0)  # $5 ATR
    }, index=dates)

    # Ensure high >= max(open, close) and low <= min(open, close)
    df['high'] = df[['open', 'high', 'close']].max(axis=1)
    df['low'] = df[['open', 'low', 'close']].min(axis=1)

    # Build targets
    builder = Phase2TargetBuilder()
    df_with_targets = builder.build_all_targets(df)

    print("\n=== Target Builder Test ===")
    print(f"Original shape: {len(df)}")
    print(f"With targets shape: {len(df_with_targets)}")
    print(f"\nTarget columns:")

    target_cols = builder.get_target_columns()
    for target_type, cols in target_cols.items():
        print(f"\n{target_type}:")
        for col in cols:
            if col in df_with_targets.columns:
                print(f"  - {col}")

    print("\n=== Target Statistics ===")
    stats = builder.get_target_statistics(df_with_targets)
    for col, stat in stats.items():
        print(f"\n{col}:")
        for k, v in stat.items():
            print(f"  {k}: {v}")

    print("\n=== Sample Data ===")
    sample_cols = ['close', 'ATR', 'delta_high_15m', 'delta_low_15m',
                   'bin_high_15m', 'tp_first_15m_rr_2_1']
    available_cols = [c for c in sample_cols if c in df_with_targets.columns]
    print(df_with_targets[available_cols].head(10))
