#!/usr/bin/env python3
"""
Sample Weighting for Large Movements
=====================================
Implements sample weighting to focus training on significant market movements.

Key Features:
- Dynamic factor based on rolling median with shift(1) to avoid leakage
- Move multiplier: actual_range / rolling_median
- Softplus attention mapping for smooth weight transition
- Higher weights for movements > factor threshold (e.g., 5 USD for XAUUSD)
- R:R ratio filtering (minimum 2:1)
- Integration with XGBoost sample_weight parameter

Author: Trading Strategist + ML Specialist
Version: 2.0.0 (2026-01-05) - Added dynamic factor weighting
"""

import numpy as np
import pandas as pd
from typing import Dict, Tuple, Optional
from dataclasses import dataclass
from loguru import logger


@dataclass
class SampleWeightConfig:
    """Configuration for sample weighting"""

    # Minimum movement threshold in USD (e.g., 5 for XAUUSD)
    min_movement_threshold: float = 5.0

    # Weight multiplier for samples above threshold
    large_movement_weight: float = 3.0

    # Weight for samples below threshold (not zero, to avoid ignoring completely)
    small_movement_weight: float = 0.3

    # Whether to use continuous weighting based on magnitude
    use_continuous_weighting: bool = True

    # Exponent for continuous weighting (higher = more emphasis on large moves)
    weight_exponent: float = 1.5

    # Minimum required TP:SL ratio to include sample
    min_rr_ratio: float = 2.0

    # Whether to filter out samples that don't meet ratio
    filter_low_ratio: bool = True

    # Minimum weight (never go below this)
    min_weight: float = 0.1

    # Maximum weight (cap to avoid extreme values)
    max_weight: float = 10.0

    # === DYNAMIC FACTOR WEIGHTING (NEW) ===
    # Use dynamic factor based on rolling median
    use_dynamic_factor: bool = True

    # Rolling window for median factor calculation
    factor_window: int = 200

    # Softplus beta (controls transition sharpness)
    # Higher beta = sharper transition from low to high attention
    softplus_beta: float = 4.0

    # Maximum attention weight from softplus
    softplus_w_max: float = 3.0


class SampleWeighter:
    """
    Calculates sample weights for XGBoost training.

    Strategy (v2.0 with dynamic factor):
    1. Compute dynamic factor: rolling median of range with shift(1) to avoid leakage
    2. Compute move multiplier: actual_movement / dynamic_factor
    3. Apply softplus attention: smooth mapping where m>1 gets attention
    4. Optionally filter samples that don't meet minimum R:R ratio

    Usage:
        weighter = SampleWeighter(SampleWeightConfig(min_movement_threshold=5.0))
        df = weighter.calculate_movement_magnitude(df, horizon_bars=3)
        weights, valid_mask = weighter.compute_sample_weights(df)

        # Use with XGBoost
        model.fit(X[valid_mask], y[valid_mask], sample_weight=weights[valid_mask])
    """

    def __init__(self, config: SampleWeightConfig = None):
        self.config = config or SampleWeightConfig()

    # =========================================================================
    # NEW: Dynamic Factor Methods
    # =========================================================================

    def compute_factor_median_range(
        self,
        df: pd.DataFrame,
        window: int = None
    ) -> pd.Series:
        """
        Compute dynamic factor as rolling median of range with shift(1).

        IMPORTANT: shift(1) avoids data leakage by only using past information.

        Args:
            df: DataFrame with High/Low columns
            window: Rolling window size (default from config)

        Returns:
            Series with dynamic factor for each row
        """
        window = window or self.config.factor_window
        min_periods = window // 2

        # Handle column name variations
        high_col = 'High' if 'High' in df.columns else 'high'
        low_col = 'Low' if 'Low' in df.columns else 'low'

        # Calculate candle range
        candle_range = (df[high_col] - df[low_col]).abs()

        # Rolling median with shift(1) to avoid leakage
        factor = candle_range.rolling(
            window=window,
            min_periods=min_periods
        ).median().shift(1)

        return factor

    def compute_move_multiplier(
        self,
        df: pd.DataFrame,
        factor: pd.Series = None
    ) -> pd.Series:
        """
        Compute movement multiplier: actual_range / factor.

        Interpretation:
        - m < 1: Movement smaller than typical (noise)
        - m = 1: Typical movement
        - m > 1: Larger than typical movement (signal)

        Args:
            df: DataFrame with High/Low columns
            factor: Pre-computed factor (computed if None)

        Returns:
            Series with move multiplier for each row
        """
        if factor is None:
            factor = self.compute_factor_median_range(df)

        # Handle column name variations
        high_col = 'High' if 'High' in df.columns else 'high'
        low_col = 'Low' if 'Low' in df.columns else 'low'

        # Calculate actual range
        delta = (df[high_col] - df[low_col]).abs()

        # Compute multiplier
        epsilon = 1e-12
        m = delta / (factor + epsilon)

        return m

    def weight_smooth(
        self,
        m: np.ndarray,
        w_max: float = None,
        beta: float = None
    ) -> np.ndarray:
        """
        Apply softplus mapping for smooth attention weights.

        Formula: w = log1p(exp(beta * (m - 1))) / beta

        Behavior:
        - m < 1: w approaches 0 (ignore noise)
        - m = 1: w ~ 0 (typical movement)
        - m > 1: w approaches (m - 1) linearly

        Args:
            m: Move multiplier array
            w_max: Maximum weight cap (default from config)
            beta: Softplus beta parameter (default from config)

        Returns:
            Array of attention weights
        """
        w_max = w_max or self.config.softplus_w_max
        beta = beta or self.config.softplus_beta

        # Convert to numpy if pandas
        if isinstance(m, pd.Series):
            m = m.values

        # Apply softplus with numerical stability
        x = beta * (m - 1.0)
        w = np.where(
            x > 20,  # For large x, softplus(x) ~ x
            x / beta,
            np.log1p(np.exp(x)) / beta
        )

        # Clip to [0, w_max]
        w = np.clip(w, 0.0, w_max)

        return w

    def compute_dynamic_weights(
        self,
        df: pd.DataFrame,
        target_high: np.ndarray = None,
        target_low: np.ndarray = None
    ) -> np.ndarray:
        """
        Compute weights using dynamic factor approach.

        Can use either:
        - Candle range (if targets not provided)
        - Target movement (if targets provided - more aligned with prediction task)

        Args:
            df: DataFrame with OHLCV data
            target_high: Optional target high values
            target_low: Optional target low values

        Returns:
            Array of sample weights
        """
        # Compute dynamic factor from candle ranges
        factor = self.compute_factor_median_range(df)

        if target_high is not None and target_low is not None:
            # Use target movement for weighting (better alignment)
            total_movement = np.abs(target_high) + np.abs(target_low)
            m = total_movement / (factor.values + 1e-12)
        else:
            # Use candle range for weighting
            m = self.compute_move_multiplier(df, factor).values

        # Apply softplus weighting
        weights = self.weight_smooth(m)

        # Handle NaN
        nan_mask = np.isnan(weights) | np.isnan(factor.values)
        weights[nan_mask] = 1.0  # Default weight

        return weights

    # =========================================================================
    # Original Methods (kept for backward compatibility)
    # =========================================================================

    def calculate_movement_magnitude(
        self,
        df: pd.DataFrame,
        horizon_bars: int = 3
    ) -> pd.DataFrame:
        """
        Calculate the actual movement magnitude for each sample.

        Uses the corrected target formula:
        - target_high = MAX(high[t+1], high[t+2], high[t+3]) - close[t]
        - target_low = close[t] - MIN(low[t+1], low[t+2], low[t+3])

        Args:
            df: DataFrame with OHLCV data (columns: open, high, low, close or Open, High, Low, Close)
            horizon_bars: Number of future bars to consider

        Returns:
            DataFrame with added columns: target_high, target_low, total_range
        """
        df = df.copy()

        # Handle column name variations
        close_col = 'close' if 'close' in df.columns else 'Close'
        high_col = 'high' if 'high' in df.columns else 'High'
        low_col = 'low' if 'low' in df.columns else 'Low'

        close = df[close_col].values
        high = df[high_col].values
        low = df[low_col].values
        n = len(df)

        target_high = np.full(n, np.nan)
        target_low = np.full(n, np.nan)

        for i in range(n - horizon_bars):
            current_close = close[i]

            # MAX of highs in [t+1, t+2, ..., t+horizon] - close[t]
            future_highs = high[i+1:i+1+horizon_bars]
            target_high[i] = np.max(future_highs) - current_close

            # close[t] - MIN of lows in [t+1, t+2, ..., t+horizon]
            future_lows = low[i+1:i+1+horizon_bars]
            target_low[i] = current_close - np.min(future_lows)

        df['target_high'] = target_high
        df['target_low'] = target_low
        df['total_range'] = target_high + target_low

        # Calculate R:R ratios
        epsilon = 0.0001
        df['rr_long'] = np.where(target_low > epsilon, target_high / target_low, 0)
        df['rr_short'] = np.where(target_high > epsilon, target_low / target_high, 0)
        df['rr_best'] = np.maximum(df['rr_long'], df['rr_short'])

        # Determine suggested direction
        df['suggested_direction'] = np.where(
            df['rr_long'] >= df['rr_short'], 'LONG', 'SHORT'
        )

        logger.info(f"Calculated movement magnitudes for {n} samples")
        logger.info(f"  Mean target_high: {np.nanmean(target_high):.2f}")
        logger.info(f"  Mean target_low: {np.nanmean(target_low):.2f}")
        logger.info(f"  Samples with R:R >= 2: {(df['rr_best'] >= 2).sum()}")

        return df

    def compute_sample_weights(
        self,
        df: pd.DataFrame,
        target_high_col: str = 'target_high',
        target_low_col: str = 'target_low'
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Compute sample weights based on movement magnitude.

        If use_dynamic_factor=True (default), uses the new softplus-based
        attention mechanism. Otherwise falls back to the original power-based
        continuous weighting.

        Args:
            df: DataFrame with target columns
            target_high_col: Column name for high target
            target_low_col: Column name for low target

        Returns:
            Tuple of (sample_weights, valid_mask)
            - sample_weights: Array of weights for each sample
            - valid_mask: Boolean mask of samples to include
        """
        target_high = df[target_high_col].values
        target_low = df[target_low_col].values
        total_movement = np.abs(target_high) + np.abs(target_low)

        n = len(df)
        weights = np.ones(n)
        valid_mask = np.ones(n, dtype=bool)

        # Handle NaN
        nan_mask = np.isnan(target_high) | np.isnan(target_low)
        valid_mask[nan_mask] = False

        # Calculate R:R ratio for filtering
        if self.config.filter_low_ratio:
            rr_col = 'rr_best' if 'rr_best' in df.columns else None

            if rr_col:
                rr_values = df[rr_col].values
            else:
                # Calculate R:R if not present
                epsilon = 0.0001
                rr_long = np.where(target_low > epsilon, target_high / target_low, 0)
                rr_short = np.where(target_high > epsilon, target_low / target_high, 0)
                rr_values = np.maximum(rr_long, rr_short)

            # Reduce weight for samples that don't meet minimum R:R
            low_rr_mask = rr_values < self.config.min_rr_ratio
            weights[low_rr_mask] = self.config.min_weight

            logger.info(f"Reduced weight for {low_rr_mask.sum()} samples with R:R < {self.config.min_rr_ratio}")

        # =====================================================================
        # NEW: Dynamic Factor Weighting (v2.0)
        # =====================================================================
        if self.config.use_dynamic_factor:
            # Compute dynamic factor from rolling median with shift(1)
            factor = self.compute_factor_median_range(df)

            # Compute move multiplier using target movement
            m = total_movement / (factor.values + 1e-12)

            # Apply softplus weighting
            dynamic_weights = self.weight_smooth(m)

            # Handle NaN from factor warmup
            factor_nan = np.isnan(factor.values)
            dynamic_weights[factor_nan] = 1.0  # Default weight for warmup period

            # Combine with existing weights (R:R filtering)
            weights[valid_mask] *= dynamic_weights[valid_mask]

            logger.info(f"Applied dynamic factor weighting (softplus)")
            logger.info(f"  Mean multiplier: {np.nanmean(m[valid_mask]):.2f}")
            logger.info(f"  Samples with m>2: {(m[valid_mask] > 2).sum()}")

        # =====================================================================
        # ORIGINAL: Power-based continuous weighting (fallback)
        # =====================================================================
        elif self.config.use_continuous_weighting:
            # Continuous weighting: weight = (movement / threshold) ^ exponent
            normalized_movement = total_movement / self.config.min_movement_threshold

            movement_weights = np.power(
                np.clip(normalized_movement, 0.1, 10),
                self.config.weight_exponent
            )

            # Only apply to valid samples
            weights[valid_mask] *= movement_weights[valid_mask]
        else:
            # Binary weighting
            large_move_mask = total_movement >= self.config.min_movement_threshold
            weights[large_move_mask & valid_mask] = self.config.large_movement_weight
            weights[~large_move_mask & valid_mask] = self.config.small_movement_weight

        # Clip weights to configured range
        weights = np.clip(weights, self.config.min_weight, self.config.max_weight)

        # Normalize to have mean = 1 for valid samples
        valid_weights = weights[valid_mask]
        if len(valid_weights) > 0 and valid_weights.mean() > 0:
            weights[valid_mask] = weights[valid_mask] / valid_weights.mean()

        # Log statistics
        valid_weights = weights[valid_mask]
        if len(valid_weights) > 0:
            logger.info(f"Sample weights computed:")
            logger.info(f"  Valid samples: {valid_mask.sum()} / {n} ({100*valid_mask.sum()/n:.1f}%)")
            logger.info(f"  Weight range: [{valid_weights.min():.3f}, {valid_weights.max():.3f}]")
            logger.info(f"  Weight mean: {valid_weights.mean():.3f}")
            logger.info(f"  High attention (w>1.5): {(valid_weights > 1.5).sum()} samples")

        return weights, valid_mask

    def apply_asymmetry_filter(
        self,
        df: pd.DataFrame,
        min_asymmetry: float = 2.0,
        target_high_col: str = 'target_high',
        target_low_col: str = 'target_low'
    ) -> Tuple[pd.DataFrame, np.ndarray]:
        """
        Filter samples to only include those with clear directional bias.

        A sample is included if:
        - target_high >= min_asymmetry * target_low (LONG opportunity)
        - OR target_low >= min_asymmetry * target_high (SHORT opportunity)

        Args:
            df: DataFrame with target columns
            min_asymmetry: Minimum asymmetry ratio (default 2.0 = 2:1)
            target_high_col: Column name for high target
            target_low_col: Column name for low target

        Returns:
            Tuple of (filtered_df, mask)
        """
        target_high = df[target_high_col].values
        target_low = df[target_low_col].values

        # Check asymmetry in both directions
        is_long_opportunity = target_high >= min_asymmetry * target_low
        is_short_opportunity = target_low >= min_asymmetry * target_high

        # Also require minimum absolute movement
        min_movement = self.config.min_movement_threshold
        has_min_movement = (target_high + target_low) >= min_movement

        # Combine conditions
        mask = (is_long_opportunity | is_short_opportunity) & has_min_movement

        # Handle NaN
        mask = mask & ~np.isnan(target_high) & ~np.isnan(target_low)

        filtered_df = df[mask].copy()

        logger.info(f"Asymmetry filter applied:")
        logger.info(f"  Original samples: {len(df)}")
        logger.info(f"  Filtered samples: {len(filtered_df)} ({100*len(filtered_df)/len(df):.1f}%)")
        logger.info(f"  LONG opportunities: {is_long_opportunity.sum()}")
        logger.info(f"  SHORT opportunities: {is_short_opportunity.sum()}")

        return filtered_df, mask


def create_sample_weights_for_training(
    df: pd.DataFrame,
    config: SampleWeightConfig = None,
    horizon_bars: int = 3
) -> Tuple[np.ndarray, np.ndarray, pd.DataFrame]:
    """
    Convenience function to create sample weights ready for XGBoost training.

    Args:
        df: Raw OHLCV DataFrame
        config: SampleWeightConfig (uses defaults if None)
        horizon_bars: Number of future bars for target calculation

    Returns:
        Tuple of (sample_weights, valid_mask, df_with_targets)
    """
    weighter = SampleWeighter(config)

    # Calculate targets
    df_with_targets = weighter.calculate_movement_magnitude(df, horizon_bars)

    # Compute weights
    weights, valid_mask = weighter.compute_sample_weights(df_with_targets)

    return weights, valid_mask, df_with_targets


if __name__ == "__main__":
    # Test the module
    import numpy as np

    print("Testing SampleWeighter...")

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

    dates = pd.date_range('2025-01-01', periods=n, freq='5min')
    price = 2650 + np.cumsum(np.random.randn(n) * 0.5)

    df = pd.DataFrame({
        'open': price,
        'high': price + np.abs(np.random.randn(n)) * 3,
        'low': price - np.abs(np.random.randn(n)) * 3,
        'close': price + np.random.randn(n) * 0.5,
        'volume': np.random.randint(100, 1000, n)
    }, index=dates)

    # Test weighter
    config = SampleWeightConfig(
        min_movement_threshold=5.0,
        min_rr_ratio=2.0,
        use_continuous_weighting=True
    )

    weights, valid_mask, df_with_targets = create_sample_weights_for_training(
        df, config, horizon_bars=3
    )

    print(f"\nResults:")
    print(f"  Total samples: {len(df)}")
    print(f"  Valid samples: {valid_mask.sum()}")
    print(f"  Weight distribution: min={weights[valid_mask].min():.2f}, max={weights[valid_mask].max():.2f}")
    print(f"  Samples with R:R >= 2: {(df_with_targets['rr_best'] >= 2).sum()}")
