#!/usr/bin/env python3
"""
Dynamic Factor-Based Sample Weighting
=====================================
Implements improved sample weighting based on dynamic volatility factor.

Key Features:
1. Factor dinamico: Mediana rolling del rango con shift(1) para evitar leakage
2. Multiplicador de movimiento: delta / factor
3. Peso suave (softplus): Atencion proporcional al multiplicador

Interpretation:
- m < 1 = ruido tipico (peso bajo)
- m ~ 2 = movimiento 2x normal (peso ~2)
- m ~ 3 = movimiento 3x normal (peso ~3)

Author: Trading Strategist + ML Specialist
Version: 1.0.0
Created: 2026-01-05
"""

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


@dataclass
class DynamicFactorConfig:
    """Configuration for dynamic factor-based weighting"""

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

    # Minimum periods for rolling calculation
    min_periods: Optional[int] = None  # Defaults to window // 2

    # Maximum attention weight
    w_max: float = 3.0

    # Softplus beta (controls transition sharpness)
    beta: float = 4.0

    # Floor epsilon to avoid division by zero
    epsilon: float = 1e-12

    # Use ATR weighting in combination
    use_atr_weighting: bool = True

    # Session weighting (disabled by default)
    use_session_weighting: bool = False

    # Minimum weight (floor)
    min_weight: float = 0.0

    def __post_init__(self):
        if self.min_periods is None:
            self.min_periods = self.factor_window // 2


class DynamicFactorWeighter:
    """
    Computes sample weights based on dynamic volatility factor.

    Algorithm:
    1. Calculate rolling median of candle range with shift(1) to avoid leakage
    2. Compute move multiplier: actual_range / rolling_median
    3. Apply softplus mapping to get smooth attention weights

    Usage:
        weighter = DynamicFactorWeighter(DynamicFactorConfig())
        weights = weighter.compute_weights(df_ohlcv)

        # For training
        model.fit(X, y, sample_weight=weights)
    """

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

    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).

        Uses shift(1) to avoid data leakage - we only use information
        that would have been available at prediction time.

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

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

        # 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
        # shift(1) ensures we only use past information
        factor = candle_range.rolling(
            window=window,
            min_periods=min_periods
        ).median().shift(1)

        logger.debug(f"Factor stats: mean={factor.mean():.4f}, std={factor.std():.4f}")

        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 with epsilon to avoid division by zero
        m = delta / (factor + self.config.epsilon)

        logger.debug(f"Multiplier stats: mean={m.mean():.2f}, max={m.max():.2f}")

        return m

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

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

        This creates a smooth transition:
        - m < 1: w approaches 0 (ignore noise)
        - m = 1: w ~ 0 (typical movement, neutral)
        - m > 1: w approaches (m - 1) linearly

        Args:
            m: Move multiplier array/series
            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.w_max
        beta = beta or self.config.beta

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

        # Apply softplus: log1p(exp(x)) is numerically stable
        x = beta * (m - 1.0)

        # Numerical stability for large x
        # softplus(x) = x for large x, log1p(exp(x)) for small x
        w = np.where(
            x > 20,  # For large x, softplus(x) ~ x
            x / beta,
            np.log1p(np.exp(x)) / beta
        )

        # Clip to [min_weight, w_max]
        w = np.clip(w, self.config.min_weight, w_max)

        return w

    def compute_weights(
        self,
        df: pd.DataFrame,
        normalize: bool = True
    ) -> np.ndarray:
        """
        Compute final sample weights combining all factors.

        Pipeline:
        1. Compute dynamic factor (rolling median with shift)
        2. Compute move multiplier
        3. Apply softplus mapping
        4. Optionally normalize to mean=1

        Args:
            df: DataFrame with OHLCV data
            normalize: Whether to normalize weights to mean=1

        Returns:
            Array of sample weights
        """
        # Step 1: Compute dynamic factor
        factor = self.compute_factor_median_range(df)

        # Step 2: Compute move multiplier
        multiplier = self.compute_move_multiplier(df, factor)

        # Step 3: Apply softplus mapping
        weights = self.weight_smooth(multiplier)

        # Handle NaN values (from rolling window warmup)
        nan_mask = np.isnan(weights) | np.isnan(multiplier.values)
        weights[nan_mask] = 1.0  # Default weight for NaN

        # Step 4: Normalize
        if normalize:
            valid_mask = ~nan_mask
            if valid_mask.sum() > 0 and weights[valid_mask].mean() > 0:
                weights[valid_mask] = weights[valid_mask] / weights[valid_mask].mean()

        # Log statistics
        valid_weights = weights[~nan_mask]
        if len(valid_weights) > 0:
            logger.info(f"Dynamic factor weights computed:")
            logger.info(f"  Valid samples: {(~nan_mask).sum()} / {len(weights)}")
            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

    def compute_combined_weights(
        self,
        df: pd.DataFrame,
        atr_weights: np.ndarray = None,
        session_weights: np.ndarray = None,
        normalize: bool = True
    ) -> np.ndarray:
        """
        Compute combined weights from multiple sources.

        Args:
            df: DataFrame with OHLCV data
            atr_weights: Optional ATR-based weights
            session_weights: Optional session-based weights (disabled by default)
            normalize: Whether to normalize final weights

        Returns:
            Array of combined sample weights
        """
        # Dynamic factor weights (primary)
        weights = self.compute_weights(df, normalize=False)

        # Combine with ATR weights if provided and enabled
        if self.config.use_atr_weighting and atr_weights is not None:
            weights = weights * atr_weights
            logger.info("Combined with ATR weights")

        # Combine with session weights if provided and enabled
        if self.config.use_session_weighting and session_weights is not None:
            weights = weights * session_weights
            logger.info("Combined with session weights")

        # Normalize
        if normalize:
            valid_mask = ~np.isnan(weights)
            if valid_mask.sum() > 0 and weights[valid_mask].mean() > 0:
                weights[valid_mask] = weights[valid_mask] / weights[valid_mask].mean()

        return weights


def compute_target_weights_for_training(
    df: pd.DataFrame,
    target_high: np.ndarray,
    target_low: np.ndarray,
    config: DynamicFactorConfig = None
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Convenience function to compute weights for ML training.

    Uses the target movements (not just candle range) to compute attention.
    This is more aligned with what we're trying to predict.

    Args:
        df: DataFrame with OHLCV data
        target_high: Target high values (delta from close)
        target_low: Target low values (delta from close)
        config: Configuration (uses defaults if None)

    Returns:
        Tuple of (weights, valid_mask)
    """
    config = config or DynamicFactorConfig()
    weighter = DynamicFactorWeighter(config)

    # Compute base factor from candle ranges
    factor = weighter.compute_factor_median_range(df)

    # Compute total target movement
    total_target = np.abs(target_high) + np.abs(target_low)

    # Compute multiplier based on target vs factor
    multiplier = total_target / (factor.values + config.epsilon)

    # Apply softplus weighting
    weights = weighter.weight_smooth(multiplier)

    # Valid mask (exclude NaN and warmup period)
    valid_mask = (
        ~np.isnan(weights) &
        ~np.isnan(target_high) &
        ~np.isnan(target_low) &
        ~np.isnan(factor.values)
    )

    # Normalize valid weights
    if valid_mask.sum() > 0 and weights[valid_mask].mean() > 0:
        weights[valid_mask] = weights[valid_mask] / weights[valid_mask].mean()

    # Set NaN weights to 0
    weights[~valid_mask] = 0.0

    logger.info(f"Target-based weights computed:")
    logger.info(f"  Valid samples: {valid_mask.sum()} / {len(weights)}")
    logger.info(f"  High attention samples (w>1.5): {(weights[valid_mask] > 1.5).sum()}")

    return weights, valid_mask


if __name__ == "__main__":
    # Test the module
    print("Testing DynamicFactorWeighter...")

    # Create sample OHLCV 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)

    # Simulate varying volatility
    volatility = np.where(
        (dates.hour >= 13) & (dates.hour < 16),  # London/NY overlap
        5.0,  # High volatility
        2.0   # Normal volatility
    )

    df = pd.DataFrame({
        'Open': price,
        'High': price + np.abs(np.random.randn(n)) * volatility,
        'Low': price - np.abs(np.random.randn(n)) * volatility,
        'Close': price + np.random.randn(n) * 0.5,
        'Volume': np.random.randint(100, 1000, n)
    }, index=dates)

    # Test weighter
    config = DynamicFactorConfig(
        factor_window=100,
        w_max=3.0,
        beta=4.0
    )

    weighter = DynamicFactorWeighter(config)

    # Compute factor
    factor = weighter.compute_factor_median_range(df)
    print(f"\nFactor stats:")
    print(f"  Mean: {factor.mean():.4f}")
    print(f"  Std: {factor.std():.4f}")

    # Compute multiplier
    multiplier = weighter.compute_move_multiplier(df, factor)
    print(f"\nMultiplier stats:")
    print(f"  Mean: {multiplier.mean():.2f}")
    print(f"  > 2x: {(multiplier > 2).sum()} samples")
    print(f"  > 3x: {(multiplier > 3).sum()} samples")

    # Compute final weights
    weights = weighter.compute_weights(df)
    print(f"\nFinal weight stats:")
    print(f"  Mean: {weights.mean():.3f}")
    print(f"  Range: [{weights.min():.3f}, {weights.max():.3f}]")
    print(f"  High attention (w>1.5): {(weights > 1.5).sum()} samples")

    # Test with target-based weighting
    print("\n" + "="*50)
    print("Testing target-based weighting...")

    # Simulate targets
    target_high = np.random.exponential(3, n)  # Mostly small, some large
    target_low = np.random.exponential(3, n)

    target_weights, valid_mask = compute_target_weights_for_training(
        df, target_high, target_low, config
    )

    print(f"Target weight stats:")
    print(f"  Valid: {valid_mask.sum()}")
    print(f"  High attention: {(target_weights[valid_mask] > 1.5).sum()}")
