#!/usr/bin/env python3
"""
Dual Horizon Ensemble Model
===========================
Combines long-term structural patterns with short-term market adaptation.

Architecture:
- Long-term model: Trained on 5 years of data (structural patterns, seasonality)
- Short-term model: Trained on 3 months of data (current market regime)
- Dynamic weighting: Adjusts based on recent performance of each model

Key Benefits:
- Long-term captures: macro patterns, historical support/resistance, seasonality
- Short-term captures: current volatility regime, recent price behavior
- Ensemble reduces overfitting to either time horizon

Author: Trading Strategist + ML Specialist
Version: 1.0.0
"""

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


@dataclass
class DualHorizonConfig:
    """Configuration for dual horizon ensemble"""

    # Data horizons
    long_term_years: float = 5.0       # Years of data for long-term model
    short_term_months: float = 3.0     # Months of data for short-term model

    # Model weights
    initial_long_weight: float = 0.6   # Initial weight for long-term model
    initial_short_weight: float = 0.4  # Initial weight for short-term model

    # Dynamic weight adjustment
    use_dynamic_weights: bool = True
    weight_adjustment_lookback: int = 100  # Samples to evaluate performance
    weight_adjustment_rate: float = 0.1    # How much to adjust weights
    min_weight: float = 0.2                # Minimum weight for either model
    max_weight: float = 0.8                # Maximum weight for either model

    # Performance tracking
    performance_window: int = 50       # Rolling window for performance metrics

    # Short-term retraining
    retrain_frequency_days: int = 7    # Retrain short-term model every N days
    min_samples_for_retrain: int = 500 # Minimum samples for short-term retrain

    # XGBoost parameters (shared base, can be overridden)
    xgb_params: Dict = field(default_factory=lambda: {
        'objective': 'reg:squarederror',
        'max_depth': 6,
        'learning_rate': 0.05,
        'n_estimators': 200,
        'min_child_weight': 10,
        'subsample': 0.8,
        'colsample_bytree': 0.8,
        'reg_alpha': 0.1,
        'reg_lambda': 1.0,
        'random_state': 42
    })

    # Long-term specific adjustments
    long_term_depth: int = 8           # Deeper for complex patterns
    long_term_estimators: int = 300    # More trees for diversity

    # Short-term specific adjustments
    short_term_depth: int = 5          # Shallower to avoid overfitting
    short_term_estimators: int = 150   # Fewer trees for faster adaptation
    short_term_learning_rate: float = 0.08  # Faster learning


@dataclass
class ModelPerformance:
    """Tracks performance metrics for a model"""
    predictions: List[float] = field(default_factory=list)
    actuals: List[float] = field(default_factory=list)
    timestamps: List[datetime] = field(default_factory=list)
    mae_rolling: float = 0.0
    mse_rolling: float = 0.0
    direction_accuracy: float = 0.0


class DualHorizonEnsemble:
    """
    Dual horizon ensemble that combines long-term and short-term models.

    Strategy:
    1. Long-term model captures structural patterns from years of data
    2. Short-term model adapts to current market regime
    3. Dynamic weighting adjusts based on recent accuracy
    4. Short-term model is periodically retrained for adaptation

    Usage:
        ensemble = DualHorizonEnsemble(DualHorizonConfig())
        ensemble.fit(X_train, y_train, timestamps_train, sample_weights)

        # For predictions
        pred_high, pred_low = ensemble.predict(X_test)

        # Periodic short-term retrain
        if ensemble.should_retrain():
            ensemble.retrain_short_term(X_recent, y_recent, timestamps_recent)
    """

    def __init__(self, config: DualHorizonConfig = None):
        self.config = config or DualHorizonConfig()

        # Initialize models
        self.long_term_high: Optional[xgb.XGBRegressor] = None
        self.long_term_low: Optional[xgb.XGBRegressor] = None
        self.short_term_high: Optional[xgb.XGBRegressor] = None
        self.short_term_low: Optional[xgb.XGBRegressor] = None

        # Current weights
        self.weight_long = self.config.initial_long_weight
        self.weight_short = self.config.initial_short_weight

        # Performance tracking
        self.long_term_perf = ModelPerformance()
        self.short_term_perf = ModelPerformance()

        # Training metadata
        self.last_retrain: Optional[datetime] = None
        self.feature_names: Optional[List[str]] = None
        self.is_fitted = False

    def _get_long_term_params(self) -> Dict:
        """Get XGBoost parameters for long-term model."""
        params = self.config.xgb_params.copy()
        params['max_depth'] = self.config.long_term_depth
        params['n_estimators'] = self.config.long_term_estimators
        return params

    def _get_short_term_params(self) -> Dict:
        """Get XGBoost parameters for short-term model."""
        params = self.config.xgb_params.copy()
        params['max_depth'] = self.config.short_term_depth
        params['n_estimators'] = self.config.short_term_estimators
        params['learning_rate'] = self.config.short_term_learning_rate
        return params

    def _split_by_horizon(
        self,
        X: np.ndarray,
        y_high: np.ndarray,
        y_low: np.ndarray,
        timestamps: pd.DatetimeIndex,
        sample_weights: np.ndarray = None
    ) -> Tuple[Dict, Dict]:
        """
        Split data into long-term and short-term sets.

        Returns:
            Tuple of (long_term_data, short_term_data) dicts
        """
        current_time = timestamps.max()

        # Calculate cutoff dates
        long_term_start = current_time - timedelta(days=self.config.long_term_years * 365)
        short_term_start = current_time - timedelta(days=self.config.short_term_months * 30)

        # Create masks
        long_term_mask = timestamps >= long_term_start
        short_term_mask = timestamps >= short_term_start

        # Handle weights
        if sample_weights is None:
            sample_weights = np.ones(len(X))

        long_term_data = {
            'X': X[long_term_mask],
            'y_high': y_high[long_term_mask],
            'y_low': y_low[long_term_mask],
            'weights': sample_weights[long_term_mask],
            'timestamps': timestamps[long_term_mask]
        }

        short_term_data = {
            'X': X[short_term_mask],
            'y_high': y_high[short_term_mask],
            'y_low': y_low[short_term_mask],
            'weights': sample_weights[short_term_mask],
            'timestamps': timestamps[short_term_mask]
        }

        logger.info(f"Data split:")
        logger.info(f"  Long-term: {len(long_term_data['X'])} samples "
                   f"({long_term_start.strftime('%Y-%m-%d')} to {current_time.strftime('%Y-%m-%d')})")
        logger.info(f"  Short-term: {len(short_term_data['X'])} samples "
                   f"({short_term_start.strftime('%Y-%m-%d')} to {current_time.strftime('%Y-%m-%d')})")

        return long_term_data, short_term_data

    def fit(
        self,
        X: np.ndarray,
        y_high: np.ndarray,
        y_low: np.ndarray,
        timestamps: pd.DatetimeIndex,
        sample_weights: np.ndarray = None,
        feature_names: List[str] = None
    ) -> 'DualHorizonEnsemble':
        """
        Fit both long-term and short-term models.

        Args:
            X: Feature matrix
            y_high: High targets (in multiplier or USD)
            y_low: Low targets (in multiplier or USD)
            timestamps: Datetime index for each sample
            sample_weights: Optional sample weights
            feature_names: Optional feature names

        Returns:
            Self for chaining
        """
        self.feature_names = feature_names

        # Split data by horizon
        long_term_data, short_term_data = self._split_by_horizon(
            X, y_high, y_low, timestamps, sample_weights
        )

        # Train long-term models
        logger.info("Training long-term models...")
        long_params = self._get_long_term_params()

        self.long_term_high = xgb.XGBRegressor(**long_params)
        self.long_term_high.fit(
            long_term_data['X'],
            long_term_data['y_high'],
            sample_weight=long_term_data['weights']
        )

        self.long_term_low = xgb.XGBRegressor(**long_params)
        self.long_term_low.fit(
            long_term_data['X'],
            long_term_data['y_low'],
            sample_weight=long_term_data['weights']
        )

        # Train short-term models
        logger.info("Training short-term models...")
        short_params = self._get_short_term_params()

        self.short_term_high = xgb.XGBRegressor(**short_params)
        self.short_term_high.fit(
            short_term_data['X'],
            short_term_data['y_high'],
            sample_weight=short_term_data['weights']
        )

        self.short_term_low = xgb.XGBRegressor(**short_params)
        self.short_term_low.fit(
            short_term_data['X'],
            short_term_data['y_low'],
            sample_weight=short_term_data['weights']
        )

        self.last_retrain = datetime.now()
        self.is_fitted = True

        logger.info("Dual horizon ensemble training complete")
        logger.info(f"  Long-term weight: {self.weight_long:.2f}")
        logger.info(f"  Short-term weight: {self.weight_short:.2f}")

        return self

    def predict(
        self,
        X: np.ndarray,
        return_components: bool = False
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Predict high and low targets using weighted ensemble.

        Args:
            X: Feature matrix
            return_components: If True, return individual model predictions

        Returns:
            Tuple of (predicted_high, predicted_low)
            If return_components: Also returns dict with individual predictions
        """
        if not self.is_fitted:
            raise ValueError("Model not fitted. Call fit() first.")

        # Get predictions from both models
        long_high = self.long_term_high.predict(X)
        long_low = self.long_term_low.predict(X)
        short_high = self.short_term_high.predict(X)
        short_low = self.short_term_low.predict(X)

        # Weighted ensemble
        pred_high = self.weight_long * long_high + self.weight_short * short_high
        pred_low = self.weight_long * long_low + self.weight_short * short_low

        # Ensure non-negative predictions
        pred_high = np.maximum(pred_high, 0)
        pred_low = np.maximum(pred_low, 0)

        if return_components:
            components = {
                'long_term_high': long_high,
                'long_term_low': long_low,
                'short_term_high': short_high,
                'short_term_low': short_low,
                'weight_long': self.weight_long,
                'weight_short': self.weight_short
            }
            return pred_high, pred_low, components

        return pred_high, pred_low

    def update_performance(
        self,
        X: np.ndarray,
        y_high_actual: np.ndarray,
        y_low_actual: np.ndarray,
        timestamps: List[datetime] = None
    ) -> None:
        """
        Update performance tracking with recent predictions and actuals.

        Args:
            X: Feature matrix for samples
            y_high_actual: Actual high targets
            y_low_actual: Actual low targets
            timestamps: Optional timestamps
        """
        # Get individual predictions
        pred_high, pred_low, components = self.predict(X, return_components=True)

        # Calculate errors for each model (combined high + low)
        long_pred = components['long_term_high'] + components['long_term_low']
        short_pred = components['short_term_high'] + components['short_term_low']
        actual_total = y_high_actual + y_low_actual

        # Update long-term performance
        for i, (lp, sp, actual) in enumerate(zip(long_pred, short_pred, actual_total)):
            self.long_term_perf.predictions.append(lp)
            self.long_term_perf.actuals.append(actual)
            self.short_term_perf.predictions.append(sp)
            self.short_term_perf.actuals.append(actual)

            if timestamps:
                self.long_term_perf.timestamps.append(timestamps[i])
                self.short_term_perf.timestamps.append(timestamps[i])

        # Keep only recent performance data
        window = self.config.performance_window
        if len(self.long_term_perf.predictions) > window * 2:
            self.long_term_perf.predictions = self.long_term_perf.predictions[-window:]
            self.long_term_perf.actuals = self.long_term_perf.actuals[-window:]
            self.short_term_perf.predictions = self.short_term_perf.predictions[-window:]
            self.short_term_perf.actuals = self.short_term_perf.actuals[-window:]

        # Calculate rolling metrics
        if len(self.long_term_perf.predictions) >= 10:
            self._update_rolling_metrics()

    def _update_rolling_metrics(self) -> None:
        """Update rolling performance metrics."""
        window = min(self.config.performance_window, len(self.long_term_perf.predictions))

        # Long-term metrics
        long_preds = np.array(self.long_term_perf.predictions[-window:])
        long_actuals = np.array(self.long_term_perf.actuals[-window:])
        self.long_term_perf.mae_rolling = np.mean(np.abs(long_preds - long_actuals))
        self.long_term_perf.mse_rolling = np.mean((long_preds - long_actuals) ** 2)

        # Short-term metrics
        short_preds = np.array(self.short_term_perf.predictions[-window:])
        short_actuals = np.array(self.short_term_perf.actuals[-window:])
        self.short_term_perf.mae_rolling = np.mean(np.abs(short_preds - short_actuals))
        self.short_term_perf.mse_rolling = np.mean((short_preds - short_actuals) ** 2)

    def adjust_weights(self) -> Tuple[float, float]:
        """
        Dynamically adjust weights based on recent performance.

        Better performing model gets higher weight.

        Returns:
            Tuple of (new_long_weight, new_short_weight)
        """
        if not self.config.use_dynamic_weights:
            return self.weight_long, self.weight_short

        if len(self.long_term_perf.predictions) < self.config.weight_adjustment_lookback:
            logger.debug("Not enough data for weight adjustment")
            return self.weight_long, self.weight_short

        # Compare MAE (lower is better)
        long_mae = self.long_term_perf.mae_rolling
        short_mae = self.short_term_perf.mae_rolling

        if long_mae == 0 and short_mae == 0:
            return self.weight_long, self.weight_short

        # Calculate relative performance
        total_mae = long_mae + short_mae
        if total_mae > 0:
            # Better model (lower MAE) gets higher weight
            long_score = 1 - (long_mae / total_mae)
            short_score = 1 - (short_mae / total_mae)

            # Apply adjustment rate
            rate = self.config.weight_adjustment_rate
            new_long = self.weight_long + rate * (long_score - 0.5)
            new_short = self.weight_short + rate * (short_score - 0.5)

            # Normalize to sum to 1
            total = new_long + new_short
            new_long = new_long / total
            new_short = new_short / total

            # Apply min/max constraints
            new_long = np.clip(new_long, self.config.min_weight, self.config.max_weight)
            new_short = np.clip(new_short, self.config.min_weight, self.config.max_weight)

            # Re-normalize after clipping
            total = new_long + new_short
            self.weight_long = new_long / total
            self.weight_short = new_short / total

            logger.info(f"Weights adjusted - Long: {self.weight_long:.3f}, Short: {self.weight_short:.3f}")
            logger.info(f"  Long MAE: {long_mae:.4f}, Short MAE: {short_mae:.4f}")

        return self.weight_long, self.weight_short

    def should_retrain(self) -> bool:
        """Check if short-term model should be retrained."""
        if self.last_retrain is None:
            return True

        days_since_retrain = (datetime.now() - self.last_retrain).days
        return days_since_retrain >= self.config.retrain_frequency_days

    def retrain_short_term(
        self,
        X: np.ndarray,
        y_high: np.ndarray,
        y_low: np.ndarray,
        sample_weights: np.ndarray = None
    ) -> None:
        """
        Retrain only the short-term model with recent data.

        Args:
            X: Recent feature matrix
            y_high: Recent high targets
            y_low: Recent low targets
            sample_weights: Optional sample weights
        """
        if len(X) < self.config.min_samples_for_retrain:
            logger.warning(f"Insufficient samples for retrain: {len(X)} < {self.config.min_samples_for_retrain}")
            return

        logger.info(f"Retraining short-term model with {len(X)} samples...")

        short_params = self._get_short_term_params()

        self.short_term_high = xgb.XGBRegressor(**short_params)
        self.short_term_high.fit(X, y_high, sample_weight=sample_weights)

        self.short_term_low = xgb.XGBRegressor(**short_params)
        self.short_term_low.fit(X, y_low, sample_weight=sample_weights)

        self.last_retrain = datetime.now()

        # Reset short-term performance tracking
        self.short_term_perf = ModelPerformance()

        logger.info("Short-term model retrained successfully")

    def get_feature_importance(
        self,
        model_type: str = 'combined'
    ) -> pd.DataFrame:
        """
        Get feature importance from models.

        Args:
            model_type: 'long_term', 'short_term', or 'combined'

        Returns:
            DataFrame with feature importances
        """
        if not self.is_fitted:
            raise ValueError("Model not fitted")

        importance_data = {}

        if model_type in ['long_term', 'combined']:
            lt_high_imp = self.long_term_high.feature_importances_
            lt_low_imp = self.long_term_low.feature_importances_
            importance_data['long_term_avg'] = (lt_high_imp + lt_low_imp) / 2

        if model_type in ['short_term', 'combined']:
            st_high_imp = self.short_term_high.feature_importances_
            st_low_imp = self.short_term_low.feature_importances_
            importance_data['short_term_avg'] = (st_high_imp + st_low_imp) / 2

        if model_type == 'combined':
            importance_data['combined'] = (
                self.weight_long * importance_data['long_term_avg'] +
                self.weight_short * importance_data['short_term_avg']
            )

        df = pd.DataFrame(importance_data)

        if self.feature_names:
            df['feature'] = self.feature_names
            df = df.set_index('feature')

        return df.sort_values(by=list(importance_data.keys())[0], ascending=False)

    def get_model_summary(self) -> Dict[str, Any]:
        """Get summary of model configuration and performance."""
        return {
            'is_fitted': self.is_fitted,
            'weight_long': self.weight_long,
            'weight_short': self.weight_short,
            'last_retrain': self.last_retrain.isoformat() if self.last_retrain else None,
            'long_term_years': self.config.long_term_years,
            'short_term_months': self.config.short_term_months,
            'performance': {
                'long_term_mae': self.long_term_perf.mae_rolling,
                'short_term_mae': self.short_term_perf.mae_rolling,
                'samples_tracked': len(self.long_term_perf.predictions)
            }
        }

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

        # Save models
        joblib.dump(self.long_term_high, path / 'long_term_high.joblib')
        joblib.dump(self.long_term_low, path / 'long_term_low.joblib')
        joblib.dump(self.short_term_high, path / 'short_term_high.joblib')
        joblib.dump(self.short_term_low, path / 'short_term_low.joblib')

        # Save metadata
        metadata = {
            'config': self.config,
            'weight_long': self.weight_long,
            'weight_short': self.weight_short,
            'last_retrain': self.last_retrain,
            'feature_names': self.feature_names,
            'is_fitted': self.is_fitted
        }
        joblib.dump(metadata, path / 'metadata.joblib')

        logger.info(f"Ensemble saved to {path}")

    @classmethod
    def load(cls, path: str) -> 'DualHorizonEnsemble':
        """Load ensemble from disk."""
        path = Path(path)

        # Load metadata
        metadata = joblib.load(path / 'metadata.joblib')

        # Create instance
        ensemble = cls(metadata['config'])
        ensemble.weight_long = metadata['weight_long']
        ensemble.weight_short = metadata['weight_short']
        ensemble.last_retrain = metadata['last_retrain']
        ensemble.feature_names = metadata['feature_names']
        ensemble.is_fitted = metadata['is_fitted']

        # Load models
        ensemble.long_term_high = joblib.load(path / 'long_term_high.joblib')
        ensemble.long_term_low = joblib.load(path / 'long_term_low.joblib')
        ensemble.short_term_high = joblib.load(path / 'short_term_high.joblib')
        ensemble.short_term_low = joblib.load(path / 'short_term_low.joblib')

        logger.info(f"Ensemble loaded from {path}")
        return ensemble


if __name__ == "__main__":
    # Test the module
    print("Testing DualHorizonEnsemble...")

    # Create sample data spanning 6 years
    np.random.seed(42)
    n = 100000  # ~6 years of 5-min data

    # Create datetime index
    start_date = datetime(2019, 1, 1)
    dates = pd.date_range(start_date, periods=n, freq='5min')

    # Simulate features (10 features for test)
    X = np.random.randn(n, 10)

    # Simulate targets with some pattern
    base_high = 1.0 + 0.3 * np.sin(np.arange(n) / 1000)  # Long-term pattern
    base_low = 0.8 + 0.2 * np.cos(np.arange(n) / 500)    # Different pattern

    y_high = base_high + 0.2 * np.random.randn(n)
    y_low = base_low + 0.2 * np.random.randn(n)

    # Ensure non-negative
    y_high = np.maximum(y_high, 0.1)
    y_low = np.maximum(y_low, 0.1)

    # Test ensemble
    config = DualHorizonConfig(
        long_term_years=5.0,
        short_term_months=3.0,
        initial_long_weight=0.6,
        initial_short_weight=0.4,
        use_dynamic_weights=True
    )

    ensemble = DualHorizonEnsemble(config)

    # Fit on training data
    train_idx = int(n * 0.8)
    ensemble.fit(
        X[:train_idx],
        y_high[:train_idx],
        y_low[:train_idx],
        dates[:train_idx],
        feature_names=[f'feature_{i}' for i in range(10)]
    )

    # Predict on test data
    pred_high, pred_low = ensemble.predict(X[train_idx:])

    print(f"\nResults:")
    print(f"  Training samples: {train_idx}")
    print(f"  Test samples: {n - train_idx}")
    print(f"  Predicted high range: [{pred_high.min():.2f}, {pred_high.max():.2f}]")
    print(f"  Predicted low range: [{pred_low.min():.2f}, {pred_low.max():.2f}]")
    print(f"  Model weights: Long={ensemble.weight_long:.2f}, Short={ensemble.weight_short:.2f}")

    # Test performance update and weight adjustment
    ensemble.update_performance(
        X[train_idx:train_idx+100],
        y_high[train_idx:train_idx+100],
        y_low[train_idx:train_idx+100]
    )
    ensemble.adjust_weights()

    print(f"\nAfter weight adjustment:")
    print(f"  Long weight: {ensemble.weight_long:.3f}")
    print(f"  Short weight: {ensemble.weight_short:.3f}")

    # Test feature importance
    importance = ensemble.get_feature_importance('combined')
    print(f"\nTop 5 features:")
    print(importance.head())
