#!/usr/bin/env python3
"""
Session and Volatility-Based Sample Weighting
==============================================
Implements attention mechanism based on trading sessions and market volatility.

Key Features:
- Higher weights during London/NY overlap (maximum liquidity)
- Lower weights during Asian session and off-hours
- ATR-based volatility weighting (ignore lateral/ranging markets)
- Cyclical time encoding for model features

Author: Trading Strategist + ML Specialist
Version: 1.0.0
"""

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


@dataclass
class SessionWeightConfig:
    """Configuration for session-based weighting"""

    # === MAIN SWITCHES ===
    # By default, only use ATR volatility weighting (not session/hour weights)
    use_session_weighting: bool = False  # Disabled by default
    use_atr_weighting: bool = True       # Enabled by default

    # Session definitions (UTC hours) - only used if use_session_weighting=True
    london_start: int = 8
    london_end: int = 16
    ny_start: int = 13
    ny_end: int = 21
    tokyo_start: int = 0
    tokyo_end: int = 8

    # Overlap periods (highest volatility)
    london_ny_overlap_start: int = 13  # 1 PM UTC
    london_ny_overlap_end: int = 16    # 4 PM UTC

    # Session weights (multipliers) - only used if use_session_weighting=True
    overlap_weight: float = 2.0      # London/NY overlap - highest priority
    london_weight: float = 1.5       # London session alone
    ny_weight: float = 1.3           # NY session alone (outside overlap)
    tokyo_weight: float = 0.7        # Asian session - lower priority
    off_hours_weight: float = 0.3    # Outside main sessions

    # Volatility-based adjustments (ATR)
    atr_high_percentile: float = 75  # Above this = high volatility
    atr_low_percentile: float = 25   # Below this = low volatility (lateral)
    atr_high_weight_boost: float = 1.5   # Multiplier for high ATR periods
    atr_low_weight_penalty: float = 0.3  # Multiplier for low ATR (lateral) periods

    # Lookback for ATR percentile calculation
    atr_lookback: int = 100


class SessionVolatilityWeighter:
    """
    Calculates sample weights based on trading session and volatility.

    Strategy:
    1. Assign base weights by trading session (London/NY overlap gets highest)
    2. Adjust weights by relative ATR (high vol = higher weight, lateral = lower)
    3. Combine both factors for final weight

    Usage:
        weighter = SessionVolatilityWeighter(SessionWeightConfig())
        session_weights = weighter.compute_session_weights(df)
        atr_weights = weighter.compute_atr_weights(df)
        combined = weighter.compute_combined_weights(df)
    """

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

    def get_session_name(self, hour: int) -> str:
        """Determine the trading session name for a given hour (UTC)."""
        cfg = self.config

        # Check overlap first (highest priority)
        if cfg.london_ny_overlap_start <= hour < cfg.london_ny_overlap_end:
            return 'overlap'
        elif cfg.london_start <= hour < cfg.london_end:
            return 'london'
        elif cfg.ny_start <= hour < cfg.ny_end:
            return 'ny'
        elif cfg.tokyo_start <= hour < cfg.tokyo_end:
            return 'tokyo'
        else:
            return 'off_hours'

    def compute_session_weights(
        self,
        df: pd.DataFrame,
        timestamp_col: str = None
    ) -> np.ndarray:
        """
        Compute weights based on trading session.

        Args:
            df: DataFrame with datetime index or timestamp column
            timestamp_col: Column name with timestamps (if not using index)

        Returns:
            Array of session-based weights
        """
        # Get hours
        if timestamp_col and timestamp_col in df.columns:
            timestamps = pd.to_datetime(df[timestamp_col])
            hours = timestamps.dt.hour.values
        elif isinstance(df.index, pd.DatetimeIndex):
            hours = df.index.hour.values
        else:
            logger.warning("No datetime index found, using uniform session weights")
            return np.ones(len(df))

        n = len(df)
        weights = np.ones(n)
        cfg = self.config

        # Identify sessions
        is_overlap = (hours >= cfg.london_ny_overlap_start) & (hours < cfg.london_ny_overlap_end)
        is_london = (hours >= cfg.london_start) & (hours < cfg.london_end) & ~is_overlap
        is_ny = (hours >= cfg.ny_start) & (hours < cfg.ny_end) & ~is_overlap
        is_tokyo = (hours >= cfg.tokyo_start) & (hours < cfg.tokyo_end)

        # Apply weights
        weights[:] = cfg.off_hours_weight  # Default
        weights[is_tokyo] = cfg.tokyo_weight
        weights[is_london] = cfg.london_weight
        weights[is_ny] = cfg.ny_weight
        weights[is_overlap] = cfg.overlap_weight

        # Log distribution
        logger.info("Session weight distribution:")
        logger.info(f"  Overlap: {is_overlap.sum()} samples ({100*is_overlap.mean():.1f}%) - weight {cfg.overlap_weight}")
        logger.info(f"  London: {is_london.sum()} samples ({100*is_london.mean():.1f}%) - weight {cfg.london_weight}")
        logger.info(f"  NY: {is_ny.sum()} samples ({100*is_ny.mean():.1f}%) - weight {cfg.ny_weight}")
        logger.info(f"  Tokyo: {is_tokyo.sum()} samples ({100*is_tokyo.mean():.1f}%) - weight {cfg.tokyo_weight}")

        return weights

    def compute_atr_weights(
        self,
        df: pd.DataFrame,
        atr_col: str = 'atr'
    ) -> np.ndarray:
        """
        Compute weights based on relative ATR (volatility).

        High ATR periods get higher weights (more informative for trends).
        Low ATR periods get lower weights (lateral/consolidation = noise).

        Args:
            df: DataFrame with ATR column
            atr_col: Name of ATR column

        Returns:
            Array of ATR-based weights
        """
        # Check for ATR column
        if atr_col not in df.columns and 'ATR' not in df.columns:
            logger.warning(f"ATR column not found, calculating from OHLC")
            df = self._calculate_atr(df)
            atr_col = 'atr'

        actual_col = atr_col if atr_col in df.columns else 'ATR'
        atr = df[actual_col].values

        n = len(df)
        weights = np.ones(n)
        cfg = self.config
        lookback = cfg.atr_lookback

        # Calculate rolling percentiles
        for i in range(lookback, n):
            window = atr[max(0, i-lookback):i]
            current_atr = atr[i]

            if np.isnan(current_atr) or len(window) == 0:
                continue

            # Clean window
            window = window[~np.isnan(window)]
            if len(window) == 0:
                continue

            # Calculate percentile of current ATR vs recent history
            percentile = (window < current_atr).mean() * 100

            if percentile >= cfg.atr_high_percentile:
                # High volatility - boost weight
                weights[i] = cfg.atr_high_weight_boost
            elif percentile <= cfg.atr_low_percentile:
                # Low volatility (lateral) - reduce weight
                weights[i] = cfg.atr_low_weight_penalty

        # Log statistics
        high_vol_mask = weights > 1.0
        low_vol_mask = weights < 1.0
        logger.info("ATR weight distribution:")
        logger.info(f"  High volatility: {high_vol_mask.sum()} samples ({100*high_vol_mask.mean():.1f}%) - weight {cfg.atr_high_weight_boost}")
        logger.info(f"  Low volatility: {low_vol_mask.sum()} samples ({100*low_vol_mask.mean():.1f}%) - weight {cfg.atr_low_weight_penalty}")
        logger.info(f"  Normal: {(~high_vol_mask & ~low_vol_mask).sum()} samples")

        return weights

    def _calculate_atr(self, df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """Calculate ATR if not present."""
        df = df.copy()

        high_col = 'high' if 'high' in df.columns else 'High'
        low_col = 'low' if 'low' in df.columns else 'Low'
        close_col = 'close' if 'close' in df.columns else 'Close'

        high = df[high_col]
        low = df[low_col]
        close = df[close_col]

        tr1 = high - low
        tr2 = abs(high - close.shift(1))
        tr3 = abs(low - close.shift(1))

        true_range = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
        df['atr'] = true_range.rolling(period).mean()

        return df

    def compute_combined_weights(
        self,
        df: pd.DataFrame,
        movement_weights: np.ndarray = None,
        atr_col: str = 'atr',
        normalize: bool = True
    ) -> np.ndarray:
        """
        Compute combined weights from ATR (and optionally session/movement).

        By default only uses ATR volatility weighting (session weighting is disabled).
        Final weight = session_weight * atr_weight * movement_weight

        Args:
            df: DataFrame with required columns
            movement_weights: Optional pre-computed movement-based weights
            atr_col: Name of ATR column
            normalize: Whether to normalize weights to mean=1

        Returns:
            Array of combined weights
        """
        n = len(df)

        # Session weights (only if enabled - disabled by default)
        if self.config.use_session_weighting:
            session_weights = self.compute_session_weights(df)
        else:
            session_weights = np.ones(n)
            logger.info("Session weighting disabled, using uniform session weights")

        # ATR weights (enabled by default)
        if self.config.use_atr_weighting:
            atr_weights = self.compute_atr_weights(df, atr_col)
        else:
            atr_weights = np.ones(n)

        # Combine multiplicatively
        combined = session_weights * atr_weights

        # Add movement weights if provided
        if movement_weights is not None:
            combined = combined * movement_weights

        # Normalize
        if normalize and combined.mean() > 0:
            combined = combined / combined.mean()

        logger.info(f"Combined weights: min={combined.min():.3f}, max={combined.max():.3f}, mean={combined.mean():.3f}")

        return combined


def create_session_features(df: pd.DataFrame) -> pd.DataFrame:
    """
    Create additional features for session-based attention.

    These features help the model learn session-specific patterns.

    Args:
        df: DataFrame with datetime index

    Returns:
        DataFrame with additional session features
    """
    df = df.copy()

    if not isinstance(df.index, pd.DatetimeIndex):
        logger.warning("DataFrame does not have DatetimeIndex, skipping session features")
        return df

    hours = df.index.hour
    dow = df.index.dayofweek

    # Binary session indicators
    df['is_london_session'] = ((hours >= 8) & (hours < 16)).astype(int)
    df['is_ny_session'] = ((hours >= 13) & (hours < 21)).astype(int)
    df['is_tokyo_session'] = ((hours >= 0) & (hours < 8)).astype(int)
    df['is_overlap_session'] = ((hours >= 13) & (hours < 16)).astype(int)

    # Cyclical encoding of hour (preserves circular nature)
    df['hour_sin'] = np.sin(2 * np.pi * hours / 24)
    df['hour_cos'] = np.cos(2 * np.pi * hours / 24)

    # Day of week (some days have different patterns)
    df['dow_sin'] = np.sin(2 * np.pi * dow / 5)  # 5-day trading week
    df['dow_cos'] = np.cos(2 * np.pi * dow / 5)

    # Session progression (how far into each session)
    df['london_progress'] = np.where(
        df['is_london_session'] == 1,
        (hours - 8) / 8,  # 0 at start, 1 at end
        0
    )
    df['ny_progress'] = np.where(
        df['is_ny_session'] == 1,
        (hours - 13) / 8,
        0
    )

    # Is it a high-activity day? (Tue-Thu typically have more volume)
    df['is_midweek'] = ((dow >= 1) & (dow <= 3)).astype(int)

    # Monday/Friday special handling (often less predictable)
    df['is_monday'] = (dow == 0).astype(int)
    df['is_friday'] = (dow == 4).astype(int)

    logger.info(f"Created {10} session features")

    return df


def identify_lateral_periods(
    df: pd.DataFrame,
    atr_col: str = 'atr',
    lookback: int = 20,
    threshold_percentile: float = 30
) -> np.ndarray:
    """
    Identify lateral/ranging periods based on ATR.

    Args:
        df: DataFrame with ATR column
        atr_col: Name of ATR column
        lookback: Lookback for percentile calculation
        threshold_percentile: Below this percentile = lateral

    Returns:
        Boolean mask where True = lateral period
    """
    if atr_col not in df.columns:
        logger.warning(f"ATR column '{atr_col}' not found")
        return np.zeros(len(df), dtype=bool)

    atr = df[atr_col].values
    n = len(df)
    is_lateral = np.zeros(n, dtype=bool)

    for i in range(lookback, n):
        window = atr[max(0, i-lookback):i]
        window = window[~np.isnan(window)]

        if len(window) == 0:
            continue

        percentile = (window < atr[i]).mean() * 100

        if percentile <= threshold_percentile:
            is_lateral[i] = True

    logger.info(f"Identified {is_lateral.sum()} lateral periods ({100*is_lateral.mean():.1f}%)")

    return is_lateral


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

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

    # Create datetime index spanning multiple days
    dates = pd.date_range('2025-01-06', periods=n, freq='15min')
    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 with default config (only ATR weighting, no session weights)
    config = SessionWeightConfig(
        use_session_weighting=False,  # Disabled by default
        use_atr_weighting=True        # Enabled by default
    )

    weighter = SessionVolatilityWeighter(config)

    # Add session features
    df = create_session_features(df)

    # Compute weights
    combined_weights = weighter.compute_combined_weights(df)

    print(f"\nResults:")
    print(f"  Total samples: {len(df)}")
    print(f"  Weight range: [{combined_weights.min():.2f}, {combined_weights.max():.2f}]")
    print(f"  Mean weight: {combined_weights.mean():.2f}")

    # Show sample of session features
    print(f"\nSession features added: {[c for c in df.columns if 'session' in c or 'sin' in c or 'cos' in c]}")
