"""Data preprocessing: normalization, feature engineering, and sequence creation."""

from typing import Tuple, Optional, Literal
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
import pandas_ta as ta


class CryptoPreprocessor:
    """Preprocess cryptocurrency data with normalization and feature engineering."""

    def __init__(
        self,
        lookback: int = 168,
        forecast_horizon: int = 24,
        price_change_threshold: float = 0.0  # 改为0，使用真实涨跌
    ):
        """Initialize preprocessor.

        Args:
            lookback: Number of historical timesteps to use
            forecast_horizon: Number of timesteps to predict ahead
            price_change_threshold: Threshold for binary classification (e.g., 0.015 = 1.5%)
        """
        self.lookback = lookback
        self.forecast_horizon = forecast_horizon
        self.price_change_threshold = price_change_threshold
        self.scaler = MinMaxScaler()
        self.feature_columns = []
        self.is_fitted = False

    def add_technical_indicators(self, df: pd.DataFrame) -> pd.DataFrame:
        """添加技术指标作为特征

        Args:
            df: 包含OHLCV数据的DataFrame

        Returns:
            添加了技术指标的DataFrame
        """
        df = df.copy()

        # 使用pandas_ta库
        # RSI (14 periods)
        df['rsi'] = ta.rsi(df['close'], length=14)

        # MACD
        macd = ta.macd(df['close'])
        if macd is not None and not macd.empty:
            # pandas-ta 返回的列名可能不同，动态获取
            macd_cols = macd.columns.tolist()
            if len(macd_cols) >= 3:
                df['macd'] = macd.iloc[:, 0]        # MACD线
                df['macd_hist'] = macd.iloc[:, 1]   # MACD柱
                df['macd_signal'] = macd.iloc[:, 2] # 信号线
            else:
                # 如果格式不对，使用None
                df['macd'] = None
                df['macd_signal'] = None
                df['macd_hist'] = None
        else:
            df['macd'] = None
            df['macd_signal'] = None
            df['macd_hist'] = None

        # Bollinger Bands
        bbands = ta.bbands(df['close'], length=20)
        if bbands is not None and not bbands.empty:
            # pandas-ta 返回的列名可能不同，动态获取
            bb_cols = bbands.columns.tolist()
            if len(bb_cols) >= 3:
                df['bb_lower'] = bbands.iloc[:, 0]   # 第一列是下轨
                df['bb_middle'] = bbands.iloc[:, 1]  # 第二列是中轨
                df['bb_upper'] = bbands.iloc[:, 2]   # 第三列是上轨
                df['bb_bandwidth'] = (df['bb_upper'] - df['bb_lower']) / (df['bb_middle'] + 1e-8)
            else:
                # 如果格式不对，手动计算
                sma = df['close'].rolling(window=20).mean()
                std = df['close'].rolling(window=20).std()
                df['bb_middle'] = sma
                df['bb_upper'] = sma + (std * 2)
                df['bb_lower'] = sma - (std * 2)
                df['bb_bandwidth'] = (df['bb_upper'] - df['bb_lower']) / (df['bb_middle'] + 1e-8)
        else:
            # pandas-ta不可用，手动计算布林带
            sma = df['close'].rolling(window=20).mean()
            std = df['close'].rolling(window=20).std()
            df['bb_middle'] = sma
            df['bb_upper'] = sma + (std * 2)
            df['bb_lower'] = sma - (std * 2)
            df['bb_bandwidth'] = (df['bb_upper'] - df['bb_lower']) / (df['bb_middle'] + 1e-8)

        # Volume indicators
        df['volume_sma'] = ta.sma(df['volume'], length=20)
        df['volume_ratio'] = df['volume'] / (df['volume_sma'] + 1e-8)

        # Price momentum
        df['returns'] = df['close'].pct_change()
        df['returns_5'] = df['close'].pct_change(5)
        df['returns_10'] = df['close'].pct_change(10)

        # Volatility
        df['volatility'] = df['returns'].rolling(window=20).std()

        return df

    def add_time_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """Add time-based features.

        Args:
            df: DataFrame with timestamp column

        Returns:
            DataFrame with added time features
        """
        df = df.copy()

        df['hour'] = df['timestamp'].dt.hour
        df['day_of_week'] = df['timestamp'].dt.dayofweek
        df['is_weekend'] = (df['day_of_week'] >= 5).astype(int)

        # Cyclical encoding for hour
        df['hour_sin'] = np.sin(2 * np.pi * df['hour'] / 24)
        df['hour_cos'] = np.cos(2 * np.pi * df['hour'] / 24)

        # Cyclical encoding for day of week
        df['dow_sin'] = np.sin(2 * np.pi * df['day_of_week'] / 7)
        df['dow_cos'] = np.cos(2 * np.pi * df['day_of_week'] / 7)

        return df

    def create_labels(
        self,
        df: pd.DataFrame,
        task: Literal['classification', 'regression']
    ) -> pd.DataFrame:
        """Create target labels for classification or regression.

        Args:
            df: DataFrame with price data
            task: Task type ('classification' or 'regression')

        Returns:
            DataFrame with added label column
        """
        df = df.copy()

        # Calculate future returns
        df['future_price'] = df['close'].shift(-self.forecast_horizon)
        df['future_return'] = (df['future_price'] - df['close']) / df['close']

        if task == 'classification':
            # Binary classification: 1 if price increases above threshold, 0 otherwise
            df['label'] = (df['future_return'] > self.price_change_threshold).astype(int)
        else:  # regression
            # Regression: predict percentage return
            df['label'] = df['future_return']

        # Drop rows without labels (last forecast_horizon rows)
        df = df[:-self.forecast_horizon].copy()

        return df

    def fit(self, df: pd.DataFrame) -> 'CryptoPreprocessor':
        """Fit the scaler on training data.

        Args:
            df: Training DataFrame with features

        Returns:
            Self
        """
        # Get feature columns (exclude timestamp and label)
        exclude_cols = ['timestamp', 'label', 'future_price', 'future_return']
        self.feature_columns = [col for col in df.columns if col not in exclude_cols]

        # Fit scaler
        self.scaler.fit(df[self.feature_columns].fillna(0))
        self.is_fitted = True

        return self

    def transform(self, df: pd.DataFrame) -> np.ndarray:
        """Transform data using fitted scaler.

        Args:
            df: DataFrame with features

        Returns:
            Scaled feature array
        """
        if not self.is_fitted:
            raise ValueError("Preprocessor must be fitted before transform")

        # Fill NaN values
        df_filled = df[self.feature_columns].fillna(0)

        # Scale features
        scaled = self.scaler.transform(df_filled)

        return scaled

    def fit_transform(self, df: pd.DataFrame) -> np.ndarray:
        """Fit and transform in one step.

        Args:
            df: DataFrame with features

        Returns:
            Scaled feature array
        """
        return self.fit(df).transform(df)

    def create_sequences(
        self,
        features: np.ndarray,
        labels: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Create sequences for time series prediction.

        Args:
            features: Feature array (n_samples, n_features)
            labels: Label array (n_samples,)

        Returns:
            Tuple of (X, y) where X is (n_sequences, lookback, n_features)
            and y is (n_sequences,)
        """
        X, y = [], []

        for i in range(len(features) - self.lookback):
            X.append(features[i:i + self.lookback])
            y.append(labels[i + self.lookback])

        return np.array(X), np.array(y)

    def preprocess_pipeline(
        self,
        df: pd.DataFrame,
        task: Literal['classification', 'regression'],
        fit: bool = True
    ) -> Tuple[np.ndarray, np.ndarray, pd.DataFrame]:
        """Full preprocessing pipeline.

        Args:
            df: Raw OHLCV DataFrame
            task: Task type
            fit: Whether to fit the scaler (True for training data)

        Returns:
            Tuple of (X, y, processed_df) where X is sequences, y is labels,
            and processed_df is the DataFrame after feature engineering
        """
        # Add features
        df = self.add_technical_indicators(df)
        df = self.add_time_features(df)
        df = self.create_labels(df, task)

        # Drop initial NaN rows from indicators
        df = df.dropna().reset_index(drop=True)

        # Scale features
        if fit:
            features = self.fit_transform(df)
        else:
            features = self.transform(df)

        # Get labels
        labels = df['label'].values

        # Create sequences
        X, y = self.create_sequences(features, labels)

        print(f"Created {len(X)} sequences with shape {X.shape}")

        return X, y, df


def split_data(
    X: np.ndarray,
    y: np.ndarray,
    train_split: float = 0.7,
    val_split: float = 0.15
) -> Tuple[Tuple[np.ndarray, np.ndarray], ...]:
    """Split data into train/val/test sets using time-based split.

    Args:
        X: Feature sequences
        y: Labels
        train_split: Fraction for training
        val_split: Fraction for validation

    Returns:
        Tuple of ((X_train, y_train), (X_val, y_val), (X_test, y_test))
    """
    n = len(X)
    train_end = int(n * train_split)
    val_end = int(n * (train_split + val_split))

    X_train, y_train = X[:train_end], y[:train_end]
    X_val, y_val = X[train_end:val_end], y[train_end:val_end]
    X_test, y_test = X[val_end:], y[val_end:]

    print(f"Train: {len(X_train)} samples")
    print(f"Val: {len(X_val)} samples")
    print(f"Test: {len(X_test)} samples")

    return (X_train, y_train), (X_val, y_val), (X_test, y_test)


if __name__ == "__main__":
    # Test the preprocessor
    from fetcher import BinanceDataFetcher

    # Fetch data
    fetcher = BinanceDataFetcher()
    df = fetcher.fetch_ohlcv('BTC', '1h')

    # Preprocess for classification
    preprocessor = CryptoPreprocessor(lookback=168, forecast_horizon=24)
    X, y, processed_df = preprocessor.preprocess_pipeline(df, task='classification')

    print(f"\nFeatures shape: {X.shape}")
    print(f"Labels shape: {y.shape}")
    print(f"Label distribution: {np.bincount(y.astype(int))}")

    # Split data
    (X_train, y_train), (X_val, y_val), (X_test, y_test) = split_data(X, y)