#!/usr/bin/env python3
"""
Enhanced Range Predictor
========================
Integrates all improvements for volatility-factor based range prediction.

Components Integrated:
1. Corrected targets (MAX/MIN formula)
2. Sample weighting (movement + session + volatility)
3. Dual horizon ensemble (5 years + 3 months)
4. R:R ratio filtering (2:1 minimum)

Key Features:
- Predictions in multipliers of base factor (e.g., 5 USD for XAUUSD)
- Focus on large movements only
- Higher weight for high-volatility sessions (London/NY overlap)
- Lower weight for lateral/ranging markets
- Automatic retraining of short-term model

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

import numpy as np
import pandas as pd
from typing import Dict, Tuple, Optional, List, Any
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from pathlib import Path
import joblib
from loguru import logger

# Local imports
from ..data.corrected_targets import (
    CorrectedTargetBuilder,
    CorrectedTargetConfig,
    TargetResult,
    calculate_volatility_factor
)
from ..training.sample_weighting import (
    SampleWeighter,
    SampleWeightConfig
)
from ..training.session_volatility_weighting import (
    SessionVolatilityWeighter,
    SessionWeightConfig,
    create_session_features
)
from .dual_horizon_ensemble import (
    DualHorizonEnsemble,
    DualHorizonConfig
)


@dataclass
class EnhancedRangePredictorConfig:
    """Master configuration for enhanced range predictor"""

    # Symbol and base parameters
    symbol: str = 'XAUUSD'
    base_factor: float = 5.0  # USD for gold

    # Timeframe and horizon
    input_timeframe: str = '15m'
    prediction_horizon_bars: int = 3  # 3 bars of 15m = 45 minutes

    # Target configuration
    target_config: CorrectedTargetConfig = field(default_factory=lambda: CorrectedTargetConfig(
        horizon_bars=3,
        start_offset=1,
        min_movement_usd=5.0,
        normalize_by_atr=True,
        min_rr_ratio=2.0,
        base_factor=5.0
    ))

    # Sample weighting configuration
    sample_weight_config: SampleWeightConfig = field(default_factory=lambda: SampleWeightConfig(
        min_movement_threshold=5.0,
        large_movement_weight=3.0,
        small_movement_weight=0.3,
        use_continuous_weighting=True,
        weight_exponent=1.5,
        min_rr_ratio=2.0
    ))

    # Session weighting configuration
    session_config: SessionWeightConfig = field(default_factory=lambda: SessionWeightConfig(
        overlap_weight=2.0,
        london_weight=1.5,
        ny_weight=1.3,
        tokyo_weight=0.7,
        off_hours_weight=0.3,
        use_atr_weighting=True,
        atr_high_weight_boost=1.5,
        atr_low_weight_penalty=0.3
    ))

    # Dual horizon configuration
    dual_horizon_config: DualHorizonConfig = field(default_factory=lambda: DualHorizonConfig(
        long_term_years=5.0,
        short_term_months=3.0,
        initial_long_weight=0.6,
        initial_short_weight=0.4,
        use_dynamic_weights=True,
        retrain_frequency_days=7
    ))

    # Prediction output mode
    output_mode: str = 'multiplier'  # 'multiplier' or 'usd'

    # Minimum confidence for predictions
    min_prediction_confidence: float = 0.6

    # Feature columns to use (if None, auto-detect)
    feature_columns: List[str] = None


@dataclass
class PredictionResult:
    """Result of a prediction"""
    pred_high_mult: float
    pred_low_mult: float
    pred_high_usd: float
    pred_low_usd: float
    suggested_direction: str  # 'LONG', 'SHORT', or 'NEUTRAL'
    rr_ratio: float
    confidence: float
    model_weights: Dict[str, float]
    timestamp: datetime


class EnhancedRangePredictor:
    """
    Main predictor class integrating all improvements.

    Pipeline:
    1. Calculate corrected targets from OHLCV data
    2. Compute combined sample weights (movement + session + volatility)
    3. Train dual horizon ensemble with weighted samples
    4. Predict range in multipliers of base factor
    5. Filter predictions by R:R ratio

    Usage:
        config = EnhancedRangePredictorConfig(symbol='XAUUSD', base_factor=5.0)
        predictor = EnhancedRangePredictor(config)

        # Train
        predictor.fit(df_ohlcv, feature_df)

        # Predict
        result = predictor.predict_single(features)
        # or
        results_df = predictor.predict_batch(features_df)
    """

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

        # Initialize components
        self.target_builder = CorrectedTargetBuilder(self.config.target_config)
        self.sample_weighter = SampleWeighter(self.config.sample_weight_config)
        self.session_weighter = SessionVolatilityWeighter(self.config.session_config)
        self.ensemble = DualHorizonEnsemble(self.config.dual_horizon_config)

        # State
        self.is_fitted = False
        self.feature_names: List[str] = []
        self.volatility_metrics: Dict = {}
        self.training_stats: Dict = {}

    def prepare_data(
        self,
        df_ohlcv: pd.DataFrame,
        df_features: pd.DataFrame = None
    ) -> Tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]:
        """
        Prepare data for training: targets, features, and weights.

        Args:
            df_ohlcv: OHLCV DataFrame with datetime index
            df_features: Optional pre-computed features DataFrame

        Returns:
            Tuple of (df_prepared, X, y_high, y_low, sample_weights)
        """
        logger.info(f"Preparing data for {self.config.symbol}...")

        # 1. Calculate corrected targets
        target_result = self.target_builder.build_targets(df_ohlcv, self.config.symbol)
        df = self.target_builder.add_targets_to_dataframe(df_ohlcv, target_result, self.config.symbol)

        # 2. Add session features
        df = create_session_features(df)

        # 3. Merge with additional features if provided
        if df_features is not None:
            # Align by index
            df = df.join(df_features, how='inner', rsuffix='_feat')
            logger.info(f"Merged with {len(df_features.columns)} additional features")

        # 4. Calculate sample weights

        # Movement-based weights
        movement_weights, movement_valid = self.sample_weighter.compute_sample_weights(
            df, 'target_high_usd', 'target_low_usd'
        )

        # Session and volatility weights
        session_weights = self.session_weighter.compute_combined_weights(df)

        # Combine weights multiplicatively
        combined_weights = movement_weights * session_weights

        # Normalize
        valid_mask = movement_valid & target_result.is_valid
        if combined_weights[valid_mask].sum() > 0:
            combined_weights[valid_mask] = combined_weights[valid_mask] / combined_weights[valid_mask].mean()

        # 5. Prepare feature matrix
        if self.config.feature_columns:
            feature_cols = [c for c in self.config.feature_columns if c in df.columns]
        else:
            # Auto-detect feature columns (exclude targets and metadata)
            exclude_patterns = [
                'target_', 'rr_', 'is_valid', 'direction', 'high', 'low', 'open', 'close', 'volume',
                'High', 'Low', 'Open', 'Close', 'Volume', 'atr', 'ATR'
            ]
            feature_cols = [
                c for c in df.columns
                if not any(p in c for p in exclude_patterns)
                and df[c].dtype in [np.float64, np.float32, np.int64, np.int32, float, int]
            ]

        self.feature_names = feature_cols
        logger.info(f"Using {len(feature_cols)} features")

        # Extract arrays
        X = df[feature_cols].values
        y_high = df['target_high_mult'].values
        y_low = df['target_low_mult'].values

        # Store stats
        self.training_stats = {
            'total_samples': len(df),
            'valid_samples': valid_mask.sum(),
            'long_opportunities': (target_result.direction == 1).sum(),
            'short_opportunities': (target_result.direction == -1).sum(),
            'feature_count': len(feature_cols),
            'mean_target_high': np.nanmean(y_high[valid_mask]),
            'mean_target_low': np.nanmean(y_low[valid_mask])
        }

        logger.info(f"Data preparation complete:")
        logger.info(f"  Total: {self.training_stats['total_samples']}")
        logger.info(f"  Valid: {self.training_stats['valid_samples']}")
        logger.info(f"  LONG: {self.training_stats['long_opportunities']}")
        logger.info(f"  SHORT: {self.training_stats['short_opportunities']}")

        return df, X, y_high, y_low, combined_weights, valid_mask

    def fit(
        self,
        df_ohlcv: pd.DataFrame,
        df_features: pd.DataFrame = None
    ) -> 'EnhancedRangePredictor':
        """
        Fit the predictor on historical data.

        Args:
            df_ohlcv: OHLCV DataFrame with datetime index
            df_features: Optional pre-computed features DataFrame

        Returns:
            Self for chaining
        """
        # Prepare data
        df, X, y_high, y_low, weights, valid_mask = self.prepare_data(df_ohlcv, df_features)

        # Filter to valid samples only
        X_valid = X[valid_mask]
        y_high_valid = y_high[valid_mask]
        y_low_valid = y_low[valid_mask]
        weights_valid = weights[valid_mask]

        # Get timestamps
        if isinstance(df.index, pd.DatetimeIndex):
            timestamps = df.index[valid_mask]
        else:
            timestamps = pd.to_datetime(df['timestamp'][valid_mask]) if 'timestamp' in df.columns else pd.date_range(
                end=datetime.now(), periods=valid_mask.sum(), freq=self.config.input_timeframe
            )

        # Calculate volatility metrics for reference
        self.volatility_metrics = calculate_volatility_factor(
            df_ohlcv, self.config.input_timeframe, lookback=min(1000, len(df_ohlcv))
        )

        # Fit ensemble
        self.ensemble.fit(
            X_valid,
            y_high_valid,
            y_low_valid,
            timestamps,
            weights_valid,
            self.feature_names
        )

        self.is_fitted = True
        logger.info("EnhancedRangePredictor training complete")

        return self

    def predict_single(
        self,
        features: np.ndarray,
        timestamp: datetime = None
    ) -> PredictionResult:
        """
        Predict range for a single sample.

        Args:
            features: Feature vector (1D array)
            timestamp: Optional timestamp for the prediction

        Returns:
            PredictionResult with predictions and metadata
        """
        if not self.is_fitted:
            raise ValueError("Model not fitted. Call fit() first.")

        # Reshape if needed
        X = features.reshape(1, -1) if features.ndim == 1 else features

        # Get predictions with components
        pred_high_mult, pred_low_mult, components = self.ensemble.predict(X, return_components=True)

        pred_high_mult = pred_high_mult[0]
        pred_low_mult = pred_low_mult[0]

        # Convert to USD
        factor = self.config.base_factor
        pred_high_usd = pred_high_mult * factor
        pred_low_usd = pred_low_mult * factor

        # Calculate R:R and direction
        epsilon = 0.0001
        rr_long = pred_high_mult / (pred_low_mult + epsilon)
        rr_short = pred_low_mult / (pred_high_mult + epsilon)

        if rr_long >= self.config.target_config.min_rr_ratio:
            direction = 'LONG'
            rr_ratio = rr_long
        elif rr_short >= self.config.target_config.min_rr_ratio:
            direction = 'SHORT'
            rr_ratio = rr_short
        else:
            direction = 'NEUTRAL'
            rr_ratio = max(rr_long, rr_short)

        # Calculate confidence based on model agreement
        long_agreement = 1 - abs(
            components['long_term_high'][0] - components['short_term_high'][0]
        ) / (pred_high_mult + epsilon)
        low_agreement = 1 - abs(
            components['long_term_low'][0] - components['short_term_low'][0]
        ) / (pred_low_mult + epsilon)
        confidence = (long_agreement + low_agreement) / 2
        confidence = np.clip(confidence, 0, 1)

        return PredictionResult(
            pred_high_mult=float(pred_high_mult),
            pred_low_mult=float(pred_low_mult),
            pred_high_usd=float(pred_high_usd),
            pred_low_usd=float(pred_low_usd),
            suggested_direction=direction,
            rr_ratio=float(rr_ratio),
            confidence=float(confidence),
            model_weights={
                'long_term': components['weight_long'],
                'short_term': components['weight_short']
            },
            timestamp=timestamp or datetime.now()
        )

    def predict_batch(
        self,
        features: np.ndarray,
        timestamps: pd.DatetimeIndex = None
    ) -> pd.DataFrame:
        """
        Predict range for multiple samples.

        Args:
            features: Feature matrix (2D array)
            timestamps: Optional timestamps for predictions

        Returns:
            DataFrame with predictions for each sample
        """
        if not self.is_fitted:
            raise ValueError("Model not fitted. Call fit() first.")

        # Get predictions
        pred_high_mult, pred_low_mult, components = self.ensemble.predict(features, return_components=True)

        # Convert to USD
        factor = self.config.base_factor
        pred_high_usd = pred_high_mult * factor
        pred_low_usd = pred_low_mult * factor

        # Calculate R:R ratios
        epsilon = 0.0001
        rr_long = pred_high_mult / (pred_low_mult + epsilon)
        rr_short = pred_low_mult / (pred_high_mult + epsilon)
        rr_best = np.maximum(rr_long, rr_short)

        # Determine directions
        min_rr = self.config.target_config.min_rr_ratio
        direction = np.where(
            rr_long >= min_rr, 'LONG',
            np.where(rr_short >= min_rr, 'SHORT', 'NEUTRAL')
        )

        # Calculate confidence
        high_diff = np.abs(components['long_term_high'] - components['short_term_high'])
        low_diff = np.abs(components['long_term_low'] - components['short_term_low'])
        confidence = 1 - (high_diff + low_diff) / (pred_high_mult + pred_low_mult + epsilon)
        confidence = np.clip(confidence, 0, 1)

        # Build result DataFrame
        results = pd.DataFrame({
            'pred_high_mult': pred_high_mult,
            'pred_low_mult': pred_low_mult,
            'pred_high_usd': pred_high_usd,
            'pred_low_usd': pred_low_usd,
            'pred_total_mult': pred_high_mult + pred_low_mult,
            'pred_total_usd': pred_high_usd + pred_low_usd,
            'rr_long': rr_long,
            'rr_short': rr_short,
            'rr_best': rr_best,
            'direction': direction,
            'confidence': confidence,
            'long_term_high': components['long_term_high'],
            'long_term_low': components['long_term_low'],
            'short_term_high': components['short_term_high'],
            'short_term_low': components['short_term_low']
        })

        if timestamps is not None:
            results.index = timestamps

        # Filter by confidence if configured
        if self.config.min_prediction_confidence > 0:
            high_conf_mask = results['confidence'] >= self.config.min_prediction_confidence
            logger.info(f"High confidence predictions: {high_conf_mask.sum()} / {len(results)}")

        return results

    def get_trading_signal(
        self,
        features: np.ndarray,
        current_price: float
    ) -> Dict[str, Any]:
        """
        Get a trading signal with entry, TP, and SL levels.

        Args:
            features: Feature vector for prediction
            current_price: Current close price

        Returns:
            Dict with signal details
        """
        pred = self.predict_single(features)

        if pred.suggested_direction == 'NEUTRAL':
            return {
                'action': 'WAIT',
                'reason': f'R:R ratio {pred.rr_ratio:.2f} below threshold',
                'prediction': pred
            }

        if pred.confidence < self.config.min_prediction_confidence:
            return {
                'action': 'WAIT',
                'reason': f'Confidence {pred.confidence:.2f} below threshold',
                'prediction': pred
            }

        # Calculate entry, TP, SL
        if pred.suggested_direction == 'LONG':
            entry = current_price
            tp = current_price + pred.pred_high_usd
            sl = current_price - pred.pred_low_usd
        else:  # SHORT
            entry = current_price
            tp = current_price - pred.pred_low_usd
            sl = current_price + pred.pred_high_usd

        return {
            'action': pred.suggested_direction,
            'entry': entry,
            'take_profit': tp,
            'stop_loss': sl,
            'rr_ratio': pred.rr_ratio,
            'confidence': pred.confidence,
            'pred_high_usd': pred.pred_high_usd,
            'pred_low_usd': pred.pred_low_usd,
            'model_weights': pred.model_weights,
            'prediction': pred
        }

    def update_with_result(
        self,
        features: np.ndarray,
        actual_high: float,
        actual_low: float
    ) -> None:
        """
        Update model performance tracking with actual results.

        Args:
            features: Feature vector used for prediction
            actual_high: Actual high target achieved
            actual_low: Actual low target achieved
        """
        X = features.reshape(1, -1) if features.ndim == 1 else features

        # Update ensemble performance
        self.ensemble.update_performance(
            X,
            np.array([actual_high]),
            np.array([actual_low])
        )

        # Check if weights should be adjusted
        self.ensemble.adjust_weights()

    def should_retrain(self) -> bool:
        """Check if short-term model should be retrained."""
        return self.ensemble.should_retrain()

    def retrain_short_term(
        self,
        df_ohlcv: pd.DataFrame,
        df_features: pd.DataFrame = None
    ) -> None:
        """
        Retrain short-term model with recent data.

        Args:
            df_ohlcv: Recent OHLCV data (last 3 months recommended)
            df_features: Optional features DataFrame
        """
        # Prepare recent data
        df, X, y_high, y_low, weights, valid_mask = self.prepare_data(df_ohlcv, df_features)

        # Filter to valid samples
        X_valid = X[valid_mask]
        y_high_valid = y_high[valid_mask]
        y_low_valid = y_low[valid_mask]
        weights_valid = weights[valid_mask]

        # Retrain short-term
        self.ensemble.retrain_short_term(
            X_valid, y_high_valid, y_low_valid, weights_valid
        )

        logger.info("Short-term model retrained with recent data")

    def get_feature_importance(self) -> pd.DataFrame:
        """Get feature importance from ensemble."""
        return self.ensemble.get_feature_importance('combined')

    def get_model_summary(self) -> Dict[str, Any]:
        """Get comprehensive model summary."""
        return {
            'config': {
                'symbol': self.config.symbol,
                'base_factor': self.config.base_factor,
                'input_timeframe': self.config.input_timeframe,
                'prediction_horizon_bars': self.config.prediction_horizon_bars
            },
            'training_stats': self.training_stats,
            'volatility_metrics': self.volatility_metrics,
            'ensemble_summary': self.ensemble.get_model_summary(),
            'is_fitted': self.is_fitted,
            'feature_count': len(self.feature_names)
        }

    def save(self, path: str) -> None:
        """Save predictor to disk."""
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)

        # Save ensemble
        self.ensemble.save(path / 'ensemble')

        # Save metadata
        metadata = {
            'config': self.config,
            'feature_names': self.feature_names,
            'volatility_metrics': self.volatility_metrics,
            'training_stats': self.training_stats,
            'is_fitted': self.is_fitted
        }
        joblib.dump(metadata, path / 'predictor_metadata.joblib')

        logger.info(f"EnhancedRangePredictor saved to {path}")

    @classmethod
    def load(cls, path: str) -> 'EnhancedRangePredictor':
        """Load predictor from disk."""
        path = Path(path)

        # Load metadata
        metadata = joblib.load(path / 'predictor_metadata.joblib')

        # Create instance
        predictor = cls(metadata['config'])
        predictor.feature_names = metadata['feature_names']
        predictor.volatility_metrics = metadata['volatility_metrics']
        predictor.training_stats = metadata['training_stats']
        predictor.is_fitted = metadata['is_fitted']

        # Load ensemble
        predictor.ensemble = DualHorizonEnsemble.load(path / 'ensemble')

        logger.info(f"EnhancedRangePredictor loaded from {path}")
        return predictor


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

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

    # Create datetime index
    dates = pd.date_range('2020-01-01', periods=n, freq='15min')
    price = 2650 + np.cumsum(np.random.randn(n) * 0.5)

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

    # Create simple feature DataFrame
    df_features = pd.DataFrame({
        'rsi': 50 + np.random.randn(n) * 10,
        'macd': np.random.randn(n) * 2,
        'bb_width': 10 + np.random.randn(n),
        'momentum': np.random.randn(n) * 5
    }, index=dates)

    # Test predictor
    config = EnhancedRangePredictorConfig(
        symbol='XAUUSD',
        base_factor=5.0,
        input_timeframe='15m',
        prediction_horizon_bars=3
    )

    predictor = EnhancedRangePredictor(config)

    # Fit
    predictor.fit(df_ohlcv, df_features)

    # Single prediction
    test_features = df_features.iloc[-1].values
    result = predictor.predict_single(test_features)

    print(f"\nSingle prediction:")
    print(f"  High (mult): {result.pred_high_mult:.2f} = {result.pred_high_usd:.2f} USD")
    print(f"  Low (mult): {result.pred_low_mult:.2f} = {result.pred_low_usd:.2f} USD")
    print(f"  Direction: {result.suggested_direction}")
    print(f"  R:R ratio: {result.rr_ratio:.2f}")
    print(f"  Confidence: {result.confidence:.2f}")

    # Batch prediction
    test_batch = df_features.iloc[-100:].values
    results_df = predictor.predict_batch(test_batch, dates[-100:])

    print(f"\nBatch prediction summary ({len(results_df)} samples):")
    print(f"  LONG signals: {(results_df['direction'] == 'LONG').sum()}")
    print(f"  SHORT signals: {(results_df['direction'] == 'SHORT').sum()}")
    print(f"  NEUTRAL: {(results_df['direction'] == 'NEUTRAL').sum()}")
    print(f"  Mean confidence: {results_df['confidence'].mean():.2f}")

    # Trading signal
    signal = predictor.get_trading_signal(test_features, df_ohlcv['close'].iloc[-1])
    print(f"\nTrading signal:")
    print(f"  Action: {signal['action']}")
    if signal['action'] != 'WAIT':
        print(f"  Entry: {signal['entry']:.2f}")
        print(f"  TP: {signal['take_profit']:.2f}")
        print(f"  SL: {signal['stop_loss']:.2f}")

    # Model summary
    summary = predictor.get_model_summary()
    print(f"\nModel summary:")
    print(f"  Training samples: {summary['training_stats']['total_samples']}")
    print(f"  Valid samples: {summary['training_stats']['valid_samples']}")
    print(f"  Features: {summary['feature_count']}")
