#!/usr/bin/env python3
"""
Range Predictor with Factor-Based Approach

Instead of predicting absolute ΔHigh/ΔLow values, this model:
1. Calculates the "base volatility factor" for each asset
2. Predicts MULTIPLIERS of this factor

Example for XAUUSD:
- Base factor = 2.5 USD (normal 5-min range)
- Model predicts multiplier = 3
- Expected range = 2.5 * 3 = 7.5 USD

Horizons:
- 5m candles → 15m prediction (3 candles ahead)
- 15m candles → 45m prediction (3 candles ahead)

Author: OrbiQuant IA
Version: 2.0.0
"""

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
from loguru import logger
from datetime import datetime, timedelta

# Import configuración centralizada de símbolos
from ..training.symbol_timeframe_trainer import SYMBOL_CONFIGS
from ..config.feature_flags import FeatureFlags

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
from sklearn.model_selection import train_test_split


@dataclass
class VolatilityFactor:
    """Volatility factor analysis for an asset"""
    symbol: str
    timeframe: str
    normal_variation: float      # Median range (typical movement)
    strong_movement: float       # 85th percentile (expansion)
    noise_factor: float          # Mode (most common rounded range)
    atr_14: float               # 14-period ATR
    recommended_factor: float    # The factor to use for predictions


@dataclass
class FactorPrediction:
    """Prediction using factor-based approach"""
    timestamp: datetime
    current_price: float
    base_factor: float
    multiplier_high: float       # Predicted multiplier for high
    multiplier_low: float        # Predicted multiplier for low
    predicted_high: float        # current_price + (factor * multiplier_high)
    predicted_low: float         # current_price - (factor * multiplier_low)
    confidence: float
    horizon: str                 # "15m" or "45m"

    def to_dict(self) -> Dict:
        return {
            'timestamp': self.timestamp.isoformat() if self.timestamp else None,
            'current_price': self.current_price,
            'base_factor': self.base_factor,
            'multiplier_high': self.multiplier_high,
            'multiplier_low': self.multiplier_low,
            'predicted_high': self.predicted_high,
            'predicted_low': self.predicted_low,
            'confidence': self.confidence,
            'horizon': self.horizon
        }


class VolatilityAnalyzer:
    """
    Analyzes price data to find the base volatility factor for an asset.
    """

    def __init__(self, symbol: str):
        self.symbol = symbol

    def calculate_factor(
        self,
        df: pd.DataFrame,
        timeframe: str = '5m'
    ) -> VolatilityFactor:
        """
        Calculate the volatility factor for the given data.

        Args:
            df: DataFrame with OHLCV data
            timeframe: Timeframe of the candles

        Returns:
            VolatilityFactor with all metrics
        """
        # Ensure we have the required columns
        required = ['Open', 'High', 'Low', 'Close']
        for col in required:
            if col not in df.columns:
                raise ValueError(f"Missing required column: {col}")

        # 1. Calculate absolute range for each candle
        df = df.copy()
        df['range'] = df['High'] - df['Low']
        df['body'] = (df['Close'] - df['Open']).abs()

        # 2. Statistical metrics
        # Median - the "normal" movement ignoring outliers
        normal_variation = df['range'].median()

        # 85th percentile - strong/expansion movement
        strong_movement = np.percentile(df['range'], 85)

        # Mode - most common rounded range (the "cluster")
        if self.symbol == 'XAUUSD':
            # Round to nearest 0.5 for gold
            rounded_ranges = (df['range'] * 2).round() / 2
        elif 'USD' in self.symbol:
            # Round to nearest pip (0.0001) for forex
            rounded_ranges = (df['range'] * 10000).round() / 10000
        else:
            rounded_ranges = df['range'].round(2)

        noise_factor = rounded_ranges.mode().iloc[0] if len(rounded_ranges.mode()) > 0 else normal_variation

        # 3. ATR (14-period)
        df['tr'] = np.maximum(
            df['High'] - df['Low'],
            np.maximum(
                abs(df['High'] - df['Close'].shift(1)),
                abs(df['Low'] - df['Close'].shift(1))
            )
        )
        atr_14 = df['tr'].rolling(14).mean().iloc[-1]

        # 4. Recommended factor (use median for stability)
        recommended_factor = normal_variation

        logger.info(f"{self.symbol} {timeframe} Factor Analysis:")
        logger.info(f"  Normal Variation (Median): {normal_variation:.4f}")
        logger.info(f"  Strong Movement (P85): {strong_movement:.4f}")
        logger.info(f"  Noise Factor (Mode): {noise_factor:.4f}")
        logger.info(f"  ATR(14): {atr_14:.4f}")
        logger.info(f"  Recommended Factor: {recommended_factor:.4f}")

        return VolatilityFactor(
            symbol=self.symbol,
            timeframe=timeframe,
            normal_variation=normal_variation,
            strong_movement=strong_movement,
            noise_factor=noise_factor,
            atr_14=atr_14,
            recommended_factor=recommended_factor
        )


class RangePredictorFactor:
    """
    Range predictor using factor-based multiplier approach.

    Models:
    - 5m → 15m: Predicts multiplier for next 3 5-minute candles
    - 15m → 45m: Predicts multiplier for next 3 15-minute candles
    """

    def __init__(self, symbol: str, config: Dict[str, Any] = None):
        """
        Initialize the factor-based range predictor.

        Args:
            symbol: Trading symbol (e.g., 'XAUUSD', 'EURUSD')
            config: Configuration dictionary
        """
        self.symbol = symbol
        self.config = config or self._default_config()

        # Volatility factors (calculated during training)
        self.factors: Dict[str, VolatilityFactor] = {}

        # Models for each timeframe
        self.models: Dict[str, Any] = {}

        # Training metrics
        self.metrics: Dict[str, Dict] = {}

        self._is_trained = False

        logger.info(f"Initialized RangePredictorFactor for {symbol}")

    def _default_config(self) -> Dict:
        """Default configuration"""
        return {
            'horizons': {
                '5m': {'lookahead': 3, 'target_minutes': 15},   # 5m → 15m
                '15m': {'lookahead': 3, 'target_minutes': 45},  # 15m → 45m
            },
            'xgboost': {
                'n_estimators': 300,
                'max_depth': 6,
                'learning_rate': 0.03,
                'subsample': 0.8,
                'colsample_bytree': 0.8,
                'min_child_weight': 3,
                'gamma': 0.1,
                'reg_alpha': 0.1,
                'reg_lambda': 1.0,
                'tree_method': 'hist',
                'random_state': 42,
                'n_jobs': -1
            },
            'feature_windows': [3, 5, 10, 20],  # Windows for rolling features
        }

    def _create_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Create features for prediction.

        Features include:
        - Price-based: returns, range ratios, body ratios
        - Volatility: rolling ATR, range percentiles
        - Session: hour, day of week
        - Technical: momentum, trend strength
        """
        df = df.copy()

        # Basic price features
        df['returns'] = df['Close'].pct_change()
        df['range'] = df['High'] - df['Low']
        df['body'] = (df['Close'] - df['Open']).abs()
        df['body_ratio'] = df['body'] / df['range'].replace(0, np.nan)
        df['upper_wick'] = df['High'] - df[['Open', 'Close']].max(axis=1)
        df['lower_wick'] = df[['Open', 'Close']].min(axis=1) - df['Low']

        # Rolling features for different windows
        for window in self.config['feature_windows']:
            # Volatility
            df[f'range_mean_{window}'] = df['range'].rolling(window).mean()
            df[f'range_std_{window}'] = df['range'].rolling(window).std()
            df[f'range_max_{window}'] = df['range'].rolling(window).max()

            # Returns
            df[f'returns_mean_{window}'] = df['returns'].rolling(window).mean()
            df[f'returns_std_{window}'] = df['returns'].rolling(window).std()

            # Momentum
            df[f'momentum_{window}'] = df['Close'] - df['Close'].shift(window)

            # ATR
            df[f'atr_{window}'] = df['range'].rolling(window).mean()

        # Session features (if datetime index)
        if isinstance(df.index, pd.DatetimeIndex):
            df['hour'] = df.index.hour
            df['day_of_week'] = df.index.dayofweek
            df['is_london'] = ((df['hour'] >= 8) & (df['hour'] < 16)).astype(int)
            df['is_ny'] = ((df['hour'] >= 14) & (df['hour'] < 21)).astype(int)
            df['is_overlap'] = ((df['hour'] >= 12) & (df['hour'] < 16)).astype(int)

        # Trend features
        df['trend_3'] = np.sign(df['Close'] - df['Close'].shift(3))
        df['trend_5'] = np.sign(df['Close'] - df['Close'].shift(5))
        df['trend_10'] = np.sign(df['Close'] - df['Close'].shift(10))

        # Range expansion/contraction
        df['range_expansion'] = df['range'] / df['range'].shift(1)
        df['is_expansion'] = (df['range'] > df[f'range_mean_{5}']).astype(int)

        return df

    def _create_targets(
        self,
        df: pd.DataFrame,
        factor: float,
        lookahead: int
    ) -> Tuple[pd.Series, pd.Series]:
        """
        Create target variables: multipliers for high and low.

        Args:
            df: DataFrame with OHLCV data
            factor: Base volatility factor
            lookahead: Number of candles to look ahead

        Returns:
            Tuple of (multiplier_high, multiplier_low)
        """
        n = len(df)

        # Calculate actual high/low in the lookahead period
        actual_highs = []
        actual_lows = []

        for i in range(n):
            end_idx = min(i + lookahead + 1, n)
            if i + 1 < n:
                future_high = df['High'].iloc[i+1:end_idx].max()
                future_low = df['Low'].iloc[i+1:end_idx].min()
            else:
                future_high = df['High'].iloc[i]
                future_low = df['Low'].iloc[i]

            actual_highs.append(future_high)
            actual_lows.append(future_low)

        df = df.copy()
        df['actual_high'] = actual_highs
        df['actual_low'] = actual_lows

        # Calculate delta from current close
        df['delta_high'] = df['actual_high'] - df['Close']
        df['delta_low'] = df['Close'] - df['actual_low']

        # Convert to multipliers of the factor
        # Ensure factor > 0
        factor = max(factor, 0.0001)

        df['multiplier_high'] = df['delta_high'] / factor
        df['multiplier_low'] = df['delta_low'] / factor

        # Clip multipliers to reasonable range [0, 10]
        df['multiplier_high'] = df['multiplier_high'].clip(0, 10)
        df['multiplier_low'] = df['multiplier_low'].clip(0, 10)

        return df['multiplier_high'], df['multiplier_low']

    def train(
        self,
        data_5m: pd.DataFrame,
        data_15m: pd.DataFrame = None,
        validation_split: float = 0.2
    ) -> Dict[str, Dict]:
        """
        Train both models (5m→15m and 15m→45m).

        Args:
            data_5m: 5-minute OHLCV data
            data_15m: 15-minute OHLCV data (will be resampled from 5m if not provided)
            validation_split: Fraction of data for validation

        Returns:
            Dictionary of training metrics
        """
        logger.info(f"Training RangePredictorFactor for {self.symbol}")

        # Resample 5m to 15m if not provided
        if data_15m is None:
            data_15m = self._resample_to_15m(data_5m)

        metrics = {}

        # Train 5m → 15m model
        logger.info("Training 5m → 15m model...")
        metrics['5m'] = self._train_model(
            data_5m, '5m',
            lookahead=self.config['horizons']['5m']['lookahead']
        )

        # Train 15m → 45m model
        logger.info("Training 15m → 45m model...")
        metrics['15m'] = self._train_model(
            data_15m, '15m',
            lookahead=self.config['horizons']['15m']['lookahead']
        )

        self._is_trained = True
        self.metrics = metrics

        return metrics

    def _train_model(
        self,
        df: pd.DataFrame,
        timeframe: str,
        lookahead: int
    ) -> Dict:
        """Train a single model for a specific timeframe."""
        # Calculate volatility factor
        analyzer = VolatilityAnalyzer(self.symbol)
        factor = analyzer.calculate_factor(df, timeframe)
        self.factors[timeframe] = factor

        # Create features
        df_features = self._create_features(df)

        # Create targets
        mult_high, mult_low = self._create_targets(
            df_features, factor.recommended_factor, lookahead
        )
        df_features['mult_high'] = mult_high
        df_features['mult_low'] = mult_low

        # Drop NaN rows
        df_features = df_features.dropna()

        if len(df_features) < 100:
            logger.warning(f"Not enough data for {timeframe}: {len(df_features)} rows")
            return {'error': 'Insufficient data'}

        # Select feature columns
        feature_cols = [col for col in df_features.columns if col not in [
            'Open', 'High', 'Low', 'Close', 'Volume',
            'mult_high', 'mult_low', 'actual_high', 'actual_low',
            'delta_high', 'delta_low'
        ]]

        X = df_features[feature_cols].values
        y_high = df_features['mult_high'].values
        y_low = df_features['mult_low'].values

        # Train/test split
        X_train, X_test, y_high_train, y_high_test, y_low_train, y_low_test = train_test_split(
            X, y_high, y_low, test_size=0.2, shuffle=False
        )

        # Train high model
        xgb_params = self.config['xgboost'].copy()
        model_high = XGBRegressor(**xgb_params)
        model_high.fit(X_train, y_high_train)

        # Train low model
        model_low = XGBRegressor(**xgb_params)
        model_low.fit(X_train, y_low_train)

        # Store models
        self.models[f'{timeframe}_high'] = model_high
        self.models[f'{timeframe}_low'] = model_low
        self.models[f'{timeframe}_features'] = feature_cols

        # Evaluate
        pred_high = model_high.predict(X_test)
        pred_low = model_low.predict(X_test)

        metrics = {
            'factor': factor.recommended_factor,
            'high': {
                'mae': mean_absolute_error(y_high_test, pred_high),
                'rmse': np.sqrt(mean_squared_error(y_high_test, pred_high)),
                'r2': r2_score(y_high_test, pred_high),
            },
            'low': {
                'mae': mean_absolute_error(y_low_test, pred_low),
                'rmse': np.sqrt(mean_squared_error(y_low_test, pred_low)),
                'r2': r2_score(y_low_test, pred_low),
            },
            'samples_train': len(X_train),
            'samples_test': len(X_test),
        }

        logger.info(f"{timeframe} Model Metrics:")
        logger.info(f"  Factor: {factor.recommended_factor:.4f}")
        logger.info(f"  High - MAE: {metrics['high']['mae']:.3f}, R2: {metrics['high']['r2']:.3f}")
        logger.info(f"  Low  - MAE: {metrics['low']['mae']:.3f}, R2: {metrics['low']['r2']:.3f}")

        return metrics

    def _resample_to_15m(self, df_5m: pd.DataFrame) -> pd.DataFrame:
        """Resample 5-minute data to 15-minute."""
        ohlc = {
            'Open': 'first',
            'High': 'max',
            'Low': 'min',
            'Close': 'last',
            'Volume': 'sum'
        }
        return df_5m.resample('15min').agg(ohlc).dropna()

    def predict(
        self,
        df: pd.DataFrame,
        timeframe: str
    ) -> List[FactorPrediction]:
        """
        Generate predictions for the given data.

        Args:
            df: OHLCV DataFrame
            timeframe: '5m' or '15m'

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

        if timeframe not in self.factors:
            raise ValueError(f"No model trained for timeframe: {timeframe}")

        factor = self.factors[timeframe]
        horizon = '15m' if timeframe == '5m' else '45m'

        # Create features
        df_features = self._create_features(df)
        df_features = df_features.dropna()

        feature_cols = self.models[f'{timeframe}_features']

        # Ensure all feature columns exist
        missing_cols = set(feature_cols) - set(df_features.columns)
        for col in missing_cols:
            df_features[col] = 0

        X = df_features[feature_cols].values

        # Predict multipliers
        mult_high = self.models[f'{timeframe}_high'].predict(X)
        mult_low = self.models[f'{timeframe}_low'].predict(X)

        # Generate predictions
        predictions = []
        for i in range(len(df_features)):
            current_price = df_features['Close'].iloc[i]
            timestamp = df_features.index[i] if isinstance(df_features.index, pd.DatetimeIndex) else None

            pred_high = current_price + (factor.recommended_factor * mult_high[i])
            pred_low = current_price - (factor.recommended_factor * mult_low[i])

            # Confidence based on multiplier magnitude (lower = more confident)
            confidence = 1.0 / (1.0 + 0.1 * (mult_high[i] + mult_low[i]))

            predictions.append(FactorPrediction(
                timestamp=timestamp,
                current_price=current_price,
                base_factor=factor.recommended_factor,
                multiplier_high=mult_high[i],
                multiplier_low=mult_low[i],
                predicted_high=pred_high,
                predicted_low=pred_low,
                confidence=confidence,
                horizon=horizon
            ))

        return predictions

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

        # Save models
        for name, model in self.models.items():
            if not isinstance(model, list):  # Skip feature lists
                joblib.dump(model, path / f'{name}.joblib')
            else:
                joblib.dump(model, path / f'{name}.joblib')

        # Save metadata
        metadata = {
            'symbol': self.symbol,
            'config': self.config,
            'factors': {k: vars(v) for k, v in self.factors.items()},
            'metrics': self.metrics,
        }
        joblib.dump(metadata, path / 'metadata.joblib')

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

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

        metadata = joblib.load(path / 'metadata.joblib')
        self.symbol = metadata['symbol']
        self.config = metadata['config']
        self.metrics = metadata['metrics']

        # Reconstruct factors
        self.factors = {}
        for k, v in metadata['factors'].items():
            self.factors[k] = VolatilityFactor(**v)

        # 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 model from {path}")


# ============================================================================
# Data Generation for Testing
# ============================================================================

class PriceDataGenerator:
    """Generates realistic price data for testing.

    Uses centralized SYMBOL_CONFIGS for dynamic configuration instead of
    hardcoded values. Supports 5+ symbols automatically.
    """

    # Precios base aproximados (actualizables)
    _CURRENT_PRICES = {
        'XAUUSD': 2650.0,
        'BTCUSD': 95000.0,
        'EURUSD': 1.0420,
        'GBPUSD': 1.2650,
        'USDJPY': 157.50,
    }

    # Legacy fallback (solo si FeatureFlags.USE_CENTRALIZED_CONFIGS = False)
    _LEGACY_SYMBOLS = {
        'XAUUSD': {'base': 2650.0, 'volatility': 0.0012, 'factor': 2.5},
        'EURUSD': {'base': 1.0420, 'volatility': 0.0004, 'factor': 0.0003},
    }

    def __init__(self, symbol: str, seed: int = 42):
        self.symbol = symbol
        self.config = self._build_config(symbol)
        np.random.seed(seed)

    def _build_config(self, symbol: str) -> dict:
        """Construir configuración desde SYMBOL_CONFIGS centralizado.

        Args:
            symbol: Símbolo de trading (e.g., 'XAUUSD', 'BTCUSD')

        Returns:
            dict con keys: base, volatility, factor
        """
        # Si feature flag está desactivado, usar legacy
        if not FeatureFlags.USE_CENTRALIZED_CONFIGS:
            return self._LEGACY_SYMBOLS.get(symbol, self._LEGACY_SYMBOLS['XAUUSD'])

        # Usar configuración centralizada
        symbol_config = SYMBOL_CONFIGS.get(symbol)
        if symbol_config:
            base_price = self._CURRENT_PRICES.get(symbol, 100.0)
            return {
                'base': base_price,
                'volatility': symbol_config.base_factor / 5000,  # Normalizado
                'factor': symbol_config.base_factor
            }
        else:
            # Fallback para símbolos desconocidos
            logger.warning(f"Unknown symbol {symbol}, using default config")
            return {'base': 100.0, 'volatility': 0.001, 'factor': 1.0}

    def generate(self, days: int = 30, timeframe: str = '5m') -> pd.DataFrame:
        """Generate OHLCV data."""
        if timeframe == '5m':
            candles_per_day = 288
        elif timeframe == '15m':
            candles_per_day = 96
        else:
            candles_per_day = 288

        total_candles = days * candles_per_day
        price = self.config['base']
        vol = self.config['volatility']

        data = []
        start_date = datetime(2026, 1, 1)

        for i in range(total_candles):
            if timeframe == '5m':
                timestamp = start_date + timedelta(minutes=i * 5)
            else:
                timestamp = start_date + timedelta(minutes=i * 15)

            # Session-based volatility
            hour = timestamp.hour
            if 8 <= hour < 16:
                session_mult = 1.3
            elif 14 <= hour < 21:
                session_mult = 1.2
            else:
                session_mult = 0.7

            # Random walk
            returns = np.random.normal(0, vol * session_mult)
            open_price = price
            close_price = price * (1 + returns)

            hl_range = abs(returns) + vol * session_mult * np.random.random()
            high_price = max(open_price, close_price) * (1 + hl_range * np.random.random())
            low_price = min(open_price, close_price) * (1 - hl_range * np.random.random())

            volume = 1000 * session_mult * (1 + np.random.random())

            data.append({
                'Open': open_price,
                'High': high_price,
                'Low': low_price,
                'Close': close_price,
                'Volume': volume
            })

            price = close_price

        df = pd.DataFrame(data, index=pd.DatetimeIndex(
            [start_date + timedelta(minutes=i * (5 if timeframe == '5m' else 15))
             for i in range(total_candles)]
        ))
        df.index.name = 'Date'

        return df


if __name__ == "__main__":
    # Test the model
    print("\n" + "="*60)
    print("Testing RangePredictorFactor")
    print("="*60)

    # Generate test data
    generator = PriceDataGenerator('XAUUSD', seed=42)
    data_5m = generator.generate(days=60, timeframe='5m')
    print(f"\nGenerated {len(data_5m)} 5-minute candles")

    # Train model
    model = RangePredictorFactor('XAUUSD')
    metrics = model.train(data_5m)

    print("\n" + "="*60)
    print("Training Complete!")
    print("="*60)

    for tf, m in metrics.items():
        if 'error' not in m:
            print(f"\n{tf} Model:")
            print(f"  Factor: {m['factor']:.4f}")
            print(f"  High MAE: {m['high']['mae']:.3f} multipliers")
            print(f"  Low MAE: {m['low']['mae']:.3f} multipliers")

    # Test prediction
    print("\n" + "="*60)
    print("Sample Predictions")
    print("="*60)

    test_data = data_5m.iloc[-100:]
    predictions = model.predict(test_data, '5m')

    for pred in predictions[:5]:
        print(f"\nPrice: {pred.current_price:.2f}")
        print(f"  Factor: {pred.base_factor:.4f}")
        print(f"  Multiplier High: {pred.multiplier_high:.2f}x → Pred High: {pred.predicted_high:.2f}")
        print(f"  Multiplier Low: {pred.multiplier_low:.2f}x → Pred Low: {pred.predicted_low:.2f}")
        print(f"  Expected Range: {pred.predicted_high - pred.predicted_low:.2f}")
