import numpy as np
import pandas as pd
from numba import njit
from numpy.fft import fft, ifft
import random

# CCI labeling
@njit
def calculate_cci(close_data, period=20):
    """Расчет индикатора CCI на основе только close"""
    cci = np.zeros(len(close_data))
    
    for i in range(period - 1, len(close_data)):
        # Простое скользящее среднее
        sma = 0.0
        for j in range(i - period + 1, i + 1):
            sma += close_data[j]
        sma = sma / period
        
        # Среднее абсолютное отклонение
        mad_sum = 0.0
        for j in range(i - period + 1, i + 1):
            mad_sum += abs(close_data[j] - sma)
        
        mad = mad_sum / period
        
        # CCI формула
        if mad == 0:
            cci[i] = 0.0
        else:
            cci[i] = (close_data[i] - sma) / (0.015 * mad)
    
    return cci

@njit
def calculate_labels_cci(close_data, cci_data, 
                         oversold_level=-100.0, overbought_level=100.0):
    """
    Разметка на основе уровней CCI.
    """
    labels = []
    
    for i in range(len(close_data)):
        curr_cci = cci_data[i]
        
        # Сигнал на покупку: CCI в зоне перепроданности
        if curr_cci < oversold_level:
            labels.append(0.0)  # Buy
        
        # Сигнал на продажу: CCI в зоне перекупленности
        elif curr_cci > overbought_level:
            labels.append(1.0)  # Sell
        
        else:
            labels.append(2.0)  # Hold - CCI в нейтральной зоне
    
    return labels

def get_labels_cci(dataset, cci_period=20,
                   oversold_level=-100.0, overbought_level=100.0) -> pd.DataFrame:
    dataset = dataset.copy()
    
    close_data = dataset['close'].values
    cci_data = calculate_cci(close_data, cci_period)
    dataset['cci'] = cci_data

    labels = calculate_labels_cci(close_data, cci_data,
                                  oversold_level, overbought_level)

    # Обрезаем датасет
    dataset = dataset.iloc[:len(labels)].copy()
    dataset['labels'] = labels
    dataset = dataset.drop(columns=['cci'])
    
    return dataset.dropna()

# Stochastic labeling
@njit
def calculate_stochastic(close_data, period=14, smooth_k=3):
    """Расчет Stochastic Oscillator (%K и %D)"""
    stoch_k = np.zeros(len(close_data))
    stoch_d = np.zeros(len(close_data))
    
    # Расчет %K
    for i in range(period - 1, len(close_data)):
        # Находим максимум и минимум за period
        highest = close_data[i - period + 1]
        lowest = close_data[i - period + 1]
        
        for j in range(i - period + 1, i + 1):
            if close_data[j] > highest:
                highest = close_data[j]
            if close_data[j] < lowest:
                lowest = close_data[j]
        
        # Формула %K
        if highest - lowest == 0:
            stoch_k[i] = 50.0
        else:
            stoch_k[i] = 100.0 * (close_data[i] - lowest) / (highest - lowest)
    
    # Сглаживание %K для получения %D (SMA)
    for i in range(period + smooth_k - 2, len(close_data)):
        k_sum = 0.0
        for j in range(i - smooth_k + 1, i + 1):
            k_sum += stoch_k[j]
        stoch_d[i] = k_sum / smooth_k
    
    return stoch_k, stoch_d

@njit
def calculate_labels_stochastic(close_data, stoch_k, stoch_d,
                                oversold_level=20.0, overbought_level=80.0):
    """
    Разметка на основе Stochastic Oscillator.
    """
    labels = []
    
    for i in range(len(close_data)):
        curr_k = stoch_k[i]
        curr_d = stoch_d[i]
        
        # Сигнал на покупку: %K в зоне перепроданности и пересекает %D снизу вверх
        if curr_k < oversold_level:
            labels.append(0.0)  # Buy
        
        # Сигнал на продажу: %K в зоне перекупленности и пересекает %D сверху вниз
        elif curr_k > overbought_level:
            labels.append(1.0)  # Sell
        
        else:
            labels.append(2.0)  # Hold
    
    return labels

def get_labels_stochastic(dataset, stoch_period=14, smooth_k=3,
                          oversold_level=20.0, overbought_level=80.0) -> pd.DataFrame:
    dataset = dataset.copy()
    
    close_data = dataset['close'].values
    stoch_k, stoch_d = calculate_stochastic(close_data, stoch_period, smooth_k)
    
    dataset['stoch_k'] = stoch_k
    dataset['stoch_d'] = stoch_d

    labels = calculate_labels_stochastic(close_data, stoch_k, stoch_d,
                                        oversold_level, overbought_level)

    # Обрезаем датасет
    dataset = dataset.iloc[:len(labels)].copy()
    dataset['labels'] = labels
    dataset = dataset.drop(columns=['stoch_k', 'stoch_d'])
    
    return dataset.dropna()

# Bollindger labeling
@njit
def calculate_bollinger_bands(close_data, period=20, num_std=2.0):
    """Расчет полос Боллинджера"""
    bb_upper = np.zeros(len(close_data))
    bb_middle = np.zeros(len(close_data))
    bb_lower = np.zeros(len(close_data))
    
    for i in range(period - 1, len(close_data)):
        # Скользящее среднее (SMA)
        sma = 0.0
        for j in range(i - period + 1, i + 1):
            sma += close_data[j]
        sma = sma / period
        
        # Стандартное отклонение
        variance = 0.0
        for j in range(i - period + 1, i + 1):
            variance += (close_data[j] - sma) ** 2
        std = np.sqrt(variance / period)
        
        bb_middle[i] = sma
        bb_upper[i] = sma + (num_std * std)
        bb_lower[i] = sma - (num_std * std)
    
    return bb_upper, bb_middle, bb_lower

@njit
def calculate_labels_bb(close_data, bb_upper, bb_lower):
    """
    Разметка на основе полос Боллинджера.
    Если цена выше верхней ленты -> продажа (1.0)
    Если цена ниже нижней ленты -> покупка (0.0)
    Иначе -> hold (2.0)
    """
    labels = []
    
    for i in range(len(close_data)):
        price = close_data[i]
        upper = bb_upper[i]
        lower = bb_lower[i]
        
        # Цена выше верхней ленты - продажа
        if price > upper:
            labels.append(1.0)  # Sell
        
        # Цена ниже нижней ленты - покупка
        elif price < lower:
            labels.append(0.0)  # Buy
        
        # Цена внутри лент
        else:
            labels.append(2.0)  # Hold
    
    return labels

def get_labels_bb(dataset, bb_period=20, num_std=2.0) -> pd.DataFrame:
    dataset = dataset.copy()
    
    close_data = dataset['close'].values
    bb_upper, bb_middle, bb_lower = calculate_bollinger_bands(close_data, bb_period, num_std)
    
    dataset['bb_upper'] = bb_upper
    dataset['bb_middle'] = bb_middle
    dataset['bb_lower'] = bb_lower

    labels = calculate_labels_bb(close_data, bb_upper, bb_lower)

    # Обрезаем датасет
    dataset = dataset.iloc[:len(labels)].copy()
    dataset['labels'] = labels
    dataset = dataset.drop(columns=['bb_upper', 'bb_middle', 'bb_lower'])
    
    return dataset.dropna()

# Fourier labeling
def get_high_frequency_price(close_data, high_pass_cutoff_idx):
    """
    Выделяет высокочастотные компоненты цены с помощью БПФ и ОБПФ.

    :param close_data: Входные данные цены (закрытие).
    :param high_pass_cutoff_idx: Индекс частоты для отсечения.
                                 Компоненты с частотой < этого индекса обнуляются.
    :return: Временной ряд, содержащий только высокочастотные компоненты.
    """
    N = len(close_data)
    
    # 1. Выполнить БПФ
    # Применяем Фурье-преобразование к отклонениям цены от среднего
    # Это помогает уменьшить эффект тренда (нулевая частота)
    mean_close = np.mean(close_data)
    fft_result = fft(close_data - mean_close)
    
    # 2. Отфильтровать низкочастотные компоненты (высокочастотный фильтр)
    # Создаем маску для обнуления низких частот
    fft_filtered = np.copy(fft_result)
    
    # Обнуляем первые 'high_pass_cutoff_idx' коэффициентов (низкие частоты)
    # и их симметричные отражения в конце массива (комплексно-сопряженные частоты).
    
    # Первая половина (положительные частоты)
    for i in range(high_pass_cutoff_idx):
        fft_filtered[i] = 0.0
    
    # Вторая половина (отрицательные частоты, симметричные)
    # Важно: для четного N, fft_result[N/2] соответствует частоте Найквиста и не имеет симметричного
    # Здесь мы обнуляем симметричные низкие частоты
    # N - high_pass_cutoff_idx - это индекс, с которого начинается обнуление симметричных частот
    for i in range(N - high_pass_cutoff_idx + 1, N):
        fft_filtered[i] = 0.0
        
    # 3. Выполнить Обратное БПФ
    # Мы берем только действительную часть, т.к. вход был действительным рядом.
    high_freq_series = np.real(ifft(fft_filtered))
    
    return high_freq_series

@njit
def calculate_labels_fourier(high_freq_data, lookback_period, std_multiplier=1.5):
    """
    Разметка на основе высокочастотных компонент (стандартное отклонение).
    Использует стандартное отклонение высокочастотных данных в скользящем окне
    для определения зон перекупленности/перепроданности.

    :param high_freq_data: Ряд, содержащий только высокочастотные компоненты.
    :param lookback_period: Период для расчета скользящего стандартного отклонения.
    :param std_multiplier: Множитель для стандартного отклонения.
    :return: Массив меток (0.0=Buy, 1.0=Sell, 2.0=Hold).
    """
    N = len(high_freq_data)
    labels = np.zeros(N)
    
    for i in range(lookback_period, N):
        # Окно данных
        window = high_freq_data[i - lookback_period : i]
        
        # Расчет среднего и стандартного отклонения в окне
        mean_window = np.mean(window)
        std_window = np.std(window)
        
        # Пороги для перекупленности/перепроданности
        overbought_threshold = mean_window + std_multiplier * std_window
        oversold_threshold = mean_window - std_multiplier * std_window
        
        current_value = high_freq_data[i]
        
        # Сигнал на покупку: цена в зоне перепроданности (слишком сильный краткосрочный спад)
        if current_value < oversold_threshold:
            labels[i] = 0.0  # Buy
        
        # Сигнал на продажу: цена в зоне перекупленности (слишком сильный краткосрочный рост)
        elif current_value > overbought_threshold:
            labels[i] = 1.0  # Sell
        
        else:
            labels[i] = 2.0  # Hold
    
    return labels

def get_labels_fourier(dataset, lookback_period=20, high_pass_cutoff_idx=5,
                       std_multiplier=1.5) -> pd.DataFrame:
    """
    Основная функция для разметки данных с использованием Фурье-анализа.
    
    :param dataset: Входной датафрейм (должен содержать колонку 'close').
    :param lookback_period: Период скользящего окна для расчета STD.
    :param high_pass_cutoff_idx: Индекс для фильтрации низких частот (чем больше, тем более
                                 высокочастотные компоненты остаются).
    :param std_multiplier: Множитель стандартного отклонения для порогов.
    :return: Датафрейм с добавленной колонкой 'labels'.
    """
    dataset = dataset.copy()
    close_data = dataset['close'].values
    
    # 1. Получаем высокочастотные компоненты
    high_freq_data = get_high_frequency_price(close_data, high_pass_cutoff_idx)
    dataset['high_freq'] = high_freq_data # Опционально для просмотра
    
    # 2. Рассчитываем метки на основе отклонения
    labels = calculate_labels_fourier(high_freq_data, lookback_period, std_multiplier)

    # 3. Обрезаем датасет и добавляем метки
    dataset['labels'] = labels
    dataset = dataset.iloc[lookback_period:].copy()
    dataset = dataset.drop(columns=['high_freq'])
    
    return dataset.dropna()

# RSI labeling
@njit
def calculate_rsi(close_data, period=14):
    """Экспоненциальное сглаживание RSI (метод Уайлдера)"""
    rsi = np.zeros(len(close_data))
    
    # Первый расчет - простое среднее
    gains = 0.0
    losses = 0.0
    for j in range(1, period + 1):
        change = close_data[j] - close_data[j - 1]
        if change > 0:
            gains += change
        else:
            losses += abs(change)
    
    avg_gain = gains / period
    avg_loss = losses / period
    
    if avg_loss == 0:
        rsi[period] = 100.0
    else:
        rs = avg_gain / avg_loss
        rsi[period] = 100.0 - (100.0 / (1.0 + rs))
    
    # Экспоненциальное сглаживание для остальных значений
    for i in range(period + 1, len(close_data)):
        change = close_data[i] - close_data[i - 1]
        gain = max(change, 0.0)
        loss = max(-change, 0.0)
        
        # Экспоненциальное сглаживание Уайлдера
        avg_gain = (avg_gain * (period - 1) + gain) / period
        avg_loss = (avg_loss * (period - 1) + loss) / period
        
        if avg_loss == 0:
            rsi[i] = 100.0
        else:
            rs = avg_gain / avg_loss
            rsi[i] = 100.0 - (100.0 / (1.0 + rs))
    
    return rsi

@njit
def calculate_labels_rsi(close_data, rsi_data, 
                         oversold_level=30.0, overbought_level=70.0):
    """
    Разметка на основе уровней RSI.
    """
    labels = []
    
    for i in range(len(close_data)):
        curr_rsi = rsi_data[i]
        
        # Сигнал на покупку: RSI в зоне перепроданности
        if curr_rsi < oversold_level:
            labels.append(0.0)  # Buy
        
        # Сигнал на продажу: RSI в зоне перекупленности
        elif curr_rsi > overbought_level:
            labels.append(1.0)  # Sell
        
        else:
            labels.append(2.0)  # Hold - RSI в нейтральной зоне
    
    return labels

def get_labels_rsi(dataset, rsi_period=14,
                   oversold_level=30.0, overbought_level=70.0) -> pd.DataFrame:
    dataset = dataset.copy()
    
    close_data = dataset['close'].values
    rsi_data = calculate_rsi(close_data, rsi_period)
    dataset['rsi'] = rsi_data

    labels = calculate_labels_rsi(close_data, rsi_data,
                                  oversold_level, overbought_level)

    # Обрезаем датасет
    dataset = dataset.iloc[:len(labels)].copy()
    dataset['labels'] = labels
    dataset = dataset.drop(columns=['rsi'])
    
    return dataset.dropna()

def get_labels_profit_rsi_profit_check(dataset, rsi_period=14,
                            oversold_level=30.0, overbought_level=70.0,
                            min_forecast_period=1, max_forecast_period=15,
                            markup=0.0) -> pd.DataFrame:
    
    dataset = dataset.copy()
    close_data = dataset['close'].values
    rsi_data = calculate_rsi(close_data, rsi_period)
    
    # Используем NaN для начального заполнения, так как для последнего max_forecast_period 
    # элементов невозможно определить будущую цену.
    labels = [np.nan] * len(close_data) 
    
    # Начинаем итерацию с индекса, после которого RSI уже посчитан
    start_index = rsi_period
    
    # Обходим до конца, который позволяет спрогнозировать будущую цену
    for i in range(start_index, len(close_data) - max_forecast_period):
        
        curr_rsi = rsi_data[i]
        curr_pr = close_data[i]
        
        # 1. Сначала определяем сигнал по RSI
        rsi_signal = 2.0 # Hold по умолчанию
        
        # Сигнал на покупку: RSI в зоне перепроданности
        if curr_rsi < oversold_level:
            rsi_signal = 0.0  # Buy
        
        # Сигнал на продажу: RSI в зоне перекупленности
        elif curr_rsi > overbought_level:
            rsi_signal = 1.0  # Sell
        
        # 2. Если есть сигнал (Buy или Sell), проверяем его прибыльность
        if rsi_signal != 2.0:
            
            # Выбираем случайный период прогноза
            rand_period = random.randint(min_forecast_period, max_forecast_period)
            future_pr = close_data[i + rand_period]
            
            # Проверка прибыльности для сигнала Buy (0.0)
            if rsi_signal == 0.0:
                # Покупка прибыльна, если будущая цена > текущей цены + markup
                if (future_pr - markup) > curr_pr:
                    labels[i] = 0.0  # Buy - Прибыльно
                else:
                    labels[i] = 2.0  # Hold - Не прибыльно
            
            # Проверка прибыльности для сигнала Sell (1.0)
            elif rsi_signal == 1.0:
                # Продажа прибыльна, если будущая цена < текущей цены - markup
                if (future_pr + markup) < curr_pr:
                    labels[i] = 1.0  # Sell - Прибыльно
                else:
                    labels[i] = 2.0  # Hold - Не прибыльно
        
        else:
            # Нет сигнала RSI -> Hold
            labels[i] = 2.0
    
    # Обрезаем датасет и добавляем метки
    dataset['labels'] = labels
    dataset = dataset.iloc[start_index:].copy()
    
    # Очищаем от NaN (в конце, где невозможно определить будущую цену)
    return dataset.dropna()
