"""
Temporal Data Splitter for ML-First Strategy
============================================
Implements out-of-sample (OOS) validation by excluding specified time periods.

Key Principle: 2025 data is NEVER seen during training - reserved for OOS validation.

Author: ML-Specialist (NEXUS v4.0)
Created: 2026-01-04
"""

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


@dataclass
class TemporalSplit:
    """Data class representing a temporal train/test split"""
    name: str
    train_data: pd.DataFrame
    test_data: pd.DataFrame
    train_start: datetime
    train_end: datetime
    test_start: datetime
    test_end: datetime
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def train_size(self) -> int:
        return len(self.train_data)

    @property
    def test_size(self) -> int:
        return len(self.test_data)

    @property
    def train_date_range(self) -> str:
        return f"{self.train_start.strftime('%Y-%m-%d')} to {self.train_end.strftime('%Y-%m-%d')}"

    @property
    def test_date_range(self) -> str:
        return f"{self.test_start.strftime('%Y-%m-%d')} to {self.test_end.strftime('%Y-%m-%d')}"

    def __repr__(self) -> str:
        return (
            f"TemporalSplit('{self.name}')\n"
            f"  Train: {self.train_date_range} ({self.train_size:,} samples)\n"
            f"  Test:  {self.test_date_range} ({self.test_size:,} samples)"
        )


class TemporalDataSplitter:
    """
    Temporal data splitter for out-of-sample validation.

    Ensures that 2025 data (or any specified period) is NEVER seen during training.
    This is critical for realistic backtesting and avoiding overfitting.
    """

    def __init__(self, config_path: str = "config/validation_oos.yaml"):
        """
        Initialize the temporal data splitter.

        Args:
            config_path: Path to validation configuration YAML file
        """
        self.config = self._load_config(config_path)
        self.splits: List[TemporalSplit] = []

    def _load_config(self, config_path: str) -> Dict[str, Any]:
        """Load configuration from YAML file"""
        config_file = Path(config_path)
        if not config_file.exists():
            logger.warning(f"Config not found: {config_path}. Using defaults.")
            return self._default_config()

        with open(config_file, 'r') as f:
            config = yaml.safe_load(f)

        logger.info(f"Loaded OOS validation config from {config_path}")
        return config

    def _default_config(self) -> Dict[str, Any]:
        """Return default configuration"""
        return {
            'validation': {
                'train': {
                    'start_date': '2023-01-01T00:00:00',
                    'end_date': '2024-12-31T23:59:59'
                },
                'test_oos': {
                    'start_date': '2025-01-01T00:00:00',
                    'end_date': '2025-12-31T23:59:59'
                },
                'exclusion_method': 'temporal'
            }
        }

    def split_temporal(
        self,
        df: pd.DataFrame,
        train_start: Optional[str] = None,
        train_end: Optional[str] = None,
        test_start: Optional[str] = None,
        test_end: Optional[str] = None
    ) -> TemporalSplit:
        """
        Split data by temporal boundaries.

        Args:
            df: DataFrame with datetime index
            train_start: Start of training period (defaults to config)
            train_end: End of training period (defaults to config)
            test_start: Start of test period (defaults to config)
            test_end: End of test period (defaults to config)

        Returns:
            TemporalSplit object containing train and test data
        """
        # Use config values if not provided
        val_config = self.config['validation']

        train_start = train_start or val_config['train']['start_date']
        train_end = train_end or val_config['train']['end_date']
        test_start = test_start or val_config['test_oos']['start_date']
        test_end = test_end or val_config['test_oos']['end_date']

        # Convert to datetime
        train_start_dt = pd.to_datetime(train_start)
        train_end_dt = pd.to_datetime(train_end)
        test_start_dt = pd.to_datetime(test_start)
        test_end_dt = pd.to_datetime(test_end)

        # Ensure index is datetime
        if not isinstance(df.index, pd.DatetimeIndex):
            raise ValueError("DataFrame must have a DatetimeIndex")

        # Filter data
        train_mask = (df.index >= train_start_dt) & (df.index <= train_end_dt)
        test_mask = (df.index >= test_start_dt) & (df.index <= test_end_dt)

        train_data = df[train_mask].copy()
        test_data = df[test_mask].copy()

        # Create split object
        split = TemporalSplit(
            name="temporal_oos",
            train_data=train_data,
            test_data=test_data,
            train_start=train_start_dt,
            train_end=train_end_dt,
            test_start=test_start_dt,
            test_end=test_end_dt,
            metadata={
                'exclusion_method': 'temporal',
                'config_used': self.config['validation']
            }
        )

        self.splits.append(split)

        logger.info(f"Created temporal split:")
        logger.info(f"  Train: {split.train_date_range} ({split.train_size:,} samples)")
        logger.info(f"  Test OOS: {split.test_date_range} ({split.test_size:,} samples)")

        # Validate split
        self._validate_split(split)

        return split

    def split_with_validation(
        self,
        df: pd.DataFrame,
        val_ratio: float = 0.15
    ) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
        """
        Create train/validation/test split with validation carved from training period.

        Args:
            df: DataFrame with datetime index
            val_ratio: Ratio of training data to use for validation

        Returns:
            Tuple of (train_df, val_df, test_df)
        """
        # First get the main temporal split
        split = self.split_temporal(df)

        # Then split training data into train/validation
        train_size = len(split.train_data)
        val_size = int(train_size * val_ratio)

        # Use the most recent training data for validation (temporal ordering)
        train_df = split.train_data.iloc[:-val_size].copy()
        val_df = split.train_data.iloc[-val_size:].copy()
        test_df = split.test_data.copy()

        logger.info(f"Train/Val/Test split:")
        logger.info(f"  Train: {len(train_df):,} samples ({train_df.index.min()} to {train_df.index.max()})")
        logger.info(f"  Val:   {len(val_df):,} samples ({val_df.index.min()} to {val_df.index.max()})")
        logger.info(f"  Test:  {len(test_df):,} samples ({test_df.index.min()} to {test_df.index.max()})")

        return train_df, val_df, test_df

    def split_walk_forward_with_oos(
        self,
        df: pd.DataFrame,
        n_splits: int = 5,
        gap_periods: int = 0,
        expanding_window: bool = False
    ) -> List[Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]]:
        """
        Create walk-forward splits where test data is always OOS (2025).

        This is the recommended method for robust validation:
        - Walk-forward on training period (2023-2024)
        - Final test always on OOS period (2025)

        Args:
            df: DataFrame with datetime index
            n_splits: Number of walk-forward splits
            gap_periods: Gap between train and validation (avoid look-ahead)
            expanding_window: If True, training window expands

        Returns:
            List of (train, val, test_oos) tuples
        """
        # Get main temporal split first
        main_split = self.split_temporal(df)
        train_data = main_split.train_data
        test_oos = main_split.test_data

        # Calculate walk-forward splits on training data
        splits = []
        train_size = len(train_data)
        step_size = train_size // (n_splits + 1)
        val_size = int(step_size * 0.2)

        for i in range(n_splits):
            if expanding_window:
                split_train_start = 0
            else:
                split_train_start = i * step_size if i > 0 else 0

            split_train_end = (i + 1) * step_size
            split_val_start = split_train_end + gap_periods
            split_val_end = min(split_val_start + val_size, train_size)

            split_train = train_data.iloc[split_train_start:split_train_end].copy()
            split_val = train_data.iloc[split_val_start:split_val_end].copy()

            splits.append((split_train, split_val, test_oos.copy()))

            logger.info(
                f"Walk-forward split {i+1}/{n_splits}: "
                f"Train {len(split_train):,}, Val {len(split_val):,}, Test OOS {len(test_oos):,}"
            )

        return splits

    def _validate_split(self, split: TemporalSplit) -> bool:
        """
        Validate that the split is correct and there's no data leakage.

        Args:
            split: TemporalSplit to validate

        Returns:
            True if valid, raises exception otherwise
        """
        # Check no overlap
        if split.train_end >= split.test_start:
            raise ValueError(
                f"Data leakage detected! Train end ({split.train_end}) >= "
                f"Test start ({split.test_start})"
            )

        # Check test is truly OOS
        if not split.test_data.empty:
            test_min = split.test_data.index.min()
            train_max = split.train_data.index.max()

            if test_min <= train_max:
                raise ValueError(
                    f"OOS violation! Test data ({test_min}) overlaps with "
                    f"train data ({train_max})"
                )

        # Check sizes
        if split.train_size < 1000:
            logger.warning(f"Training set is small: {split.train_size} samples")

        if split.test_size < 100:
            logger.warning(f"Test set is small: {split.test_size} samples")

        # Calculate ratio
        ratio = split.test_size / (split.train_size + split.test_size)
        logger.info(f"Train/Test ratio: {1-ratio:.1%}/{ratio:.1%}")

        return True

    def exclude_year(
        self,
        df: pd.DataFrame,
        year: int = 2025
    ) -> pd.DataFrame:
        """
        Simple method to exclude a specific year from data.

        Args:
            df: DataFrame with datetime index
            year: Year to exclude

        Returns:
            DataFrame without the specified year
        """
        if not isinstance(df.index, pd.DatetimeIndex):
            raise ValueError("DataFrame must have a DatetimeIndex")

        original_size = len(df)
        df_filtered = df[df.index.year != year].copy()
        excluded = original_size - len(df_filtered)

        logger.info(f"Excluded {excluded:,} samples from year {year}")
        logger.info(f"Remaining: {len(df_filtered):,} samples")

        return df_filtered

    def get_oos_data(
        self,
        df: pd.DataFrame
    ) -> pd.DataFrame:
        """
        Get only the OOS (out-of-sample) data for final evaluation.

        Args:
            df: DataFrame with datetime index

        Returns:
            DataFrame containing only OOS period
        """
        val_config = self.config['validation']
        test_start = pd.to_datetime(val_config['test_oos']['start_date'])
        test_end = pd.to_datetime(val_config['test_oos']['end_date'])

        oos_mask = (df.index >= test_start) & (df.index <= test_end)
        oos_data = df[oos_mask].copy()

        logger.info(f"OOS data: {len(oos_data):,} samples ({test_start} to {test_end})")

        return oos_data

    def get_training_data(
        self,
        df: pd.DataFrame
    ) -> pd.DataFrame:
        """
        Get only the training data (excluding OOS period).

        Args:
            df: DataFrame with datetime index

        Returns:
            DataFrame containing only training period
        """
        val_config = self.config['validation']
        train_start = pd.to_datetime(val_config['train']['start_date'])
        train_end = pd.to_datetime(val_config['train']['end_date'])

        train_mask = (df.index >= train_start) & (df.index <= train_end)
        train_data = df[train_mask].copy()

        logger.info(f"Training data: {len(train_data):,} samples ({train_start} to {train_end})")

        return train_data

    def print_data_summary(self, df: pd.DataFrame):
        """Print summary of data distribution by year"""
        if not isinstance(df.index, pd.DatetimeIndex):
            logger.warning("Cannot summarize: index is not DatetimeIndex")
            return

        logger.info("=" * 60)
        logger.info("DATA DISTRIBUTION BY YEAR")
        logger.info("=" * 60)

        year_counts = df.groupby(df.index.year).size()
        total = len(df)

        for year, count in year_counts.items():
            pct = count / total * 100
            bar = "#" * int(pct / 2)
            logger.info(f"  {year}: {count:>10,} ({pct:>5.1f}%) {bar}")

        logger.info("-" * 60)
        logger.info(f"  Total: {total:,} samples")
        logger.info(f"  Date range: {df.index.min()} to {df.index.max()}")
        logger.info("=" * 60)


def create_ml_first_splits(
    df: pd.DataFrame,
    config_path: str = "config/validation_oos.yaml"
) -> Dict[str, pd.DataFrame]:
    """
    Convenience function to create ML-First train/val/test splits.

    This is the recommended entry point for preparing data for ML training.

    Args:
        df: DataFrame with datetime index containing all features
        config_path: Path to validation config

    Returns:
        Dictionary with 'train', 'val', 'test_oos' DataFrames
    """
    splitter = TemporalDataSplitter(config_path)

    # Show data summary
    splitter.print_data_summary(df)

    # Create splits
    train_df, val_df, test_df = splitter.split_with_validation(df)

    return {
        'train': train_df,
        'val': val_df,
        'test_oos': test_df
    }


if __name__ == "__main__":
    # Test the splitter
    from datetime import timedelta

    # Create sample data spanning multiple years
    dates = pd.date_range(start='2023-01-01', end='2025-12-31', freq='5min')
    np.random.seed(42)

    df = pd.DataFrame({
        'open': 2000 + np.random.randn(len(dates)).cumsum() * 0.1,
        'high': 2000 + np.random.randn(len(dates)).cumsum() * 0.1 + 1,
        'low': 2000 + np.random.randn(len(dates)).cumsum() * 0.1 - 1,
        'close': 2000 + np.random.randn(len(dates)).cumsum() * 0.1,
        'volume': np.random.randint(100, 10000, len(dates)),
        'rsi': np.random.uniform(20, 80, len(dates)),
        'macd': np.random.randn(len(dates))
    }, index=dates)

    print(f"Total data: {len(df):,} samples")
    print(f"Date range: {df.index.min()} to {df.index.max()}")

    # Initialize splitter
    splitter = TemporalDataSplitter()

    # Show data summary
    splitter.print_data_summary(df)

    # Test temporal split
    print("\n" + "=" * 60)
    print("TESTING TEMPORAL SPLIT")
    print("=" * 60)

    split = splitter.split_temporal(df)
    print(split)

    # Test train/val/test split
    print("\n" + "=" * 60)
    print("TESTING TRAIN/VAL/TEST SPLIT")
    print("=" * 60)

    train_df, val_df, test_df = splitter.split_with_validation(df)
    print(f"Train: {len(train_df):,}")
    print(f"Val: {len(val_df):,}")
    print(f"Test OOS: {len(test_df):,}")

    # Test convenience function
    print("\n" + "=" * 60)
    print("TESTING CONVENIENCE FUNCTION")
    print("=" * 60)

    splits = create_ml_first_splits(df)
    for name, data in splits.items():
        print(f"{name}: {len(data):,} samples")
