"""
Range Predictor - Phase 2
Predicts ΔHigh and ΔLow (price ranges) for multiple horizons
"""

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

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


@dataclass
class RangePrediction:
    """Single range prediction result"""
    horizon: str                    # "15m" or "1h"
    delta_high: float              # Predicted ΔHigh
    delta_low: float               # Predicted ΔLow
    delta_high_bin: Optional[int] = None  # Bin classification (0-3)
    delta_low_bin: Optional[int] = None
    confidence_high: float = 0.0   # Confidence for high prediction
    confidence_low: float = 0.0    # Confidence for low prediction
    timestamp: Optional[pd.Timestamp] = None

    def to_dict(self) -> Dict:
        """Convert to dictionary"""
        return {
            'horizon': self.horizon,
            'delta_high': float(self.delta_high),
            'delta_low': float(self.delta_low),
            'delta_high_bin': int(self.delta_high_bin) if self.delta_high_bin is not None else None,
            'delta_low_bin': int(self.delta_low_bin) if self.delta_low_bin is not None else None,
            'confidence_high': float(self.confidence_high),
            'confidence_low': float(self.confidence_low)
        }


@dataclass
class RangeModelMetrics:
    """Metrics for range prediction model"""
    horizon: str
    target_type: str  # 'high' or 'low'

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

    # Classification metrics (for bins)
    bin_accuracy: float = 0.0
    bin_f1: float = 0.0

    # Sample counts
    n_train: int = 0
    n_test: int = 0


class RangePredictor:
    """
    Predictor for price ranges (ΔHigh/ΔLow)

    Creates separate models for each:
    - Horizon (15m, 1h)
    - Target type (high, low)
    - Task (regression for values, classification for bins)
    """

    def __init__(self, config: Dict[str, Any] = None):
        """
        Initialize range predictor

        Args:
            config: Configuration dictionary
        """
        self.config = config or self._default_config()
        self.horizons = self.config.get('horizons', ['15m', '1h'])
        self.models = {}
        self.metrics = {}
        self.feature_importance = {}
        self._is_trained = False

        # Initialize models
        self._init_models()

    def _default_config(self) -> Dict:
        """Default configuration"""
        return {
            'horizons': ['15m', '1h'],
            'include_bins': True,
            'xgboost': {
                'n_estimators': 200,
                'max_depth': 5,
                'learning_rate': 0.05,
                '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
            }
        }

    def _init_models(self):
        """Initialize all models"""
        if not HAS_XGBOOST:
            raise ImportError("XGBoost is required for RangePredictor")

        xgb_params = self.config.get('xgboost', {})

        # Check GPU availability
        try:
            import torch
            if torch.cuda.is_available():
                xgb_params['device'] = 'cuda'
                logger.info("Using GPU for XGBoost")
        except:
            pass

        for horizon in self.horizons:
            # Regression models for delta values
            self.models[f'{horizon}_high_reg'] = XGBRegressor(**xgb_params)
            self.models[f'{horizon}_low_reg'] = XGBRegressor(**xgb_params)

            # Classification models for bins (if enabled)
            if self.config.get('include_bins', True):
                bin_params = xgb_params.copy()
                bin_params['objective'] = 'multi:softprob'
                bin_params['num_class'] = 4
                bin_params.pop('n_jobs', None)  # Not compatible with multiclass

                self.models[f'{horizon}_high_bin'] = XGBClassifier(**bin_params)
                self.models[f'{horizon}_low_bin'] = XGBClassifier(**bin_params)

        logger.info(f"Initialized {len(self.models)} models for {len(self.horizons)} horizons")

    def train(
        self,
        X_train: Union[pd.DataFrame, np.ndarray],
        y_train: Dict[str, Union[pd.Series, np.ndarray]],
        X_val: Optional[Union[pd.DataFrame, np.ndarray]] = None,
        y_val: Optional[Dict[str, Union[pd.Series, np.ndarray]]] = None,
        early_stopping_rounds: int = 50
    ) -> Dict[str, RangeModelMetrics]:
        """
        Train all range prediction models

        Args:
            X_train: Training features
            y_train: Dictionary of training targets with keys like:
                     'delta_high_15m', 'delta_low_15m', 'bin_high_15m', etc.
            X_val: Validation features (optional)
            y_val: Validation targets (optional)
            early_stopping_rounds: Early stopping patience

        Returns:
            Dictionary of metrics for each model
        """
        logger.info(f"Training range predictor with {len(X_train)} samples")

        # Convert to numpy if needed
        X_train_np = X_train.values if isinstance(X_train, pd.DataFrame) else X_train

        if X_val is not None:
            X_val_np = X_val.values if isinstance(X_val, pd.DataFrame) else X_val
            eval_set = [(X_val_np, None)]  # Will be updated per model
        else:
            eval_set = None

        metrics = {}

        for horizon in self.horizons:
            # Train regression models
            for target_type in ['high', 'low']:
                model_key = f'{horizon}_{target_type}_reg'
                target_key = f'delta_{target_type}_{horizon}'

                if target_key not in y_train:
                    logger.warning(f"Target {target_key} not found, skipping")
                    continue

                y_train_target = y_train[target_key]
                y_train_np = y_train_target.values if isinstance(y_train_target, pd.Series) else y_train_target

                # Prepare validation data
                fit_params = {}
                if X_val is not None and y_val is not None and target_key in y_val:
                    y_val_target = y_val[target_key]
                    y_val_np = y_val_target.values if isinstance(y_val_target, pd.Series) else y_val_target
                    fit_params['eval_set'] = [(X_val_np, y_val_np)]

                # Train model
                logger.info(f"Training {model_key}...")
                self.models[model_key].fit(X_train_np, y_train_np, **fit_params)

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

                # Calculate metrics
                train_pred = self.models[model_key].predict(X_train_np)
                metrics[model_key] = self._calculate_regression_metrics(
                    y_train_np, train_pred, horizon, target_type, len(X_train_np)
                )

                if X_val is not None and y_val is not None and target_key in y_val:
                    val_pred = self.models[model_key].predict(X_val_np)
                    val_metrics = self._calculate_regression_metrics(
                        y_val_np, val_pred, horizon, target_type, len(X_val_np)
                    )
                    metrics[f'{model_key}_val'] = val_metrics

            # Train classification models (bins)
            if self.config.get('include_bins', True):
                for target_type in ['high', 'low']:
                    model_key = f'{horizon}_{target_type}_bin'
                    target_key = f'bin_{target_type}_{horizon}'

                    if target_key not in y_train:
                        logger.warning(f"Target {target_key} not found, skipping")
                        continue

                    y_train_target = y_train[target_key]
                    y_train_np = y_train_target.values if isinstance(y_train_target, pd.Series) else y_train_target

                    # Remove NaN values
                    valid_mask = ~np.isnan(y_train_np)
                    X_train_valid = X_train_np[valid_mask]
                    y_train_valid = y_train_np[valid_mask].astype(int)

                    if len(X_train_valid) == 0:
                        logger.warning(f"No valid samples for {model_key}")
                        continue

                    # Train model
                    logger.info(f"Training {model_key}...")
                    self.models[model_key].fit(X_train_valid, y_train_valid)

                    # Calculate metrics
                    train_pred = self.models[model_key].predict(X_train_valid)
                    metrics[model_key] = self._calculate_classification_metrics(
                        y_train_valid, train_pred, horizon, target_type, len(X_train_valid)
                    )

        self._is_trained = True
        self.metrics = metrics

        logger.info(f"Training complete. Trained {len([k for k in metrics.keys() if '_val' not in k])} models")
        return metrics

    def predict(
        self,
        X: Union[pd.DataFrame, np.ndarray],
        include_bins: bool = True
    ) -> List[RangePrediction]:
        """
        Generate range predictions

        Args:
            X: Features for prediction
            include_bins: Include bin predictions

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

        X_np = X.values if isinstance(X, pd.DataFrame) else X

        # Handle single sample
        if X_np.ndim == 1:
            X_np = X_np.reshape(1, -1)

        predictions = []

        for horizon in self.horizons:
            # Regression predictions
            delta_high = self.models[f'{horizon}_high_reg'].predict(X_np)
            delta_low = self.models[f'{horizon}_low_reg'].predict(X_np)

            # Bin predictions
            bin_high = None
            bin_low = None
            conf_high = 0.0
            conf_low = 0.0

            if include_bins and self.config.get('include_bins', True):
                bin_high_model = self.models.get(f'{horizon}_high_bin')
                bin_low_model = self.models.get(f'{horizon}_low_bin')

                if bin_high_model is not None:
                    bin_high = bin_high_model.predict(X_np)
                    proba_high = bin_high_model.predict_proba(X_np)
                    conf_high = np.max(proba_high, axis=1)

                if bin_low_model is not None:
                    bin_low = bin_low_model.predict(X_np)
                    proba_low = bin_low_model.predict_proba(X_np)
                    conf_low = np.max(proba_low, axis=1)

            # Create predictions for each sample
            for i in range(len(X_np)):
                pred = RangePrediction(
                    horizon=horizon,
                    delta_high=float(delta_high[i]),
                    delta_low=float(delta_low[i]),
                    delta_high_bin=int(bin_high[i]) if bin_high is not None else None,
                    delta_low_bin=int(bin_low[i]) if bin_low is not None else None,
                    confidence_high=float(conf_high[i]) if isinstance(conf_high, np.ndarray) else conf_high,
                    confidence_low=float(conf_low[i]) if isinstance(conf_low, np.ndarray) else conf_low
                )
                predictions.append(pred)

        return predictions

    def predict_single(
        self,
        X: Union[pd.DataFrame, np.ndarray]
    ) -> Dict[str, RangePrediction]:
        """
        Predict for a single sample, return dict keyed by horizon

        Args:
            X: Single sample features

        Returns:
            Dictionary with horizon as key and RangePrediction as value
        """
        preds = self.predict(X)
        return {pred.horizon: pred for pred in preds}

    def evaluate(
        self,
        X_test: Union[pd.DataFrame, np.ndarray],
        y_test: Dict[str, Union[pd.Series, np.ndarray]]
    ) -> Dict[str, RangeModelMetrics]:
        """
        Evaluate model on test data

        Args:
            X_test: Test features
            y_test: Test targets

        Returns:
            Dictionary of metrics
        """
        X_np = X_test.values if isinstance(X_test, pd.DataFrame) else X_test
        metrics = {}

        for horizon in self.horizons:
            for target_type in ['high', 'low']:
                # Regression evaluation
                model_key = f'{horizon}_{target_type}_reg'
                target_key = f'delta_{target_type}_{horizon}'

                if target_key in y_test and model_key in self.models:
                    y_true = y_test[target_key]
                    y_true_np = y_true.values if isinstance(y_true, pd.Series) else y_true

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

                    metrics[model_key] = self._calculate_regression_metrics(
                        y_true_np, y_pred, horizon, target_type, len(X_np)
                    )

                # Classification evaluation
                if self.config.get('include_bins', True):
                    model_key = f'{horizon}_{target_type}_bin'
                    target_key = f'bin_{target_type}_{horizon}'

                    if target_key in y_test and model_key in self.models:
                        y_true = y_test[target_key]
                        y_true_np = y_true.values if isinstance(y_true, pd.Series) else y_true

                        # Remove NaN
                        valid_mask = ~np.isnan(y_true_np)
                        if valid_mask.sum() > 0:
                            y_pred = self.models[model_key].predict(X_np[valid_mask])

                            metrics[model_key] = self._calculate_classification_metrics(
                                y_true_np[valid_mask].astype(int), y_pred,
                                horizon, target_type, valid_mask.sum()
                            )

        return metrics

    def _calculate_regression_metrics(
        self,
        y_true: np.ndarray,
        y_pred: np.ndarray,
        horizon: str,
        target_type: str,
        n_samples: int
    ) -> RangeModelMetrics:
        """Calculate regression metrics"""
        # Avoid division by zero in MAPE
        mask = y_true != 0
        if mask.sum() > 0:
            mape = np.mean(np.abs((y_true[mask] - y_pred[mask]) / y_true[mask])) * 100
        else:
            mape = 0.0

        return RangeModelMetrics(
            horizon=horizon,
            target_type=target_type,
            mae=mean_absolute_error(y_true, y_pred),
            mape=mape,
            rmse=np.sqrt(mean_squared_error(y_true, y_pred)),
            r2=r2_score(y_true, y_pred),
            n_test=n_samples
        )

    def _calculate_classification_metrics(
        self,
        y_true: np.ndarray,
        y_pred: np.ndarray,
        horizon: str,
        target_type: str,
        n_samples: int
    ) -> RangeModelMetrics:
        """Calculate classification metrics"""
        return RangeModelMetrics(
            horizon=horizon,
            target_type=target_type,
            bin_accuracy=accuracy_score(y_true, y_pred),
            bin_f1=f1_score(y_true, y_pred, average='weighted'),
            n_test=n_samples
        )

    def get_feature_importance(
        self,
        model_key: str = None,
        top_n: int = 20
    ) -> Dict[str, float]:
        """
        Get feature importance for a model

        Args:
            model_key: Specific model key, or None for average across all
            top_n: Number of top features to return

        Returns:
            Dictionary of feature importances
        """
        if model_key is not None:
            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)

        # Sort and return top N
        sorted_imp = dict(sorted(importance.items(), key=lambda x: x[1], reverse=True)[:top_n])
        return sorted_imp

    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():
            joblib.dump(model, path / f'{name}.joblib')

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

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

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

        # Load metadata
        metadata = joblib.load(path / 'metadata.joblib')
        self.config = metadata['config']
        self.horizons = metadata['horizons']
        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 range predictor from {path}")


if __name__ == "__main__":
    # Test range predictor
    import numpy as np

    # Create sample data
    np.random.seed(42)
    n_samples = 1000
    n_features = 20

    X = np.random.randn(n_samples, n_features)
    y = {
        'delta_high_15m': np.random.randn(n_samples) * 5 + 2,
        'delta_low_15m': np.random.randn(n_samples) * 5 + 2,
        'delta_high_1h': np.random.randn(n_samples) * 8 + 3,
        'delta_low_1h': np.random.randn(n_samples) * 8 + 3,
        'bin_high_15m': np.random.randint(0, 4, n_samples).astype(float),
        'bin_low_15m': np.random.randint(0, 4, n_samples).astype(float),
        'bin_high_1h': np.random.randint(0, 4, n_samples).astype(float),
        'bin_low_1h': np.random.randint(0, 4, n_samples).astype(float),
    }

    # Split data
    train_size = 800
    X_train, X_test = X[:train_size], X[train_size:]
    y_train = {k: v[:train_size] for k, v in y.items()}
    y_test = {k: v[train_size:] for k, v in y.items()}

    # Train predictor
    predictor = RangePredictor()
    metrics = predictor.train(X_train, y_train)

    print("\n=== Training Metrics ===")
    for name, m in metrics.items():
        if hasattr(m, 'mae') and m.mae > 0:
            print(f"{name}: MAE={m.mae:.4f}, RMSE={m.rmse:.4f}, R2={m.r2:.4f}")
        elif hasattr(m, 'bin_accuracy') and m.bin_accuracy > 0:
            print(f"{name}: Accuracy={m.bin_accuracy:.4f}, F1={m.bin_f1:.4f}")

    # Evaluate on test
    test_metrics = predictor.evaluate(X_test, y_test)
    print("\n=== Test Metrics ===")
    for name, m in test_metrics.items():
        if hasattr(m, 'mae') and m.mae > 0:
            print(f"{name}: MAE={m.mae:.4f}, RMSE={m.rmse:.4f}, R2={m.r2:.4f}")
        elif hasattr(m, 'bin_accuracy') and m.bin_accuracy > 0:
            print(f"{name}: Accuracy={m.bin_accuracy:.4f}, F1={m.bin_f1:.4f}")

    # Test prediction
    predictions = predictor.predict(X_test[:5])
    print("\n=== Sample Predictions ===")
    for pred in predictions:
        print(pred.to_dict())
