import pandas as pd
import numpy as np


def add_returns(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df['return_1'] = df['Close'].pct_change(1)
    df['return_5'] = df['Close'].pct_change(5)
    df['return_20'] = df['Close'].pct_change(20)
    return df


def add_sma(df: pd.DataFrame, windows: list[int] | None = None) -> pd.DataFrame:
    if windows is None:
        windows = [5, 10, 20, 50]
    df = df.copy()
    for w in windows:
        df[f'sma_{w}'] = df['Close'].rolling(window=w).mean()
        df[f'close_to_sma_{w}'] = df['Close'] / df[f'sma_{w}'] - 1
    return df


def add_ema(df: pd.DataFrame, windows: list[int] | None = None) -> pd.DataFrame:
    if windows is None:
        windows = [12, 26]
    df = df.copy()
    for w in windows:
        df[f'ema_{w}'] = df['Close'].ewm(span=w, adjust=False).mean()
    if 12 in windows and 26 in windows:
        df['macd'] = df['ema_12'] - df['ema_26']
        df['macd_signal'] = df['macd'].ewm(span=9, adjust=False).mean()
        df['macd_hist'] = df['macd'] - df['macd_signal']
    return df


def add_rsi(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
    df = df.copy()
    delta = df['Close'].diff()
    gain = delta.where(delta > 0, 0.0)
    loss = (-delta).where(delta < 0, 0.0)
    avg_gain = gain.rolling(window=period, min_periods=1).mean()
    avg_loss = loss.rolling(window=period, min_periods=1).mean()
    rs = avg_gain / avg_loss.replace(0, np.nan)
    df['rsi'] = 100 - (100 / (1 + rs))
    df['rsi'] = df['rsi'].fillna(50)
    return df


def add_bollinger(df: pd.DataFrame, period: int = 20, num_std: float = 2.0) -> pd.DataFrame:
    df = df.copy()
    df['bb_mid'] = df['Close'].rolling(window=period).mean()
    bb_std = df['Close'].rolling(window=period).std()
    df['bb_upper'] = df['bb_mid'] + num_std * bb_std
    df['bb_lower'] = df['bb_mid'] - num_std * bb_std
    df['bb_width'] = (df['bb_upper'] - df['bb_lower']) / df['bb_mid']
    df['bb_position'] = (df['Close'] - df['bb_lower']) / (df['bb_upper'] - df['bb_lower'] + 1e-10)
    return df


def add_atr(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
    df = df.copy()
    high_low = df['High'] - df['Low']
    high_close = (df['High'] - df['Close'].shift()).abs()
    low_close = (df['Low'] - df['Close'].shift()).abs()
    tr = pd.concat([high_low, high_close, low_close], axis=1).max(axis=1)
    df['atr'] = tr.rolling(window=period, min_periods=1).mean()
    df['atr_pct'] = df['atr'] / df['Close']
    return df


def add_volume_features(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df['volume_sma_5'] = df['Volume'].rolling(window=5).mean()
    df['volume_ratio'] = df['Volume'] / df['volume_sma_5'].replace(0, np.nan)
    df['volume_ratio'] = df['volume_ratio'].fillna(1)
    return df


def add_interaction_features(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    req = ['bb_width', 'atr_pct', 'rsi', 'bb_position', 'close_to_sma_50', 'return_20', 'volume_ratio']
    if all(c in df.columns for c in req):
        df['bb_width_x_atr_pct'] = df['bb_width'] * df['atr_pct']
        df['rsi_x_bb_pos'] = (df['rsi'] - 50) * (df['bb_position'] - 0.5)
        df['cts50_x_ret20'] = df['close_to_sma_50'] * df['return_20']
        df['bb_width_x_rsi'] = df['bb_width'] * (df['rsi'] - 50)
        df['atr_pct_x_vol_ratio'] = df['atr_pct'] * df['volume_ratio']
    return df


INTERACTION_COLS = [
    'bb_width_x_atr_pct',
    'rsi_x_bb_pos',
    'cts50_x_ret20',
    'bb_width_x_rsi',
    'atr_pct_x_vol_ratio',
]


def add_target(df: pd.DataFrame, horizon: int = 1) -> pd.DataFrame:
    df = df.copy()
    df['target'] = df['Close'].shift(-horizon) / df['Close'] - 1
    df['target_direction'] = (df['target'] > 0).astype(int)
    return df


FEATURE_COLS = [
    'return_1', 'return_5', 'return_20',
    'close_to_sma_5', 'close_to_sma_10', 'close_to_sma_20', 'close_to_sma_50',
    'rsi',
    'bb_width', 'bb_position',
    'macd', 'macd_signal', 'macd_hist',
    'atr_pct',
    'volume_ratio',
    'kst_norm', 'efficiency_ratio', 'choppiness',
]


def add_temporal_features(df: pd.DataFrame) -> pd.DataFrame:
    """Добавляет временные признаки: час, день недели, месяц, циклические кодировки."""
    df = df.copy()

    # Парсим дату и время из колонок Date (YYYY.MM.DD) и Time (HH:MM)
    # Если есть timestamp, используем его
    if 'timestamp' in df.columns:
        dt = pd.to_datetime(df['timestamp'], unit='s', utc=True)
        # Смещение на MSK (UTC+3) для MOEX
        dt_msk = dt.dt.tz_convert('Europe/Moscow') if hasattr(dt.dt, 'tz_convert') else dt
    else:
        dt = pd.to_datetime(df['Date'] + ' ' + df['Time'], format='%Y.%m.%d %H:%M')
        dt_msk = dt

    df['hour'] = dt_msk.dt.hour.astype(float)
    df['day_of_week'] = dt_msk.dt.dayofweek.astype(float)  # 0=Monday
    df['month'] = dt_msk.dt.month.astype(float)
    df['day_of_month'] = dt_msk.dt.day.astype(float)
    df['quarter'] = dt_msk.dt.quarter.astype(float)
    df['is_weekend'] = (dt_msk.dt.dayofweek >= 5).astype(float)

    # Циклическое кодирование часа (для учёта цикличности времени суток)
    hour_rad = 2 * np.pi * df['hour'] / 24
    df['hour_sin'] = np.sin(hour_rad)
    df['hour_cos'] = np.cos(hour_rad)

    # Циклическое кодирование дня недели
    dow_rad = 2 * np.pi * df['day_of_week'] / 7
    df['dow_sin'] = np.sin(dow_rad)
    df['dow_cos'] = np.cos(dow_rad)

    # Циклическое кодирование месяца
    month_rad = 2 * np.pi * df['month'] / 12
    df['month_sin'] = np.sin(month_rad)
    df['month_cos'] = np.cos(month_rad)

    # Флаги торговых сессий MOEX (MSK)
    # Основная сессия: 10:00-18:45
    # Вечерняя сессия: 19:00-23:50
    df['is_main_session'] = ((df['hour'] >= 10) & (df['hour'] < 19)).astype(float)
    df['is_evening_session'] = ((df['hour'] >= 19) & (df['hour'] < 24)).astype(float)
    df['is_market_open'] = ((df['is_main_session'] + df['is_evening_session']) >= 0.5).astype(float)

    # Время от открытия сессии (в часах)
    # Основная сессия: 10:00 → session_hour=0, 18:00 → session_hour=8
    # Вечерняя сессия: 19:00 → session_hour=0, 23:00 → session_hour=4
    df['session_hour'] = np.where(
        df['is_main_session'] == 1, df['hour'] - 10,
        np.where(df['is_evening_session'] == 1, df['hour'] - 19, -1)
    )

    return df


TEMPORAL_COLS = [
    'hour', 'day_of_week', 'month', 'day_of_month', 'quarter', 'is_weekend',
    'hour_sin', 'hour_cos', 'dow_sin', 'dow_cos', 'month_sin', 'month_cos',
    'is_main_session', 'is_evening_session', 'is_market_open', 'session_hour',
]


def compute_kst_array(df: pd.DataFrame) -> np.ndarray:
    """KST (Know Sure Thing) oscillator — взвешенная сумма сглаженных ROC.

    ROC periods: 10, 15, 20, 30
    SMA (smoothing) periods: 10, 10, 10, 15
    Weights: 1, 2, 3, 4
    KST = RCMA1 + 2*RCMA2 + 3*RCMA3 + 4*RCMA4

    Returns np.ndarray normalized to ~[-1, +1].
    """
    close = df['Close'].values
    n = len(close)

    roc10 = np.zeros(n)
    roc15 = np.zeros(n)
    roc20 = np.zeros(n)
    roc30 = np.zeros(n)
    for i in range(30, n):
        if close[i - 10] > 1e-10:
            roc10[i] = (close[i] / close[i - 10] - 1) * 100
        if close[i - 15] > 1e-10:
            roc15[i] = (close[i] / close[i - 15] - 1) * 100
        if close[i - 20] > 1e-10:
            roc20[i] = (close[i] / close[i - 20] - 1) * 100
        if close[i - 30] > 1e-10:
            roc30[i] = (close[i] / close[i - 30] - 1) * 100

    s_roc10 = pd.Series(roc10).rolling(10, min_periods=1).mean().values
    s_roc15 = pd.Series(roc15).rolling(10, min_periods=1).mean().values
    s_roc20 = pd.Series(roc20).rolling(10, min_periods=1).mean().values
    s_roc30 = pd.Series(roc30).rolling(15, min_periods=1).mean().values

    kst = s_roc10 + 2 * s_roc15 + 3 * s_roc20 + 4 * s_roc30
    return kst / 100.0


def add_kst(df: pd.DataFrame) -> pd.DataFrame:
    """KST (Know Sure Thing) oscillator — добавляет kst_norm в DataFrame."""
    df = df.copy()
    df['kst_norm'] = compute_kst_array(df)
    return df


def compute_efficiency_ratio_array(df: pd.DataFrame, period: int = 10) -> np.ndarray:
    """Kaufman Efficiency Ratio. Возвращает np.ndarray [0, 1]."""
    close = df['Close'].values
    n = len(close)
    er = np.zeros(n)
    for i in range(period, n):
        total_move = abs(close[i] - close[i - period])
        changes = np.abs(np.diff(close[i - period:i + 1]))
        noise = np.sum(changes)
        er[i] = total_move / noise if noise > 1e-10 else 0.0
    return np.clip(er, 0, 1)


def add_efficiency_ratio(df: pd.DataFrame, period: int = 10) -> pd.DataFrame:
    """Kaufman Efficiency Ratio — добавляет efficiency_ratio в DataFrame."""
    df = df.copy()
    df['efficiency_ratio'] = compute_efficiency_ratio_array(df, period)
    return df


def compute_choppiness_array(df: pd.DataFrame, period: int = 14) -> np.ndarray:
    """Choppiness Index. Возвращает np.ndarray [0, 1]."""
    high = df['High'].values
    low = df['Low'].values
    close = df['Close'].values
    n = len(high)
    chop = np.zeros(n)
    for i in range(period, n):
        tr_sum = 0.0
        for j in range(i - period + 1, i + 1):
            hl = high[j] - low[j]
            hc = abs(high[j] - close[j - 1]) if j > 0 else 0
            lc = abs(low[j] - close[j - 1]) if j > 0 else 0
            tr_sum += max(hl, hc, lc)
        hh = max(high[i - period + 1:i + 1])
        ll = min(low[i - period + 1:i + 1])
        if hh - ll > 1e-10 and tr_sum > 1e-10:
            chop[i] = 100 * np.log10(tr_sum / (hh - ll)) / np.log10(period)
        else:
            chop[i] = 50
    return np.clip(chop / 100, 0, 1)


def add_choppiness(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
    """Choppiness Index — добавляет choppiness в DataFrame."""
    df = df.copy()
    df['choppiness'] = compute_choppiness_array(df, period)
    return df


def engineer_features(df: pd.DataFrame, skip_temporal: bool = False) -> pd.DataFrame:
    """
    Добавляет все технические индикаторы.

    Parameters
    ----------
    skip_temporal : bool
        Если True, пропускает временные признаки (hour, day_of_week и т.д.).
        Рекомендуется для D1/W1, где час всегда 0, день недели не информативен.
    """
    df = df.copy()
    df = add_returns(df)
    df = add_sma(df)
    df = add_ema(df)
    df = add_rsi(df)
    df = add_bollinger(df)
    df = add_atr(df)
    df = add_volume_features(df)
    df = add_interaction_features(df)
    df = add_kst(df)
    df = add_efficiency_ratio(df)
    df = add_choppiness(df)
    # SMA 200 / EMA 200 для режимного фильтра MoE (не входят в FEATURE_COLS)
    df['sma_200'] = df['Close'].rolling(window=200).mean()
    df['ema_200'] = df['Close'].ewm(span=200, adjust=False).mean()
    if not skip_temporal:
        df = add_temporal_features(df)
    return df