"""
Модуль кодирования целевых переменных для обучения нейронных сетей.

Поддерживаемые методы:
    - direction: 3 класса (0=HOLD, 1=BUY, 2=SELL) с порогом
    - binary: 2 класса (0=DOWN, 1=UP)
    - multihorizon: доходности на нескольких горизонтах
    - regression: непрерывное значение доходности

Функции:
    encode_target — главная функция кодирования цели
"""

from typing import Optional, Union

import numpy as np
import pandas as pd


def encode_target_direction(
    df: pd.DataFrame,
    forecast_horizon: int = 1,
    threshold_pct: float = 0.5,
) -> pd.DataFrame:
    """
    Закодировать целевую переменную как 3 класса.

    Классы:
        0 — HOLD (изменение цены в пределах порога)
        1 — BUY (рост больше порога)
        2 — SELL (падение больше порога)

    Args:
        df: DataFrame с колонкой 'Close'.
        forecast_horizon: горизонт прогноза в свечах вперёд.
        threshold_pct: порог доходности в процентах для BUY/SELL.

    Returns:
        DataFrame с колонкой 'target' (int).
    """
    df = df.copy()

    # Будущая цена через forecast_horizon свечей
    future_close = df['Close'].shift(-forecast_horizon)
    future_return = (future_close - df['Close']) / df['Close'] * 100

    # Кодирование
    df['target'] = 0  # HOLD по умолчанию
    df.loc[future_return > threshold_pct, 'target'] = 1   # BUY
    df.loc[future_return < -threshold_pct, 'target'] = 2  # SELL

    # Последние forecast_horizon строк не имеют таргета
    df.loc[df.index[-forecast_horizon:], 'target'] = np.nan

    # Гарантируем int тип
    df['target'] = df['target'].astype('Int64')
    return df


def encode_target_binary(
    df: pd.DataFrame,
    forecast_horizon: int = 1,
) -> pd.DataFrame:
    """
    Закодировать целевую переменную как 2 класса.

    Классы:
        0 — DOWN (цена упадёт или не изменится)
        1 — UP (цена вырастет)

    Args:
        df: DataFrame с колонкой 'Close'.
        forecast_horizon: горизонт прогноза в свечах вперёд.

    Returns:
        DataFrame с колонкой 'target' (int).
    """
    df = df.copy()

    future_close = df['Close'].shift(-forecast_horizon)
    future_return = (future_close - df['Close']) / df['Close']

    df['target'] = (future_return > 0).astype(int)
    df.loc[df.index[-forecast_horizon:], 'target'] = np.nan
    df['target'] = df['target'].astype('Int64')
    return df


def encode_target_multihorizon(
    df: pd.DataFrame,
    horizons: Optional[list[int]] = None,
    threshold_pct: float = 0.5,
) -> pd.DataFrame:
    """
    Закодировать целевые переменные для нескольких горизонтов.

    Добавляет колонки target_{h}d для каждого горизонта.

    Args:
        df: DataFrame с колонкой 'Close'.
        horizons: список горизонтов (по умолчанию [1, 3, 5, 10, 21]).
        threshold_pct: порог для классификации.

    Returns:
        DataFrame с колонками target_{h}d.
    """
    if horizons is None:
        horizons = [1, 3, 5, 10, 21]

    df = df.copy()

    for h in horizons:
        future_close = df['Close'].shift(-h)
        future_return = (future_close - df['Close']) / df['Close'] * 100

        target_col = f'target_{h}d'
        df[target_col] = 0  # HOLD
        df.loc[future_return > threshold_pct, target_col] = 1   # BUY
        df.loc[future_return < -threshold_pct, target_col] = 2  # SELL
        df.loc[df.index[-h:], target_col] = np.nan
        df[target_col] = df[target_col].astype('Int64')

    return df


def encode_target_regression(
    df: pd.DataFrame,
    forecast_horizon: int = 1,
) -> pd.DataFrame:
    """
    Закодировать целевую переменную как непрерывное значение (регрессия).

    Таргет: будущая доходность в процентах.

    Args:
        df: DataFrame с колонкой 'Close'.
        forecast_horizon: горизонт прогноза.

    Returns:
        DataFrame с колонкой 'target' (float).
    """
    df = df.copy()

    future_close = df['Close'].shift(-forecast_horizon)
    df['target'] = (future_close - df['Close']) / df['Close'] * 100

    df.loc[df.index[-forecast_horizon:], 'target'] = np.nan
    return df


def encode_target(
    df: pd.DataFrame,
    forecast_horizon: int = 1,
    method: str = 'direction',
    threshold_pct: float = 0.5,
) -> pd.DataFrame:
    """
    Главная функция кодирования целевой переменной.

    Args:
        df: DataFrame с колонкой 'Close'.
        forecast_horizon: горизонт прогноза в свечах.
        method: метод кодирования.
            'direction' — 3 класса (BUY/HOLD/SELL) с порогом.
            'binary' — 2 класса (UP/DOWN).
            'regression' — непрерывная доходность.
            'multihorizon' — несколько горизонтов.
        threshold_pct: порог доходности для direction (в процентах).

    Returns:
        DataFrame с закодированной целевой переменной.

    Raises:
        ValueError: если method не поддерживается.
    """
    if method == 'direction':
        return encode_target_direction(df, forecast_horizon, threshold_pct)
    elif method == 'binary':
        return encode_target_binary(df, forecast_horizon)
    elif method == 'regression':
        return encode_target_regression(df, forecast_horizon)
    elif method == 'multihorizon':
        return encode_target_multihorizon(df, threshold_pct=threshold_pct)
    else:
        raise ValueError(
            f"Неизвестный метод кодирования: '{method}'. "
            f"Доступные: 'direction', 'binary', 'regression', 'multihorizon'"
        )