"""
TP vs SL Classifier - Phase 2
Binary classifier to predict if Take Profit or Stop Loss will be hit first
"""

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 XGBClassifier
    HAS_XGBOOST = True
except ImportError:
    HAS_XGBOOST = False
    logger.warning("XGBoost not available")

from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score,
    roc_auc_score, confusion_matrix, classification_report
)
from sklearn.calibration import CalibratedClassifierCV


@dataclass
class TPSLPrediction:
    """Single TP/SL prediction result"""
    horizon: str                    # "15m" or "1h"
    rr_config: str                  # "rr_2_1" or "rr_3_1"
    prob_tp_first: float           # P(TP hits first)
    prob_sl_first: float           # P(SL hits first) = 1 - prob_tp_first
    recommended_action: str         # "long", "short", "hold"
    confidence: float              # Confidence level
    entry_price: Optional[float] = None
    sl_price: Optional[float] = None
    tp_price: Optional[float] = None
    sl_distance: Optional[float] = None
    tp_distance: Optional[float] = None

    def to_dict(self) -> Dict:
        """Convert to dictionary"""
        return {
            'horizon': self.horizon,
            'rr_config': self.rr_config,
            'prob_tp_first': float(self.prob_tp_first),
            'prob_sl_first': float(self.prob_sl_first),
            'recommended_action': self.recommended_action,
            'confidence': float(self.confidence),
            'entry_price': float(self.entry_price) if self.entry_price else None,
            'sl_price': float(self.sl_price) if self.sl_price else None,
            'tp_price': float(self.tp_price) if self.tp_price else None,
            'sl_distance': float(self.sl_distance) if self.sl_distance else None,
            'tp_distance': float(self.tp_distance) if self.tp_distance else None
        }


@dataclass
class TPSLMetrics:
    """Metrics for TP/SL classifier"""
    horizon: str
    rr_config: str

    # Classification metrics
    accuracy: float = 0.0
    precision: float = 0.0
    recall: float = 0.0
    f1: float = 0.0
    roc_auc: float = 0.0

    # Class distribution
    tp_rate: float = 0.0           # Rate of TP outcomes
    sl_rate: float = 0.0           # Rate of SL outcomes

    # Confusion matrix
    true_positives: int = 0
    true_negatives: int = 0
    false_positives: int = 0
    false_negatives: int = 0

    # Sample counts
    n_samples: int = 0

    def to_dict(self) -> Dict:
        return {
            'horizon': self.horizon,
            'rr_config': self.rr_config,
            'accuracy': self.accuracy,
            'precision': self.precision,
            'recall': self.recall,
            'f1': self.f1,
            'roc_auc': self.roc_auc,
            'tp_rate': self.tp_rate,
            'n_samples': self.n_samples
        }


class TPSLClassifier:
    """
    Binary classifier for TP vs SL prediction

    Predicts the probability that Take Profit will be hit before Stop Loss
    for a given entry point and R:R configuration.
    """

    def __init__(self, config: Dict[str, Any] = None):
        """
        Initialize TP/SL classifier

        Args:
            config: Configuration dictionary
        """
        self.config = config or self._default_config()
        self.horizons = self.config.get('horizons', ['15m', '1h'])
        self.rr_configs = self.config.get('rr_configs', [
            {'name': 'rr_2_1', 'sl': 5.0, 'tp': 10.0},
            {'name': 'rr_3_1', 'sl': 5.0, 'tp': 15.0}
        ])

        self.probability_threshold = self.config.get('probability_threshold', 0.55)
        self.use_calibration = self.config.get('use_calibration', True)
        self.calibration_method = self.config.get('calibration_method', 'isotonic')

        self.models = {}
        self.calibrated_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'],
            'rr_configs': [
                {'name': 'rr_2_1', 'sl': 5.0, 'tp': 10.0},
                {'name': 'rr_3_1', 'sl': 5.0, 'tp': 15.0}
            ],
            'probability_threshold': 0.55,
            'use_calibration': True,
            'calibration_method': 'isotonic',
            '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,
                'scale_pos_weight': 1.0,
                'objective': 'binary:logistic',
                'eval_metric': 'auc',
                '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 TPSLClassifier")

        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:
            for rr in self.rr_configs:
                model_key = f'{horizon}_{rr["name"]}'
                self.models[model_key] = XGBClassifier(**xgb_params)

        logger.info(f"Initialized {len(self.models)} TP/SL classifiers")

    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,
        range_predictions: Optional[Dict[str, np.ndarray]] = None,
        sample_weights: Optional[np.ndarray] = None
    ) -> Dict[str, TPSLMetrics]:
        """
        Train all TP/SL classifiers

        Args:
            X_train: Training features
            y_train: Dictionary of training targets with keys like:
                     'tp_first_15m_rr_2_1', 'tp_first_1h_rr_2_1', etc.
            X_val: Validation features (optional)
            y_val: Validation targets (optional)
            range_predictions: Optional range predictions to use as features (stacking)
            sample_weights: Optional sample weights

        Returns:
            Dictionary of metrics for each model
        """
        logger.info(f"Training TP/SL classifier with {len(X_train)} samples")

        # Convert to numpy
        X_train_np = X_train.values if isinstance(X_train, pd.DataFrame) else X_train.copy()
        feature_names = X_train.columns.tolist() if isinstance(X_train, pd.DataFrame) else None

        # Add range predictions as features if provided (stacking)
        if range_predictions is not None:
            logger.info("Adding range predictions as features (stacking)")
            range_features = []
            range_names = []
            for name, pred in range_predictions.items():
                range_features.append(pred.reshape(-1, 1) if pred.ndim == 1 else pred)
                range_names.append(name)
            X_train_np = np.hstack([X_train_np] + range_features)
            if feature_names:
                feature_names = feature_names + range_names

        if X_val is not None:
            X_val_np = X_val.values if isinstance(X_val, pd.DataFrame) else X_val.copy()

        metrics = {}

        for horizon in self.horizons:
            for rr in self.rr_configs:
                model_key = f'{horizon}_{rr["name"]}'
                target_key = f'tp_first_{horizon}_{rr["name"]}'

                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

                # Adjust scale_pos_weight for class imbalance
                pos_rate = y_train_valid.mean()
                if pos_rate > 0 and pos_rate < 1:
                    scale_pos_weight = (1 - pos_rate) / pos_rate
                    self.models[model_key].set_params(scale_pos_weight=scale_pos_weight)
                    logger.info(f"{model_key}: TP rate={pos_rate:.2%}, scale_pos_weight={scale_pos_weight:.2f}")

                # 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
                    valid_val_mask = ~np.isnan(y_val_np)
                    if valid_val_mask.sum() > 0:
                        fit_params['eval_set'] = [(X_val_np[valid_val_mask], y_val_np[valid_val_mask].astype(int))]

                # Prepare sample weights
                weights = None
                if sample_weights is not None:
                    weights = sample_weights[valid_mask]

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

                # Calibrate probabilities if enabled
                if self.use_calibration and X_val is not None and y_val is not None:
                    logger.info(f"Calibrating {model_key}...")
                    self.calibrated_models[model_key] = CalibratedClassifierCV(
                        self.models[model_key],
                        method=self.calibration_method,
                        cv='prefit'
                    )
                    if target_key in y_val:
                        y_val_np = y_val[target_key]
                        y_val_np = y_val_np.values if isinstance(y_val_np, pd.Series) else y_val_np
                        valid_val_mask = ~np.isnan(y_val_np)
                        if valid_val_mask.sum() > 0:
                            self.calibrated_models[model_key].fit(
                                X_val_np[valid_val_mask],
                                y_val_np[valid_val_mask].astype(int)
                            )

                # Store feature importance
                if feature_names:
                    self.feature_importance[model_key] = dict(
                        zip(feature_names, self.models[model_key].feature_importances_)
                    )

                # Calculate metrics
                train_pred = self.models[model_key].predict(X_train_valid)
                train_prob = self.models[model_key].predict_proba(X_train_valid)[:, 1]

                metrics[model_key] = self._calculate_metrics(
                    y_train_valid, train_pred, train_prob,
                    horizon, rr['name']
                )

        self._is_trained = True
        self.metrics = metrics

        logger.info(f"Training complete. Trained {len(metrics)} classifiers")
        return metrics

    def predict_proba(
        self,
        X: Union[pd.DataFrame, np.ndarray],
        horizon: str = '15m',
        rr_config: str = 'rr_2_1',
        use_calibrated: bool = True
    ) -> np.ndarray:
        """
        Predict probability of TP hitting first

        Args:
            X: Features
            horizon: Prediction horizon
            rr_config: R:R configuration name
            use_calibrated: Use calibrated model if available

        Returns:
            Array of probabilities
        """
        if not self._is_trained:
            raise RuntimeError("Model must be trained before prediction")

        model_key = f'{horizon}_{rr_config}'
        X_np = X.values if isinstance(X, pd.DataFrame) else X

        # Use calibrated model if available
        if use_calibrated and model_key in self.calibrated_models:
            return self.calibrated_models[model_key].predict_proba(X_np)[:, 1]
        else:
            return self.models[model_key].predict_proba(X_np)[:, 1]

    def predict(
        self,
        X: Union[pd.DataFrame, np.ndarray],
        current_price: Optional[float] = None,
        direction: str = 'long'
    ) -> List[TPSLPrediction]:
        """
        Generate TP/SL predictions for all horizons and R:R configs

        Args:
            X: Features (single sample or batch)
            current_price: Current price for SL/TP calculation
            direction: Trade direction ('long' or 'short')

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

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

        predictions = []

        for horizon in self.horizons:
            for rr in self.rr_configs:
                model_key = f'{horizon}_{rr["name"]}'

                if model_key not in self.models:
                    continue

                # Get probabilities
                proba = self.predict_proba(X_np, horizon, rr['name'])

                for i in range(len(X_np)):
                    prob_tp = float(proba[i])
                    prob_sl = 1.0 - prob_tp

                    # Determine recommended action
                    if prob_tp >= self.probability_threshold:
                        action = direction
                    elif prob_sl >= self.probability_threshold:
                        action = 'short' if direction == 'long' else 'long'
                    else:
                        action = 'hold'

                    # Confidence based on how far from 0.5
                    confidence = abs(prob_tp - 0.5) * 2

                    # Calculate prices if current_price provided
                    entry_price = current_price
                    sl_price = None
                    tp_price = None

                    if current_price is not None:
                        if direction == 'long':
                            sl_price = current_price - rr['sl']
                            tp_price = current_price + rr['tp']
                        else:
                            sl_price = current_price + rr['sl']
                            tp_price = current_price - rr['tp']

                    pred = TPSLPrediction(
                        horizon=horizon,
                        rr_config=rr['name'],
                        prob_tp_first=prob_tp,
                        prob_sl_first=prob_sl,
                        recommended_action=action,
                        confidence=confidence,
                        entry_price=entry_price,
                        sl_price=sl_price,
                        tp_price=tp_price,
                        sl_distance=rr['sl'],
                        tp_distance=rr['tp']
                    )
                    predictions.append(pred)

        return predictions

    def predict_single(
        self,
        X: Union[pd.DataFrame, np.ndarray],
        current_price: Optional[float] = None,
        direction: str = 'long'
    ) -> Dict[str, TPSLPrediction]:
        """
        Predict for single sample, return dict keyed by model

        Args:
            X: Single sample features
            current_price: Current price
            direction: Trade direction

        Returns:
            Dictionary with (horizon, rr_config) as key
        """
        preds = self.predict(X, current_price, direction)
        return {f'{p.horizon}_{p.rr_config}': p for p in preds}

    def evaluate(
        self,
        X_test: Union[pd.DataFrame, np.ndarray],
        y_test: Dict[str, Union[pd.Series, np.ndarray]]
    ) -> Dict[str, TPSLMetrics]:
        """
        Evaluate classifier 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 rr in self.rr_configs:
                model_key = f'{horizon}_{rr["name"]}'
                target_key = f'tp_first_{horizon}_{rr["name"]}'

                if target_key not in y_test or model_key not in self.models:
                    continue

                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:
                    continue

                y_true_valid = y_true_np[valid_mask].astype(int)
                X_valid = X_np[valid_mask]

                y_pred = self.models[model_key].predict(X_valid)
                y_prob = self.predict_proba(X_valid, horizon, rr['name'])

                metrics[model_key] = self._calculate_metrics(
                    y_true_valid, y_pred, y_prob,
                    horizon, rr['name']
                )

        return metrics

    def _calculate_metrics(
        self,
        y_true: np.ndarray,
        y_pred: np.ndarray,
        y_prob: np.ndarray,
        horizon: str,
        rr_config: str
    ) -> TPSLMetrics:
        """Calculate all metrics"""
        cm = confusion_matrix(y_true, y_pred)

        # Handle case where one class is missing
        if cm.shape == (1, 1):
            if y_true[0] == 1:
                tn, fp, fn, tp = 0, 0, 0, cm[0, 0]
            else:
                tn, fp, fn, tp = cm[0, 0], 0, 0, 0
        else:
            tn, fp, fn, tp = cm.ravel()

        return TPSLMetrics(
            horizon=horizon,
            rr_config=rr_config,
            accuracy=accuracy_score(y_true, y_pred),
            precision=precision_score(y_true, y_pred, zero_division=0),
            recall=recall_score(y_true, y_pred, zero_division=0),
            f1=f1_score(y_true, y_pred, zero_division=0),
            roc_auc=roc_auc_score(y_true, y_prob) if len(np.unique(y_true)) > 1 else 0.5,
            tp_rate=y_true.mean(),
            sl_rate=1 - y_true.mean(),
            true_positives=int(tp),
            true_negatives=int(tn),
            false_positives=int(fp),
            false_negatives=int(fn),
            n_samples=len(y_true)
        )

    def get_feature_importance(
        self,
        model_key: str = None,
        top_n: int = 20
    ) -> Dict[str, float]:
        """Get feature importance"""
        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)

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

    def save(self, path: str):
        """Save classifier 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 calibrated models
        for name, model in self.calibrated_models.items():
            joblib.dump(model, path / f'{name}_calibrated.joblib')

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

        logger.info(f"Saved TP/SL classifier to {path}")

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

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

        # Load models
        self.models = {}
        self.calibrated_models = {}
        for model_file in path.glob('*.joblib'):
            if model_file.name == 'metadata.joblib':
                continue
            name = model_file.stem
            if name.endswith('_calibrated'):
                self.calibrated_models[name.replace('_calibrated', '')] = joblib.load(model_file)
            else:
                self.models[name] = joblib.load(model_file)

        self._is_trained = True
        logger.info(f"Loaded TP/SL classifier from {path}")


if __name__ == "__main__":
    # Test TP/SL classifier
    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 = {
        'tp_first_15m_rr_2_1': (np.random.rand(n_samples) > 0.55).astype(float),
        'tp_first_15m_rr_3_1': (np.random.rand(n_samples) > 0.65).astype(float),
        'tp_first_1h_rr_2_1': (np.random.rand(n_samples) > 0.50).astype(float),
        'tp_first_1h_rr_3_1': (np.random.rand(n_samples) > 0.60).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 classifier
    classifier = TPSLClassifier()
    metrics = classifier.train(X_train, y_train, X_test, y_test)

    print("\n=== Training Metrics ===")
    for name, m in metrics.items():
        print(f"{name}: Accuracy={m.accuracy:.4f}, ROC-AUC={m.roc_auc:.4f}, "
              f"TP Rate={m.tp_rate:.2%}")

    # Evaluate on test
    test_metrics = classifier.evaluate(X_test, y_test)
    print("\n=== Test Metrics ===")
    for name, m in test_metrics.items():
        print(f"{name}: Accuracy={m.accuracy:.4f}, ROC-AUC={m.roc_auc:.4f}")

    # Test prediction
    predictions = classifier.predict(X_test[:3], current_price=2000.0)
    print("\n=== Sample Predictions ===")
    for pred in predictions:
        print(f"{pred.horizon}_{pred.rr_config}: P(TP)={pred.prob_tp_first:.3f}, "
              f"Action={pred.recommended_action}, Entry={pred.entry_price}, "
              f"SL={pred.sl_price}, TP={pred.tp_price}")
