"""
Range Predictor V2 - Multi-Timeframe Extension
===============================================
Predicts price ranges (high/low) across multiple timeframes for ML-First strategy.

Supports timeframes: 5m, 15m, 1H, 4H, D, W
Target: 80% win rate with proper TP/SL placement

Author: ML-Specialist (NEXUS v4.0)
Created: 2026-01-04
"""

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

try:
    from xgboost import XGBRegressor, XGBClassifier
    HAS_XGBOOST = True
except ImportError:
    HAS_XGBOOST = False
    logger.warning("XGBoost not available")

from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.metrics import accuracy_score, f1_score, classification_report
from sklearn.model_selection import TimeSeriesSplit


@dataclass
class TimeframeConfig:
    """Configuration for a specific timeframe"""
    name: str                          # e.g., "5m", "1H", "D"
    horizons: Dict[str, int]           # e.g., {"scalping": 6, "intraday": 2}
    min_samples: int = 5000            # Minimum samples required
    feature_windows: List[int] = field(default_factory=lambda: [12, 24, 48])
    target_multiplier: float = 1.0     # Scale targets by timeframe


@dataclass
class RangePredictionV2:
    """Enhanced range prediction with multi-timeframe support"""
    timeframe: str
    horizon: str
    delta_high: float
    delta_low: float
    confidence: float = 0.0
    direction_prob: float = 0.5        # Probability of upward movement
    expected_range: float = 0.0        # Expected total range
    risk_reward_long: float = 0.0      # R:R for long position
    risk_reward_short: float = 0.0     # R:R for short position
    timestamp: Optional[datetime] = None

    def to_dict(self) -> Dict:
        return {
            'timeframe': self.timeframe,
            'horizon': self.horizon,
            'delta_high': float(self.delta_high),
            'delta_low': float(self.delta_low),
            'confidence': float(self.confidence),
            'direction_prob': float(self.direction_prob),
            'expected_range': float(self.expected_range),
            'risk_reward_long': float(self.risk_reward_long),
            'risk_reward_short': float(self.risk_reward_short),
            'timestamp': self.timestamp.isoformat() if self.timestamp else None
        }

    @property
    def suggested_direction(self) -> str:
        """Suggest trading direction based on prediction"""
        if self.direction_prob > 0.6 and self.risk_reward_long > 1.5:
            return "LONG"
        elif self.direction_prob < 0.4 and self.risk_reward_short > 1.5:
            return "SHORT"
        return "NEUTRAL"


@dataclass
class RangeMetricsV2:
    """Enhanced metrics for range prediction"""
    timeframe: str
    horizon: str
    target_type: str  # 'high', 'low', 'direction'

    # Regression metrics
    mae: float = 0.0
    mape: float = 0.0
    rmse: float = 0.0
    r2: float = 0.0

    # Directional accuracy
    directional_accuracy: float = 0.0

    # Trading-specific metrics
    profitable_predictions: float = 0.0  # % of predictions that would be profitable
    avg_edge: float = 0.0                # Average edge per prediction

    # Sample info
    n_samples: int = 0
    date_range: str = ""


class RangePredictorV2:
    """
    Multi-Timeframe Range Predictor for ML-First Trading Strategy

    Key Features:
    - Supports 6 timeframes: 5m, 15m, 1H, 4H, D, W
    - Multiple horizons per timeframe (scalping, intraday, swing, position)
    - GPU-accelerated XGBoost training
    - Walk-forward validation with OOS testing
    - Designed for 80% win rate objective
    """

    # Default timeframe configurations
    TIMEFRAME_CONFIGS = {
        '5m': TimeframeConfig(
            name='5m',
            horizons={'scalping': 6},  # 30 min
            min_samples=10000,
            feature_windows=[12, 48, 96]  # 1h, 4h, 8h
        ),
        '15m': TimeframeConfig(
            name='15m',
            horizons={'scalping': 4, 'intraday': 8},  # 1h, 2h
            min_samples=8000,
            feature_windows=[4, 16, 32]  # 1h, 4h, 8h
        ),
        '1H': TimeframeConfig(
            name='1H',
            horizons={'intraday': 4, 'swing': 8},  # 4h, 8h
            min_samples=5000,
            feature_windows=[4, 12, 24]  # 4h, 12h, 24h
        ),
        '4H': TimeframeConfig(
            name='4H',
            horizons={'swing': 6, 'position': 12},  # 1d, 2d
            min_samples=2000,
            feature_windows=[6, 12, 24]  # 1d, 2d, 4d
        ),
        'D': TimeframeConfig(
            name='D',
            horizons={'position': 5, 'weekly': 10},  # 1w, 2w
            min_samples=500,
            feature_windows=[5, 10, 20]  # 1w, 2w, 1m
        ),
        'W': TimeframeConfig(
            name='W',
            horizons={'weekly': 4},  # 1m
            min_samples=100,
            feature_windows=[4, 8, 12]  # 1m, 2m, 3m
        )
    }

    def __init__(
        self,
        timeframes: Optional[List[str]] = None,
        config: Optional[Dict[str, Any]] = None,
        use_gpu: bool = True
    ):
        """
        Initialize multi-timeframe range predictor.

        Args:
            timeframes: List of timeframes to support (default: all)
            config: Custom configuration
            use_gpu: Whether to use GPU acceleration
        """
        self.timeframes = timeframes or list(self.TIMEFRAME_CONFIGS.keys())
        self.config = config or self._default_config()
        self.use_gpu = use_gpu

        self.models: Dict[str, Any] = {}
        self.metrics: Dict[str, RangeMetricsV2] = {}
        self.feature_importance: Dict[str, Dict[str, float]] = {}
        self.training_history: List[Dict] = []
        self._is_trained = False

        # Initialize models for each timeframe
        self._init_models()

        logger.info(
            f"Initialized RangePredictorV2 with {len(self.timeframes)} timeframes: "
            f"{', '.join(self.timeframes)}"
        )

    def _default_config(self) -> Dict[str, Any]:
        """Default XGBoost configuration optimized for trading"""
        return {
            'xgboost': {
                'n_estimators': 300,
                'max_depth': 6,
                'learning_rate': 0.03,
                'subsample': 0.8,
                'colsample_bytree': 0.8,
                'min_child_weight': 5,
                'gamma': 0.1,
                'reg_alpha': 0.1,
                'reg_lambda': 1.0,
                'tree_method': 'hist',
                'random_state': 42,
                'early_stopping_rounds': 50
            },
            'training': {
                'val_size': 0.15,
                'walk_forward_splits': 5,
                'min_train_samples': 1000
            },
            'prediction': {
                'ensemble_method': 'weighted_average',
                'confidence_threshold': 0.6
            }
        }

    def _init_models(self):
        """Initialize XGBoost models for each timeframe and horizon"""
        if not HAS_XGBOOST:
            raise ImportError("XGBoost is required for RangePredictorV2")

        xgb_params = self.config['xgboost'].copy()

        # GPU configuration
        if self.use_gpu:
            try:
                import torch
                if torch.cuda.is_available():
                    xgb_params['device'] = 'cuda'
                    xgb_params['tree_method'] = 'gpu_hist'
                    logger.info("GPU acceleration enabled for XGBoost")
            except ImportError:
                pass

        # Remove early_stopping from base params (used in fit)
        early_stopping = xgb_params.pop('early_stopping_rounds', 50)
        self.early_stopping = early_stopping

        for tf in self.timeframes:
            tf_config = self.TIMEFRAME_CONFIGS.get(tf)
            if tf_config is None:
                logger.warning(f"Unknown timeframe: {tf}")
                continue

            for horizon_name in tf_config.horizons.keys():
                # Regression models for delta_high and delta_low
                for target in ['high', 'low']:
                    model_key = f"{tf}_{horizon_name}_{target}"
                    self.models[model_key] = XGBRegressor(**xgb_params)

                # Classification model for direction
                dir_params = xgb_params.copy()
                dir_params['objective'] = 'binary:logistic'
                dir_key = f"{tf}_{horizon_name}_direction"
                self.models[dir_key] = XGBClassifier(**dir_params)

        total_models = len(self.models)
        logger.info(f"Initialized {total_models} models across {len(self.timeframes)} timeframes")

    def get_target_columns(self, timeframe: str) -> Dict[str, List[str]]:
        """Get target column names for a timeframe"""
        tf_config = self.TIMEFRAME_CONFIGS.get(timeframe)
        if tf_config is None:
            return {}

        targets = {}
        for horizon_name in tf_config.horizons.keys():
            targets[horizon_name] = [
                f"target_delta_high_{horizon_name}",
                f"target_delta_low_{horizon_name}",
                f"target_direction_{horizon_name}"
            ]

        return targets

    def train(
        self,
        X_train: pd.DataFrame,
        y_train: pd.DataFrame,
        X_val: Optional[pd.DataFrame] = None,
        y_val: Optional[pd.DataFrame] = None,
        timeframe: Optional[str] = None,
        verbose: bool = True
    ) -> Dict[str, RangeMetricsV2]:
        """
        Train models for specified timeframe(s).

        Args:
            X_train: Training features
            y_train: Training targets
            X_val: Validation features (optional)
            y_val: Validation targets (optional)
            timeframe: Specific timeframe to train (None = all)
            verbose: Print training progress

        Returns:
            Dictionary of metrics for trained models
        """
        timeframes_to_train = [timeframe] if timeframe else self.timeframes

        all_metrics = {}

        for tf in timeframes_to_train:
            tf_config = self.TIMEFRAME_CONFIGS.get(tf)
            if tf_config is None:
                continue

            if verbose:
                logger.info(f"\n{'='*60}")
                logger.info(f"Training models for timeframe: {tf}")
                logger.info(f"{'='*60}")

            target_cols = self.get_target_columns(tf)

            for horizon_name, cols in target_cols.items():
                if verbose:
                    logger.info(f"  Horizon: {horizon_name}")

                # Train high predictor
                high_col = cols[0]  # target_delta_high_{horizon}
                if high_col in y_train.columns:
                    metrics = self._train_single_model(
                        X_train, y_train[high_col],
                        X_val, y_val[high_col] if y_val is not None and high_col in y_val.columns else None,
                        f"{tf}_{horizon_name}_high",
                        tf, horizon_name, 'high'
                    )
                    all_metrics[f"{tf}_{horizon_name}_high"] = metrics

                # Train low predictor
                low_col = cols[1]  # target_delta_low_{horizon}
                if low_col in y_train.columns:
                    metrics = self._train_single_model(
                        X_train, y_train[low_col],
                        X_val, y_val[low_col] if y_val is not None and low_col in y_val.columns else None,
                        f"{tf}_{horizon_name}_low",
                        tf, horizon_name, 'low'
                    )
                    all_metrics[f"{tf}_{horizon_name}_low"] = metrics

                # Train direction classifier
                dir_col = cols[2]  # target_direction_{horizon}
                if dir_col in y_train.columns:
                    metrics = self._train_single_model(
                        X_train, y_train[dir_col],
                        X_val, y_val[dir_col] if y_val is not None and dir_col in y_val.columns else None,
                        f"{tf}_{horizon_name}_direction",
                        tf, horizon_name, 'direction',
                        is_classifier=True
                    )
                    all_metrics[f"{tf}_{horizon_name}_direction"] = metrics

        self._is_trained = True
        self.metrics = all_metrics

        # Log summary
        if verbose:
            self._print_training_summary(all_metrics)

        return all_metrics

    def _train_single_model(
        self,
        X_train: pd.DataFrame,
        y_train: pd.Series,
        X_val: Optional[pd.DataFrame],
        y_val: Optional[pd.Series],
        model_key: str,
        timeframe: str,
        horizon: str,
        target_type: str,
        is_classifier: bool = False
    ) -> RangeMetricsV2:
        """Train a single model"""
        # Get model
        model = self.models.get(model_key)
        if model is None:
            logger.warning(f"Model not found: {model_key}")
            return RangeMetricsV2(timeframe=timeframe, horizon=horizon, target_type=target_type)

        # Prepare data
        X_train_np = X_train.values
        y_train_np = y_train.values

        # Handle NaN
        valid_mask = ~np.isnan(y_train_np)
        X_train_clean = X_train_np[valid_mask]
        y_train_clean = y_train_np[valid_mask]

        if len(X_train_clean) < 100:
            logger.warning(f"Insufficient samples for {model_key}: {len(X_train_clean)}")
            return RangeMetricsV2(timeframe=timeframe, horizon=horizon, target_type=target_type)

        # Prepare validation set
        fit_params = {}
        if X_val is not None and y_val is not None:
            X_val_np = X_val.values
            y_val_np = y_val.values
            valid_mask_val = ~np.isnan(y_val_np)
            if valid_mask_val.sum() > 0:
                fit_params['eval_set'] = [(X_val_np[valid_mask_val], y_val_np[valid_mask_val])]

        # Train
        try:
            if is_classifier:
                y_train_clean = y_train_clean.astype(int)
            model.fit(X_train_clean, y_train_clean, **fit_params)
        except Exception as e:
            logger.error(f"Training failed for {model_key}: {e}")
            return RangeMetricsV2(timeframe=timeframe, horizon=horizon, target_type=target_type)

        # Store feature importance
        if isinstance(X_train, pd.DataFrame):
            self.feature_importance[model_key] = dict(
                zip(X_train.columns, model.feature_importances_)
            )

        # Calculate metrics
        y_pred = model.predict(X_train_clean)

        if is_classifier:
            metrics = RangeMetricsV2(
                timeframe=timeframe,
                horizon=horizon,
                target_type=target_type,
                directional_accuracy=accuracy_score(y_train_clean, y_pred.round()),
                n_samples=len(X_train_clean)
            )
        else:
            metrics = RangeMetricsV2(
                timeframe=timeframe,
                horizon=horizon,
                target_type=target_type,
                mae=mean_absolute_error(y_train_clean, y_pred),
                rmse=np.sqrt(mean_squared_error(y_train_clean, y_pred)),
                r2=r2_score(y_train_clean, y_pred),
                n_samples=len(X_train_clean)
            )

            # Calculate directional accuracy for regression
            if target_type in ['high', 'low']:
                direction_correct = (np.sign(y_train_clean) == np.sign(y_pred)).mean()
                metrics.directional_accuracy = direction_correct

        logger.info(
            f"    {target_type}: MAE={metrics.mae:.4f}, R2={metrics.r2:.4f}, "
            f"DirAcc={metrics.directional_accuracy:.2%}"
        )

        return metrics

    def predict(
        self,
        X: pd.DataFrame,
        timeframe: Optional[str] = None
    ) -> List[RangePredictionV2]:
        """
        Generate predictions for given features.

        Args:
            X: Features DataFrame
            timeframe: Specific timeframe (None = all trained)

        Returns:
            List of RangePredictionV2 objects
        """
        if not self._is_trained:
            raise RuntimeError("Model must be trained before prediction")

        timeframes_to_predict = [timeframe] if timeframe else self.timeframes
        predictions = []

        X_np = X.values if isinstance(X, pd.DataFrame) else X
        if X_np.ndim == 1:
            X_np = X_np.reshape(1, -1)

        for tf in timeframes_to_predict:
            tf_config = self.TIMEFRAME_CONFIGS.get(tf)
            if tf_config is None:
                continue

            for horizon_name in tf_config.horizons.keys():
                # Get predictions from each model
                high_key = f"{tf}_{horizon_name}_high"
                low_key = f"{tf}_{horizon_name}_low"
                dir_key = f"{tf}_{horizon_name}_direction"

                if high_key not in self.models or low_key not in self.models:
                    continue

                delta_high = self.models[high_key].predict(X_np)
                delta_low = self.models[low_key].predict(X_np)

                # Direction probability
                if dir_key in self.models:
                    try:
                        dir_prob = self.models[dir_key].predict_proba(X_np)[:, 1]
                    except:
                        dir_prob = np.full(len(X_np), 0.5)
                else:
                    dir_prob = np.full(len(X_np), 0.5)

                # Create predictions for each sample
                for i in range(len(X_np)):
                    dh = float(delta_high[i])
                    dl = float(delta_low[i])
                    dp = float(dir_prob[i])

                    # Calculate risk/reward ratios
                    expected_range = abs(dh) + abs(dl)
                    rr_long = abs(dh) / abs(dl) if abs(dl) > 0 else 0
                    rr_short = abs(dl) / abs(dh) if abs(dh) > 0 else 0

                    # Confidence based on model agreement
                    confidence = self._calculate_confidence(dh, dl, dp)

                    pred = RangePredictionV2(
                        timeframe=tf,
                        horizon=horizon_name,
                        delta_high=dh,
                        delta_low=dl,
                        confidence=confidence,
                        direction_prob=dp,
                        expected_range=expected_range,
                        risk_reward_long=rr_long,
                        risk_reward_short=rr_short,
                        timestamp=datetime.now()
                    )
                    predictions.append(pred)

        return predictions

    def _calculate_confidence(
        self,
        delta_high: float,
        delta_low: float,
        direction_prob: float
    ) -> float:
        """Calculate prediction confidence"""
        # Direction consistency
        predicted_direction = delta_high > abs(delta_low)  # More upside = bullish
        direction_agreement = (
            (predicted_direction and direction_prob > 0.5) or
            (not predicted_direction and direction_prob < 0.5)
        )

        # Base confidence from direction probability
        base_conf = abs(direction_prob - 0.5) * 2  # 0-1 scale

        # Boost if predictions agree
        if direction_agreement:
            return min(1.0, base_conf * 1.2)
        else:
            return base_conf * 0.8

    def evaluate(
        self,
        X_test: pd.DataFrame,
        y_test: pd.DataFrame,
        timeframe: Optional[str] = None
    ) -> Dict[str, RangeMetricsV2]:
        """
        Evaluate models on test data.

        Args:
            X_test: Test features
            y_test: Test targets
            timeframe: Specific timeframe (None = all)

        Returns:
            Dictionary of metrics
        """
        timeframes_to_eval = [timeframe] if timeframe else self.timeframes
        all_metrics = {}

        for tf in timeframes_to_eval:
            target_cols = self.get_target_columns(tf)

            for horizon_name, cols in target_cols.items():
                for i, target_type in enumerate(['high', 'low', 'direction']):
                    col = cols[i]
                    if col not in y_test.columns:
                        continue

                    model_key = f"{tf}_{horizon_name}_{target_type}"
                    if model_key not in self.models:
                        continue

                    y_true = y_test[col].values
                    valid_mask = ~np.isnan(y_true)

                    if valid_mask.sum() == 0:
                        continue

                    X_test_valid = X_test.values[valid_mask]
                    y_true_valid = y_true[valid_mask]

                    y_pred = self.models[model_key].predict(X_test_valid)

                    if target_type == 'direction':
                        metrics = RangeMetricsV2(
                            timeframe=tf,
                            horizon=horizon_name,
                            target_type=target_type,
                            directional_accuracy=accuracy_score(y_true_valid.astype(int), y_pred.round()),
                            n_samples=len(X_test_valid)
                        )
                    else:
                        metrics = RangeMetricsV2(
                            timeframe=tf,
                            horizon=horizon_name,
                            target_type=target_type,
                            mae=mean_absolute_error(y_true_valid, y_pred),
                            rmse=np.sqrt(mean_squared_error(y_true_valid, y_pred)),
                            r2=r2_score(y_true_valid, y_pred),
                            directional_accuracy=(np.sign(y_true_valid) == np.sign(y_pred)).mean(),
                            n_samples=len(X_test_valid)
                        )

                    all_metrics[model_key] = metrics

        return all_metrics

    def _print_training_summary(self, metrics: Dict[str, RangeMetricsV2]):
        """Print training summary"""
        logger.info("\n" + "=" * 70)
        logger.info("TRAINING SUMMARY")
        logger.info("=" * 70)

        # Group by timeframe
        by_timeframe: Dict[str, List] = {}
        for key, m in metrics.items():
            tf = m.timeframe
            if tf not in by_timeframe:
                by_timeframe[tf] = []
            by_timeframe[tf].append((key, m))

        for tf, items in by_timeframe.items():
            logger.info(f"\n{tf}:")
            for key, m in items:
                if m.target_type == 'direction':
                    logger.info(f"  {key}: DirAcc={m.directional_accuracy:.2%}")
                else:
                    logger.info(
                        f"  {key}: MAE={m.mae:.4f}, R2={m.r2:.4f}, "
                        f"DirAcc={m.directional_accuracy:.2%}"
                    )

    def get_feature_importance(
        self,
        model_key: Optional[str] = None,
        top_n: int = 20
    ) -> Dict[str, float]:
        """Get feature importance"""
        if model_key:
            importance = self.feature_importance.get(model_key, {})
        else:
            # Average across all models
            all_features = set()
            for fi in self.feature_importance.values():
                all_features.update(fi.keys())

            importance = {}
            for feat in all_features:
                values = [fi.get(feat, 0) for fi in self.feature_importance.values()]
                importance[feat] = np.mean(values)

        return dict(sorted(importance.items(), key=lambda x: x[1], reverse=True)[:top_n])

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

        # Save each model
        for name, model in self.models.items():
            joblib.dump(model, path / f'{name}.joblib')

        # Save metadata
        metadata = {
            'config': self.config,
            'timeframes': self.timeframes,
            'metrics': {k: vars(v) for k, v in self.metrics.items()},
            'feature_importance': self.feature_importance,
            'training_history': self.training_history,
            'version': 'v2.0'
        }
        joblib.dump(metadata, path / 'metadata.joblib')

        logger.info(f"Saved RangePredictorV2 to {path}")

    def load(self, path: str):
        """Load models from disk"""
        path = Path(path)

        # Load metadata
        metadata = joblib.load(path / 'metadata.joblib')
        self.config = metadata['config']
        self.timeframes = metadata['timeframes']
        self.feature_importance = metadata['feature_importance']

        # Load models
        self.models = {}
        for model_file in path.glob('*.joblib'):
            if model_file.name != 'metadata.joblib':
                name = model_file.stem
                self.models[name] = joblib.load(model_file)

        self._is_trained = True
        logger.info(f"Loaded RangePredictorV2 from {path}")


if __name__ == "__main__":
    # Test RangePredictorV2
    np.random.seed(42)

    # Create sample data
    n_samples = 2000
    n_features = 30

    X = pd.DataFrame(
        np.random.randn(n_samples, n_features),
        columns=[f'feature_{i}' for i in range(n_features)]
    )

    # Create targets for 15m timeframe
    y = pd.DataFrame({
        'target_delta_high_scalping': np.random.randn(n_samples) * 0.01,
        'target_delta_low_scalping': np.random.randn(n_samples) * 0.01,
        'target_direction_scalping': (np.random.rand(n_samples) > 0.5).astype(int),
        'target_delta_high_intraday': np.random.randn(n_samples) * 0.02,
        'target_delta_low_intraday': np.random.randn(n_samples) * 0.02,
        'target_direction_intraday': (np.random.rand(n_samples) > 0.5).astype(int),
    })

    # Split data
    train_size = 1600
    X_train, X_val = X.iloc[:train_size], X.iloc[train_size:]
    y_train, y_val = y.iloc[:train_size], y.iloc[train_size:]

    # Train
    predictor = RangePredictorV2(timeframes=['15m'], use_gpu=False)
    metrics = predictor.train(X_train, y_train, X_val, y_val)

    # Evaluate
    test_metrics = predictor.evaluate(X_val, y_val)
    print("\n=== Test Metrics ===")
    for key, m in test_metrics.items():
        print(f"{key}: MAE={m.mae:.4f}, DirAcc={m.directional_accuracy:.2%}")

    # Predict
    predictions = predictor.predict(X_val.iloc[:5], timeframe='15m')
    print("\n=== Sample Predictions ===")
    for pred in predictions:
        print(f"{pred.timeframe}/{pred.horizon}: "
              f"High={pred.delta_high:.4f}, Low={pred.delta_low:.4f}, "
              f"Direction={pred.suggested_direction}")
