"""
AMD Detector ML - Machine Learning Based Phase Detection
=========================================================
Trainable classifier for AMD (Accumulation, Manipulation, Distribution) phases.

Combines rule-based heuristics with ML for robust phase detection.

Enhanced with:
- Sample weighting by movement magnitude
- Session and volatility-based attention
- Improved session features (cyclical encoding)

Author: ML-Specialist (NEXUS v4.0)
Created: 2026-01-04
Version: 2.0.0 (2026-01-05) - Added weighting support
"""

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
from enum import IntEnum

try:
    from xgboost import XGBClassifier
    HAS_XGBOOST = True
except ImportError:
    HAS_XGBOOST = False

from sklearn.metrics import (
    accuracy_score, f1_score, classification_report,
    confusion_matrix, precision_recall_fscore_support
)
from sklearn.preprocessing import LabelEncoder

# Import weighting modules
try:
    from ..training.sample_weighting import SampleWeighter, SampleWeightConfig
    from ..training.session_volatility_weighting import (
        SessionVolatilityWeighter, SessionWeightConfig, create_session_features
    )
    HAS_WEIGHTING = True
except ImportError:
    HAS_WEIGHTING = False
    logger.warning("Weighting modules not available, using uniform weights")


class AMDPhaseLabel(IntEnum):
    """AMD Phase labels for classification"""
    UNKNOWN = 0
    ACCUMULATION = 1
    MANIPULATION = 2
    DISTRIBUTION = 3


@dataclass
class AMDPrediction:
    """AMD phase prediction result"""
    phase: str
    phase_label: int
    confidence: float
    probabilities: Dict[str, float]
    features_used: List[str]
    timestamp: Optional[datetime] = None

    def to_dict(self) -> Dict:
        return {
            'phase': self.phase,
            'phase_label': self.phase_label,
            'confidence': float(self.confidence),
            'probabilities': self.probabilities,
            'timestamp': self.timestamp.isoformat() if self.timestamp else None
        }

    @property
    def is_high_confidence(self) -> bool:
        return self.confidence >= 0.7

    @property
    def trading_bias(self) -> str:
        """Get trading bias based on phase"""
        if self.phase == 'accumulation' and self.is_high_confidence:
            return 'LONG_BIAS'
        elif self.phase == 'distribution' and self.is_high_confidence:
            return 'SHORT_BIAS'
        elif self.phase == 'manipulation':
            return 'AVOID'
        return 'NEUTRAL'


@dataclass
class AMDMetrics:
    """Metrics for AMD classification"""
    accuracy: float = 0.0
    macro_f1: float = 0.0
    weighted_f1: float = 0.0
    per_class_f1: Dict[str, float] = field(default_factory=dict)
    confusion_matrix: Optional[np.ndarray] = None
    n_samples: int = 0


class AMDFeatureExtractor:
    """
    Extract features for AMD phase detection.

    Features are designed to capture:
    - Volume patterns
    - Price action characteristics
    - Market structure
    - Smart money footprints
    """

    def __init__(self, lookback: int = 50):
        """
        Initialize feature extractor.

        Args:
            lookback: Number of periods for feature calculation
        """
        self.lookback = lookback

    def extract_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Extract AMD-specific features from OHLCV data.

        Args:
            df: DataFrame with OHLCV data

        Returns:
            DataFrame with extracted features
        """
        features = pd.DataFrame(index=df.index)

        # Volume features
        features = self._add_volume_features(df, features)

        # Price action features
        features = self._add_price_features(df, features)

        # Market structure features
        features = self._add_structure_features(df, features)

        # Order flow features
        features = self._add_order_flow_features(df, features)

        # Volatility features
        features = self._add_volatility_features(df, features)

        # Clean NaN values
        features = features.fillna(0)

        return features

    def _add_volume_features(self, df: pd.DataFrame, features: pd.DataFrame) -> pd.DataFrame:
        """Add volume-based features"""
        if 'volume' not in df.columns:
            return features

        vol = df['volume']

        # Volume moving averages
        features['vol_ma_5'] = vol.rolling(5).mean()
        features['vol_ma_20'] = vol.rolling(20).mean()
        features['vol_ma_50'] = vol.rolling(50).mean()

        # Volume ratios
        features['vol_ratio_5_20'] = features['vol_ma_5'] / (features['vol_ma_20'] + 1e-8)
        features['vol_ratio_20_50'] = features['vol_ma_20'] / (features['vol_ma_50'] + 1e-8)

        # Volume trend
        features['vol_trend'] = vol.rolling(10).mean() - vol.rolling(30).mean()

        # Volume z-score
        vol_std = vol.rolling(20).std()
        features['vol_zscore'] = (vol - features['vol_ma_20']) / (vol_std + 1e-8)

        # Volume spikes
        features['vol_spike'] = (features['vol_zscore'] > 2).astype(float)
        features['vol_spike_count_10'] = features['vol_spike'].rolling(10).sum()

        # Up/Down volume ratio
        price_change = df['close'].diff()
        up_volume = vol.where(price_change > 0, 0)
        down_volume = vol.where(price_change < 0, 0)
        features['up_down_vol_ratio'] = (
            up_volume.rolling(10).sum() /
            (down_volume.rolling(10).sum() + 1e-8)
        )

        return features

    def _add_price_features(self, df: pd.DataFrame, features: pd.DataFrame) -> pd.DataFrame:
        """Add price action features"""
        close = df['close']
        high = df['high']
        low = df['low']

        # Price range
        price_range = high - low
        features['range_pct'] = price_range / (close + 1e-8)
        features['range_ma_10'] = features['range_pct'].rolling(10).mean()
        features['range_compression'] = features['range_pct'] / (features['range_ma_10'] + 1e-8)

        # Price momentum
        features['return_1'] = close.pct_change(1)
        features['return_5'] = close.pct_change(5)
        features['return_10'] = close.pct_change(10)
        features['return_20'] = close.pct_change(20)

        # Price position within range
        features['close_position'] = (close - low) / (high - low + 1e-8)

        # Candle body analysis
        body = abs(close - df['open'])
        features['body_ratio'] = body / (price_range + 1e-8)
        features['upper_wick'] = (high - close.where(close > df['open'], df['open'])) / (price_range + 1e-8)
        features['lower_wick'] = (close.where(close < df['open'], df['open']) - low) / (price_range + 1e-8)

        # Trend features
        features['sma_10'] = close.rolling(10).mean()
        features['sma_20'] = close.rolling(20).mean()
        features['sma_50'] = close.rolling(50).mean()
        features['price_vs_sma_20'] = (close - features['sma_20']) / (features['sma_20'] + 1e-8)
        features['sma_slope'] = features['sma_20'].diff(5) / (features['sma_20'] + 1e-8)

        return features

    def _add_structure_features(self, df: pd.DataFrame, features: pd.DataFrame) -> pd.DataFrame:
        """Add market structure features"""
        high = df['high']
        low = df['low']
        close = df['close']

        # Swing points
        features['swing_high_20'] = high.rolling(20).max()
        features['swing_low_20'] = low.rolling(20).min()

        # Distance from swings
        features['dist_from_high'] = (features['swing_high_20'] - close) / (close + 1e-8)
        features['dist_from_low'] = (close - features['swing_low_20']) / (close + 1e-8)

        # Higher highs / Lower lows count
        features['hh_count'] = (high > high.shift(1)).astype(int).rolling(10).sum()
        features['ll_count'] = (low < low.shift(1)).astype(int).rolling(10).sum()

        # Trend strength
        features['trend_strength'] = features['hh_count'] - features['ll_count']

        # Breakout proximity
        features['near_resistance'] = (close > features['swing_high_20'] * 0.995).astype(float)
        features['near_support'] = (close < features['swing_low_20'] * 1.005).astype(float)

        # Consolidation detection
        range_20 = high.rolling(20).max() - low.rolling(20).min()
        range_5 = high.rolling(5).max() - low.rolling(5).min()
        features['consolidation_ratio'] = range_5 / (range_20 + 1e-8)

        return features

    def _add_order_flow_features(self, df: pd.DataFrame, features: pd.DataFrame) -> pd.DataFrame:
        """Add order flow proxy features"""
        close = df['close']
        high = df['high']
        low = df['low']
        volume = df.get('volume', pd.Series(1, index=df.index))

        # Buying/Selling pressure
        features['buying_pressure'] = (close - low) / (high - low + 1e-8)
        features['selling_pressure'] = (high - close) / (high - low + 1e-8)

        # Weighted buying/selling
        features['buy_pressure_vol'] = features['buying_pressure'] * volume
        features['sell_pressure_vol'] = features['selling_pressure'] * volume

        # Rolling buy/sell balance
        buy_sum = features['buy_pressure_vol'].rolling(10).sum()
        sell_sum = features['sell_pressure_vol'].rolling(10).sum()
        features['order_flow_balance'] = (buy_sum - sell_sum) / (buy_sum + sell_sum + 1e-8)

        # Delta (approximation)
        features['delta_approx'] = features['buying_pressure'] - features['selling_pressure']
        features['cumulative_delta'] = features['delta_approx'].rolling(20).sum()

        return features

    def _add_volatility_features(self, df: pd.DataFrame, features: pd.DataFrame) -> pd.DataFrame:
        """Add volatility features"""
        close = df['close']
        high = df['high']
        low = df['low']

        # ATR
        tr1 = high - low
        tr2 = abs(high - close.shift(1))
        tr3 = abs(low - close.shift(1))
        true_range = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)

        features['atr_14'] = true_range.rolling(14).mean()
        features['atr_ratio'] = features['atr_14'] / (features['atr_14'].rolling(50).mean() + 1e-8)

        # Historical volatility
        returns = close.pct_change()
        features['volatility_10'] = returns.rolling(10).std()
        features['volatility_20'] = returns.rolling(20).std()
        features['vol_expansion'] = features['volatility_10'] / (features['volatility_20'] + 1e-8)

        # Bollinger Band width
        sma_20 = close.rolling(20).mean()
        std_20 = close.rolling(20).std()
        features['bb_width'] = (2 * std_20) / (sma_20 + 1e-8)

        return features


class AMDLabelGenerator:
    """
    Generate labels for AMD phase training.

    Uses forward-looking price action to label historical data.
    This is for training only - not for real-time use.
    """

    def __init__(
        self,
        forward_periods: int = 20,
        volatility_threshold: float = 1.5,
        trend_threshold: float = 0.02
    ):
        """
        Initialize label generator.

        Args:
            forward_periods: Periods to look forward for labeling
            volatility_threshold: ATR ratio threshold for manipulation
            trend_threshold: Return threshold for trend detection
        """
        self.forward_periods = forward_periods
        self.volatility_threshold = volatility_threshold
        self.trend_threshold = trend_threshold

    def generate_labels(self, df: pd.DataFrame) -> pd.Series:
        """
        Generate AMD phase labels based on forward price action.

        Uses percentile-based thresholds for balanced label distribution.

        Args:
            df: DataFrame with OHLCV data

        Returns:
            Series with AMD phase labels (0-3)
        """
        labels = pd.Series(AMDPhaseLabel.UNKNOWN, index=df.index)

        close = df['close']
        high = df['high']
        low = df['low']
        volume = df.get('volume', pd.Series(1, index=df.index))

        # Calculate forward metrics
        forward_return = close.shift(-self.forward_periods) / close - 1
        forward_high = high.rolling(self.forward_periods).max().shift(-self.forward_periods)
        forward_low = low.rolling(self.forward_periods).min().shift(-self.forward_periods)
        forward_range = (forward_high - forward_low) / close

        # Current volatility (shorter window for more sensitivity)
        returns = close.pct_change()
        current_vol = returns.rolling(5).std()
        avg_vol = returns.rolling(30).std()
        vol_ratio = current_vol / (avg_vol + 1e-8)

        # Volume patterns
        vol_ma = volume.rolling(20).mean()
        vol_ratio_cur = volume / (vol_ma + 1e-8)

        # Range compression
        range_cur = (high - low) / close
        range_ma = range_cur.rolling(15).mean()
        range_ratio = range_cur / (range_ma + 1e-8)

        # Price position in recent range
        rolling_high = high.rolling(20).max()
        rolling_low = low.rolling(20).min()
        price_position = (close - rolling_low) / (rolling_high - rolling_low + 1e-8)

        # Calculate percentile thresholds for balanced labeling
        fwd_ret_valid = forward_return.dropna()
        fwd_ret_p75 = fwd_ret_valid.quantile(0.75)  # Top 25% moves up
        fwd_ret_p25 = fwd_ret_valid.quantile(0.25)  # Bottom 25% moves down

        range_ratio_p30 = range_ratio.dropna().quantile(0.30)  # Range compression
        vol_ratio_p70 = vol_ratio.dropna().quantile(0.70)  # High volatility

        for i in range(50, len(df) - self.forward_periods):
            # Get current metrics
            fwd_ret = forward_return.iloc[i]
            fwd_rng = forward_range.iloc[i]
            v_ratio = vol_ratio.iloc[i]
            r_ratio = range_ratio.iloc[i]
            p_position = price_position.iloc[i]
            vr_cur = vol_ratio_cur.iloc[i]

            if pd.isna(fwd_ret) or pd.isna(fwd_rng) or pd.isna(v_ratio):
                continue

            # ACCUMULATION: Range compression + price near lows + followed by upward move
            if (r_ratio < range_ratio_p30 and  # Range compression (bottom 30%)
                p_position < 0.4 and  # Price near recent lows
                fwd_ret > fwd_ret_p75 * 0.5):  # Followed by decent up move
                labels.iloc[i] = AMDPhaseLabel.ACCUMULATION

            # DISTRIBUTION: Range compression + price near highs + followed by downward move
            elif (r_ratio < range_ratio_p30 and  # Range compression
                  p_position > 0.6 and  # Price near recent highs
                  fwd_ret < fwd_ret_p25 * 0.5):  # Followed by decent down move
                labels.iloc[i] = AMDPhaseLabel.DISTRIBUTION

            # MANIPULATION: High volatility with no clear direction
            elif (v_ratio > vol_ratio_p70 and  # High volatility
                  abs(fwd_ret) < abs(fwd_ret_p75) * 0.3):  # No significant move
                labels.iloc[i] = AMDPhaseLabel.MANIPULATION

            # Strong bullish move -> Label previous consolidation as accumulation
            elif fwd_ret > fwd_ret_p75:
                if r_ratio < 1.0:  # Some consolidation
                    labels.iloc[i] = AMDPhaseLabel.ACCUMULATION
                else:
                    labels.iloc[i] = AMDPhaseLabel.MANIPULATION  # Trending, not consolidation

            # Strong bearish move -> Label previous consolidation as distribution
            elif fwd_ret < fwd_ret_p25:
                if r_ratio < 1.0:  # Some consolidation
                    labels.iloc[i] = AMDPhaseLabel.DISTRIBUTION
                else:
                    labels.iloc[i] = AMDPhaseLabel.MANIPULATION

        # Log distribution
        label_counts = labels.value_counts()
        logger.info(f"Label distribution: {label_counts.to_dict()}")

        return labels


class AMDDetectorML:
    """
    Machine Learning based AMD Phase Detector.

    Combines extracted features with XGBoost classification
    for robust phase detection.
    """

    PHASE_NAMES = {
        0: 'unknown',
        1: 'accumulation',
        2: 'manipulation',
        3: 'distribution'
    }

    def __init__(
        self,
        lookback: int = 50,
        config: Optional[Dict[str, Any]] = None,
        use_gpu: bool = True,
        use_sample_weighting: bool = True,
        use_session_weighting: bool = False,  # Disabled by default - only ATR volatility weighting
        sample_weight_config: Optional[Dict] = None,
        session_weight_config: Optional[Dict] = None,
    ):
        """
        Initialize AMD ML Detector.

        Args:
            lookback: Periods for feature extraction
            config: Model configuration
            use_gpu: Use GPU acceleration
            use_sample_weighting: Enable movement-based sample weighting
            use_session_weighting: Enable session/hour-based weighting (disabled by default)
            sample_weight_config: Configuration for SampleWeighter
            session_weight_config: Configuration for SessionVolatilityWeighter (uses ATR by default)
        """
        self.lookback = lookback
        self.config = config or self._default_config()
        self.use_gpu = use_gpu

        self.feature_extractor = AMDFeatureExtractor(lookback)
        self.label_generator = AMDLabelGenerator()
        self.model: Optional[XGBClassifier] = None
        self.feature_columns: List[str] = []
        self.metrics: Optional[AMDMetrics] = None
        self._is_trained = False

        # Weighting configuration
        self.use_sample_weighting = use_sample_weighting and HAS_WEIGHTING
        self.use_session_weighting = use_session_weighting and HAS_WEIGHTING

        # Initialize weighters if available
        if HAS_WEIGHTING:
            sample_cfg = sample_weight_config or {}
            if isinstance(sample_cfg, dict):
                sample_cfg = SampleWeightConfig(**sample_cfg)
            self.sample_weighter = SampleWeighter(sample_cfg)

            session_cfg = session_weight_config or {}
            if isinstance(session_cfg, dict):
                session_cfg = SessionWeightConfig(**session_cfg)
            self.session_weighter = SessionVolatilityWeighter(session_cfg)

            logger.info(f"AMD Detector - Sample weighting: {self.use_sample_weighting}, Session weighting: {self.use_session_weighting}")
        else:
            self.sample_weighter = None
            self.session_weighter = None

        self._init_model()

    def _default_config(self) -> Dict[str, Any]:
        """Default XGBoost configuration"""
        return {
            '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,
            'objective': 'multi:softprob',
            'num_class': 4,
            'tree_method': 'hist',
            'random_state': 42
        }

    def _init_model(self):
        """Initialize XGBoost classifier"""
        if not HAS_XGBOOST:
            raise ImportError("XGBoost is required for AMDDetectorML")

        params = self.config.copy()

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

        self.model = XGBClassifier(**params)

    def compute_sample_weights(
        self,
        df: pd.DataFrame,
        labels: pd.Series
    ) -> np.ndarray:
        """
        Compute combined sample weights for AMD classification.

        For AMD detection, we weight samples based on:
        1. Session (London/NY overlap more important)
        2. Volatility (higher ATR = more relevant)
        3. Label clarity (higher forward movement = clearer signal)

        Args:
            df: DataFrame with OHLCV data
            labels: Series with AMD phase labels

        Returns:
            Array of sample weights
        """
        n_samples = len(df)
        weights = np.ones(n_samples)

        # 1. Session and volatility weights
        if self.use_session_weighting and self.session_weighter is not None:
            session_weights = self.session_weighter.compute_combined_weights(df)
            # Align weights with labels index
            if len(session_weights) == len(df):
                # Map session weights to valid labels
                valid_indices = labels.index
                if isinstance(df.index, pd.DatetimeIndex):
                    session_df = pd.Series(session_weights, index=df.index)
                    session_weights = session_df.reindex(valid_indices).fillna(1.0).values
            weights = session_weights
            logger.info(f"Applied session weights: mean={weights.mean():.3f}")

        # 2. Movement-based weights (for samples with forward movement data)
        if self.use_sample_weighting and self.sample_weighter is not None:
            # Calculate forward movement for weighting
            close = df['close']
            high = df['high']
            low = df['low']
            forward_periods = self.label_generator.forward_periods

            # Forward high/low moves
            target_high = high.rolling(forward_periods).max().shift(-forward_periods) - close
            target_low = close - low.rolling(forward_periods).min().shift(-forward_periods)

            # Create temp df for weighter
            df_temp = df.copy()
            df_temp['target_high'] = target_high
            df_temp['target_low'] = target_low

            # Get movement weights
            movement_weights, valid_mask = self.sample_weighter.compute_sample_weights(
                df_temp, 'target_high', 'target_low'
            )

            # Combine with existing weights
            weights = weights * movement_weights
            logger.info(f"Applied movement weights: mean={movement_weights.mean():.3f}")

        # Normalize to mean=1
        if weights.mean() > 0:
            weights = weights / weights.mean()

        logger.info(f"Combined weights: min={weights.min():.3f}, max={weights.max():.3f}, mean={weights.mean():.3f}")

        return weights

    def prepare_training_data(
        self,
        df: pd.DataFrame
    ) -> Tuple[pd.DataFrame, pd.Series]:
        """
        Prepare features and labels for training.

        Args:
            df: Raw OHLCV DataFrame

        Returns:
            Tuple of (features, labels)
        """
        logger.info(f"Preparing training data from {len(df)} samples")

        # Extract features
        features = self.feature_extractor.extract_features(df)
        self.feature_columns = features.columns.tolist()

        # Generate labels
        labels = self.label_generator.generate_labels(df)

        # Remove unknown labels
        valid_mask = labels != AMDPhaseLabel.UNKNOWN
        features = features[valid_mask]
        labels = labels[valid_mask]

        # Remove NaN
        valid_mask = features.notna().all(axis=1)
        features = features[valid_mask]
        labels = labels[valid_mask]

        logger.info(f"Prepared {len(features)} samples with valid labels")
        logger.info(f"Label distribution: {labels.value_counts().to_dict()}")

        return features, labels

    def train(
        self,
        df_train: pd.DataFrame,
        df_val: Optional[pd.DataFrame] = None,
        verbose: bool = True,
        sample_weight: Optional[np.ndarray] = None
    ) -> AMDMetrics:
        """
        Train the AMD detector.

        Args:
            df_train: Training OHLCV data
            df_val: Validation OHLCV data (optional)
            verbose: Print training progress
            sample_weight: Pre-computed sample weights (optional, auto-computed if None)

        Returns:
            Training metrics
        """
        logger.info("Training AMD Detector ML...")

        # Add enhanced session features if weighting module available
        if HAS_WEIGHTING and isinstance(df_train.index, pd.DatetimeIndex):
            df_train = create_session_features(df_train)

        # Prepare training data
        X_train, y_train = self.prepare_training_data(df_train)

        # Remap labels to start from 0 for XGBoost compatibility
        unique_labels = sorted(y_train.unique())
        self._label_mapping = {old: new for new, old in enumerate(unique_labels)}
        self._reverse_mapping = {new: old for old, new in self._label_mapping.items()}
        y_train = y_train.map(self._label_mapping)
        logger.info(f"Label mapping: {self._label_mapping}")

        # Reinitialize model with correct number of classes
        n_classes = len(unique_labels)
        if n_classes > 2:
            self.config['num_class'] = n_classes
            self.config['objective'] = 'multi:softmax'
        else:
            # For binary classification, remove num_class
            self.config.pop('num_class', None)
            self.config['objective'] = 'binary:logistic'
        self._init_model()
        logger.info(f"Model configured for {n_classes} classes")

        # Compute sample weights if not provided
        if sample_weight is not None:
            weights = sample_weight
        elif self.use_sample_weighting or self.use_session_weighting:
            # Need to align weights with the filtered training data
            weights = self.compute_sample_weights(df_train, y_train)
            # Slice to match X_train length
            if len(weights) != len(X_train):
                # Recompute with aligned data
                weights_df = pd.DataFrame({'weight': weights}, index=df_train.index)
                weights = weights_df.loc[X_train.index, 'weight'].values
        else:
            weights = None

        # Prepare fit params
        fit_params = {}
        if weights is not None:
            fit_params['sample_weight'] = weights
            logger.info(f"Training with sample weights (weighted)" if weights is not None else "Training without weights")

        # Prepare validation data if provided
        if df_val is not None:
            if HAS_WEIGHTING and isinstance(df_val.index, pd.DatetimeIndex):
                df_val = create_session_features(df_val)
            X_val, y_val = self.prepare_training_data(df_val)
            if len(X_val) > 0:
                y_val = y_val.map(lambda x: self._label_mapping.get(x, 0))
                fit_params['eval_set'] = [(X_val.values, y_val.values)]

        # Train model
        self.model.fit(X_train.values, y_train.values, **fit_params)
        self._is_trained = True

        # Calculate metrics
        y_pred = self.model.predict(X_train.values)
        self.metrics = self._calculate_metrics(y_train.values, y_pred)

        if verbose:
            self._print_metrics(self.metrics, "Training")

        # Validation metrics
        if df_val is not None and len(X_val) > 0:
            y_val_pred = self.model.predict(X_val.values)
            val_metrics = self._calculate_metrics(y_val.values, y_val_pred)
            if verbose:
                self._print_metrics(val_metrics, "Validation")

        return self.metrics

    def predict(
        self,
        df: pd.DataFrame
    ) -> List[AMDPrediction]:
        """
        Predict AMD phases for given data.

        Args:
            df: OHLCV DataFrame

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

        # Extract features
        features = self.feature_extractor.extract_features(df)

        # Ensure same columns as training
        for col in self.feature_columns:
            if col not in features.columns:
                features[col] = 0
        features = features[self.feature_columns]

        # Predict
        X = features.values
        y_pred = self.model.predict(X)
        y_proba = self.model.predict_proba(X)

        predictions = []
        for i in range(len(X)):
            phase_label = int(y_pred[i])
            phase_name = self.PHASE_NAMES.get(phase_label, 'unknown')
            confidence = float(y_proba[i].max())

            probs = {
                self.PHASE_NAMES[j]: float(y_proba[i][j])
                for j in range(4)
            }

            pred = AMDPrediction(
                phase=phase_name,
                phase_label=phase_label,
                confidence=confidence,
                probabilities=probs,
                features_used=self.feature_columns,
                timestamp=df.index[i] if isinstance(df.index, pd.DatetimeIndex) else None
            )
            predictions.append(pred)

        return predictions

    def predict_single(self, df: pd.DataFrame) -> AMDPrediction:
        """Predict for the most recent bar"""
        predictions = self.predict(df.tail(self.lookback + 1))
        return predictions[-1] if predictions else None

    def _calculate_metrics(
        self,
        y_true: np.ndarray,
        y_pred: np.ndarray
    ) -> AMDMetrics:
        """Calculate classification metrics"""
        # Ensure arrays are 1D
        y_true = np.asarray(y_true).ravel()
        y_pred = np.asarray(y_pred).ravel()

        # Use actual labels present in data
        unique_labels = sorted(set(y_true.tolist()) | set(y_pred.tolist()))

        precision, recall, f1, _ = precision_recall_fscore_support(
            y_true, y_pred, average=None, labels=unique_labels, zero_division=0
        )

        # Map back to original phase names if we have reverse mapping
        per_class_f1 = {}
        for i, label in enumerate(unique_labels):
            if hasattr(self, '_reverse_mapping') and label in self._reverse_mapping:
                orig_label = self._reverse_mapping[label]
                phase_name = self.PHASE_NAMES.get(orig_label, f"Class_{label}")
            else:
                phase_name = self.PHASE_NAMES.get(label, f"Class_{label}")
            per_class_f1[phase_name] = float(f1[i])

        return AMDMetrics(
            accuracy=accuracy_score(y_true, y_pred),
            macro_f1=f1_score(y_true, y_pred, average='macro', zero_division=0),
            weighted_f1=f1_score(y_true, y_pred, average='weighted', zero_division=0),
            per_class_f1=per_class_f1,
            confusion_matrix=confusion_matrix(y_true, y_pred),
            n_samples=len(y_true)
        )

    def _print_metrics(self, metrics: AMDMetrics, prefix: str = ""):
        """Print metrics summary"""
        logger.info(f"\n{prefix} Metrics:")
        logger.info(f"  Accuracy: {metrics.accuracy:.2%}")
        logger.info(f"  Macro F1: {metrics.macro_f1:.4f}")
        logger.info(f"  Weighted F1: {metrics.weighted_f1:.4f}")
        logger.info(f"  Per-class F1:")
        for phase, f1 in metrics.per_class_f1.items():
            logger.info(f"    {phase}: {f1:.4f}")

    def get_feature_importance(self, top_n: int = 20) -> Dict[str, float]:
        """Get feature importance"""
        if not self._is_trained:
            return {}

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

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

        joblib.dump(self.model, path / 'amd_model.joblib')

        metadata = {
            'config': self.config,
            'lookback': self.lookback,
            'feature_columns': self.feature_columns,
            'metrics': vars(self.metrics) if self.metrics else None,
            'version': '1.0'
        }
        joblib.dump(metadata, path / 'metadata.joblib')

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

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

        self.model = joblib.load(path / 'amd_model.joblib')

        metadata = joblib.load(path / 'metadata.joblib')
        self.config = metadata['config']
        self.lookback = metadata['lookback']
        self.feature_columns = metadata['feature_columns']

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


if __name__ == "__main__":
    # Test AMD Detector ML
    np.random.seed(42)

    # Create sample OHLCV data
    n_samples = 2000
    dates = pd.date_range('2023-01-01', periods=n_samples, freq='5min')

    price = 2000 + np.random.randn(n_samples).cumsum() * 0.5
    df = pd.DataFrame({
        'open': price + np.random.randn(n_samples) * 0.5,
        'high': price + abs(np.random.randn(n_samples)) * 1.0,
        'low': price - abs(np.random.randn(n_samples)) * 1.0,
        'close': price + np.random.randn(n_samples) * 0.5,
        'volume': np.random.randint(100, 10000, n_samples)
    }, index=dates)

    # Ensure high > low
    df['high'] = df[['open', 'high', 'close']].max(axis=1)
    df['low'] = df[['open', 'low', 'close']].min(axis=1)

    # Split data
    train_size = 1600
    df_train = df.iloc[:train_size]
    df_val = df.iloc[train_size:]

    # Train
    detector = AMDDetectorML(use_gpu=False)
    metrics = detector.train(df_train, df_val)

    # Predict
    predictions = detector.predict(df_val.tail(100))
    print(f"\n=== Sample Predictions ===")
    for pred in predictions[-5:]:
        print(f"{pred.timestamp}: {pred.phase} (conf={pred.confidence:.2f}) -> {pred.trading_bias}")

    # Feature importance
    print(f"\n=== Top Features ===")
    for feat, imp in detector.get_feature_importance(10).items():
        print(f"  {feat}: {imp:.4f}")
