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):
    """Calculate CCI indicator based on close only"""
    cci = np.zeros(len(close_data))

    for i in range(period - 1, len(close_data)):
        # Simple moving average
        sma = 0.0
        for j in range(i - period + 1, i + 1):
            sma += close_data[j]
        sma = sma / period

        # Average absolute deviation
        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 equation
        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 level-based labeling.
    """
    labels = []

    for i in range(len(close_data)):
        curr_cci = cci_data[i]

        # Buy signal: CCI in the oversold zone
        if curr_cci < oversold_level:
            labels.append(0.0)  # Buy

        # Sell signal: CCI in the overbought zone
        elif curr_cci > overbought_level:
            labels.append(1.0)  # Sell

        else:
            labels.append(2.0)  # Hold - CCI in the neutral zone

    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
    )

    # Cut the dataset
    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):
    """Calculate Stochastic Oscillator (%K and %D)"""
    stoch_k = np.zeros(len(close_data))
    stoch_d = np.zeros(len(close_data))

    # Calculate %K
    for i in range(period - 1, len(close_data)):
        # Find max and min within '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 equation
        if highest - lowest == 0:
            stoch_k[i] = 50.0
        else:
            stoch_k[i] = 100.0 * (close_data[i] - lowest) / (highest - lowest)

    # Smooth %K for receiving %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-based labeling.
    """
    labels = []

    for i in range(len(close_data)):
        curr_k = stoch_k[i]
        curr_d = stoch_d[i]

        # Buy signal: %K in the oversold zone and crosses %D upwards
        if curr_k < oversold_level:
            labels.append(0.0)  # Buy

        # Sell signal: %K in the overbought zone and crosses %D downwards
        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
    )

    # Cut dataset
    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):
    """Calculate Bollinger bands"""
    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)):
        # Moving average (SMA)
        sma = 0.0
        for j in range(i - period + 1, i + 1):
            sma += close_data[j]
        sma = sma / period

        # Standard deviation
        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):
    """
    Bollinger bands-based labeling.
    If the price is above the upper band -> sell (1.0)
    If the price is below the lower band -> buy (0.0)
    Otherwise -> hold (2.0)
    """
    labels = []

    for i in range(len(close_data)):
        price = close_data[i]
        upper = bb_upper[i]
        lower = bb_lower[i]

        # Price above the upper band - sell
        if price > upper:
            labels.append(1.0)  # Sell

        # Price below the lower band - buy
        elif price < lower:
            labels.append(0.0)  # Buy

        # Price inside the bands
        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)

    # Cut dataset
    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):
    """
    Highlight high-frequency price components using FFT and RFFT.

    :param close_data: Price incoming data (closing).
    :param high_pass_cutoff_idx: Frequency index for cutting.
                                 Components with the frequency < this index are reset.
    :return: TIme series containing only high-frequency components.
    """
    N = len(close_data)

    # 1. Perform FFT
    # Apply Fourier transform to price deviations from the average
    # This decreases the trend effect (zero frequency)
    mean_close = np.mean(close_data)
    fft_result = fft(close_data - mean_close)

    # 2. Sort out low-frequency components (high-frequency filter)
    # Create the mask for resetting low frequencies
    fft_filtered = np.copy(fft_result)

    # Reset the first 'high_pass_cutoff_idx' coefficients (low frequencies)
    # and their symmetrical reflections at the end of the array (complex conjugate frequencies).

    # First half (positive frequencies)
    for i in range(high_pass_cutoff_idx):
        fft_filtered[i] = 0.0

    # Second half (negative frequencies, symmetrical)
    # Important: for even N, fft_result[N/2] corresponds to the Nyquist frequency and has no symmetrical
    # Here we reset symmetrical low frequencies
    # N - high_pass_cutoff_idx - this is the index the reset of symmetrical frequencies starts from
    for i in range(N - high_pass_cutoff_idx + 1, N):
        fft_filtered[i] = 0.0

    # 3. Perform Reverse FFT
    # We take only the real part, since the entruy was real nearby.
    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):
    """
    Labeling based on high-frequency components (standard deviation).
    Uses standard deviation of high-frequency data in a sliding window
    to define overbought/oversold zones.

    :param high_freq_data: Series containing only high-frequency components.
    :param lookback_period: Period for calculating only moving standard deviation.
    :param std_multiplier: Standard deviation multiplier.
    :return: Array of labels (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):
        # Data window
        window = high_freq_data[i - lookback_period : i]

        # Calculate average and standard deviation in the window
        mean_window = np.mean(window)
        std_window = np.std(window)

        # Overbought/oversold thresholds
        overbought_threshold = mean_window + std_multiplier * std_window
        oversold_threshold = mean_window - std_multiplier * std_window

        current_value = high_freq_data[i]

        # Buy signal: price in the oversold zone (too strong short-term fall)
        if current_value < oversold_threshold:
            labels[i] = 0.0  # Buy

        # Sell signal: price in the overbought zone (too strong short-term rise)
        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:
    """
    Main function for data labeling using Fourier analysis.

    :param dataset: Input dataframe (should contain the 'close' column).
    :param lookback_period: Sliding window period for calculating STD.
    :param high_pass_cutoff_idx: Index for sorting low frequencies (the more the higher the
                                 frequency of the remaining components).
    :param std_multiplier: Standard deviation multiplier for thresholds.
    :return: Dataframe with the added 'labels' column.
    """
    dataset = dataset.copy()
    close_data = dataset["close"].values

    # 1. Get high-frequency components
    high_freq_data = get_high_frequency_price(close_data, high_pass_cutoff_idx)
    dataset["high_freq"] = high_freq_data  # Optionally for viewing

    # 2. Calculate deviation-based labels
    labels = calculate_labels_fourier(high_freq_data, lookback_period, std_multiplier)

    # 3. Cut dataset and add labels
    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 exponential smoothing (Wilder method)"""
    rsi = np.zeros(len(close_data))

    # FIrst calculation - simple average
    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))

    # Exponential smoothing for the remaining values
    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)

        # Wilder exponential smoothing
        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 level-based labeling.
    """
    labels = []

    for i in range(len(close_data)):
        curr_rsi = rsi_data[i]

        # Buy signal: RSI in the oversold zone
        if curr_rsi < oversold_level:
            labels.append(0.0)  # Buy

        # Sell signal: RSI in the overbought zone
        elif curr_rsi > overbought_level:
            labels.append(1.0)  # Sell

        else:
            labels.append(2.0)  # Hold - RSI in the neutral zone

    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
    )

    # Cut dataset
    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)

    # Use NaN for initial filling, since for the last max_forecast_period
    # elements, it is impossible to define the future price.
    labels = [np.nan] * len(close_data)

    # Start iteration from the index, after which RSI has already been calculated
    start_index = rsi_period

    # Iterate till the end for forecasting the future price
    for i in range(start_index, len(close_data) - max_forecast_period):

        curr_rsi = rsi_data[i]
        curr_pr = close_data[i]

        # 1. First define the signal by RSI
        rsi_signal = 2.0  # Default Hold

        # Buy signal: RSI in the oversold zone
        if curr_rsi < oversold_level:
            rsi_signal = 0.0  # Buy

        # Sell signal: RSI in the overbought zone
        elif curr_rsi > overbought_level:
            rsi_signal = 1.0  # Sell

        # 2. If there is a signal (Buy or Sell), check its profitability
        if rsi_signal != 2.0:

            # Select random forecast period
            rand_period = random.randint(min_forecast_period, max_forecast_period)
            future_pr = close_data[i + rand_period]

            # Check profitability for Buy signal (0.0)
            if rsi_signal == 0.0:
                # Purchase is profitable if the future price > current price + markup
                if (future_pr - markup) > curr_pr:
                    labels[i] = 0.0  # Buy - Profitable
                else:
                    labels[i] = 2.0  # Hold - Unprofitable

            # Check profitability for Sell signal (1.0)
            elif rsi_signal == 1.0:
                # Purchase is profitable if the future price < current price - markup
                if (future_pr + markup) < curr_pr:
                    labels[i] = 1.0  # Sell - Profitable
                else:
                    labels[i] = 2.0  # Hold - Unprofitable

        else:
            # No RSI signal -> Hold
            labels[i] = 2.0

    # Cut dataset and add labels
    dataset["labels"] = labels
    dataset = dataset.iloc[start_index:].copy()

    # Clear from NaN (in the end where it is impossible to define the future price)
    return dataset.dropna()
