"""
Модуль календарных признаков для временных рядов.

Функции:
    add_calendar_features — добавление всех календарных признаков
    add_datetime_features — признаки из даты (день недели, месяц, квартал)
    add_cyclical_features — циклическое кодирование времени (sin/cos)
    add_session_features — признаки торговой сессии (для H1)
"""

from typing import Optional

import numpy as np
import pandas as pd


def add_datetime_features(
    df: pd.DataFrame,
) -> pd.DataFrame:
    """
    Добавить признаки на основе даты.

    Рассчитываются:
        - day_of_week: день недели (0=пн, 6=вс)
        - month: месяц (1-12)
        - quarter: квартал (1-4)
        - is_month_end: конец месяца (0/1)
        - is_quarter_end: конец квартала (0/1)
        - day_of_month: день месяца (1-31)

    Args:
        df: DataFrame с колонками 'Date' (формат 'YYYY.MM.DD')
            или с DatetimeIndex.

    Returns:
        DataFrame с добавленными дата-признаками.
    """
    df = df.copy()
    dt = _get_datetime_series(df)
    if dt is None:
        return df

    df['day_of_week'] = dt.dt.dayofweek.astype(float)
    df['month'] = dt.dt.month.astype(float)
    df['quarter'] = dt.dt.quarter.astype(float)
    df['day_of_month'] = dt.dt.day.astype(float)
    df['is_month_end'] = dt.dt.is_month_end.astype(float)
    df['is_quarter_end'] = dt.dt.is_quarter_end.astype(float)
    return df


def add_cyclical_features(
    df: pd.DataFrame,
) -> pd.DataFrame:
    """
    Добавить циклическое кодирование времени (sin/cos).

    Циклические признаки позволяют модели понимать,
    что "понедельник" и "воскресенье" близки друг к другу.

    Рассчитываются:
        - day_sin, day_cos: циклическое кодирование дня недели
        - month_sin, month_cos: циклическое кодирование месяца

    Args:
        df: DataFrame с колонками 'day_of_week' и 'month'.

    Returns:
        DataFrame с sin/cos признаками.
    """
    df = df.copy()

    if 'day_of_week' in df.columns:
        df['day_sin'] = np.sin(2 * np.pi * df['day_of_week'] / 7)
        df['day_cos'] = np.cos(2 * np.pi * df['day_of_week'] / 7)

    if 'month' in df.columns:
        df['month_sin'] = np.sin(2 * np.pi * (df['month'] - 1) / 12)
        df['month_cos'] = np.cos(2 * np.pi * (df['month'] - 1) / 12)

    return df


def add_session_features(
    df: pd.DataFrame,
) -> pd.DataFrame:
    """
    Добавить признаки торговой сессии.

    Для D1 и W1 таймфреймов — полный день (1.0).
    Для H1 определяет: основная сессия, вечерняя сессия.

    Args:
        df: DataFrame с колонкой 'Time' (формат 'HH:MM').

    Returns:
        DataFrame с признаками сессии.
    """
    df = df.copy()

    if 'Time' not in df.columns:
        # Для дневных/недельных свечей — весь день
        df['is_main_session'] = 1.0
        df['is_evening_session'] = 0.0
        return df

    try:
        hour = pd.to_datetime(df['Time'], format='%H:%M').dt.hour
        df['is_main_session'] = ((hour >= 10) & (hour < 19)).astype(float)
        df['is_evening_session'] = ((hour >= 19) | (hour < 7)).astype(float)
    except (ValueError, KeyError):
        df['is_main_session'] = 1.0
        df['is_evening_session'] = 0.0

    return df


def _get_datetime_series(
    df: pd.DataFrame,
) -> Optional[pd.Series]:
    """
    Внутренняя функция: получить Series datetime из DataFrame.

    Args:
        df: DataFrame с 'Date' колонкой или DatetimeIndex.

    Returns:
        Series datetime или None, если не удалось распарсить.
    """
    if 'Date' in df.columns:
        try:
            return pd.to_datetime(df['Date'], format='%Y.%m.%d')
        except (ValueError, TypeError):
            try:
                return pd.to_datetime(df['Date'])
            except (ValueError, TypeError):
                return None

    if isinstance(df.index, pd.DatetimeIndex):
        return df.index.to_series()

    return None


def add_calendar_features(
    df: pd.DataFrame,
    include_datetime: bool = True,
    include_cyclical: bool = True,
    include_sessions: bool = True,
) -> pd.DataFrame:
    """
    Добавить все календарные признаки.

    Args:
        df: DataFrame с колонками 'Date' и опционально 'Time'.
        include_datetime: добавлять признаки из даты.
        include_cyclical: добавлять sin/cos кодирование.
        include_sessions: добавлять признаки сессии.

    Returns:
        DataFrame со всеми календарными признаками.
    """
    if include_datetime:
        df = add_datetime_features(df)
    if include_cyclical:
        df = add_cyclical_features(df)
    if include_sessions:
        df = add_session_features(df)
    return df