"""
RSI-Gaussian Channel Strategy Implementation

This module implements the trading logic for the RSI-Gaussian Channel strategy.
It is designed to be used by the validation, optimization, and backtesting scripts.

Key Function:
    generate_signals(df_price, **kwargs) -> pd.DataFrame

Expected Parameters (kwargs):
    gaussian_per (int): Period for the Gaussian Channel.
    gaussian_mult (float): Multiplier for the Gaussian Channel.
    stoch_lower (int): Lower threshold for Stochastic RSI.
    stoch_upper (int): Upper threshold for Stochastic RSI.
    stoch_threshold (int, optional): Symmetric threshold (overrides lower/upper).
"""
import numpy as np
import pandas as pd
from numba import njit
import gaussian_logic

@njit
def rma_numba(src, length):
    """
    Pine Script's RMA (Running Moving Average).
    Equivalent to Wilder's Smoothing.
    alpha = 1 / length
    """
    alpha = 1.0 / length
    out = np.full_like(src, np.nan)
    
    # Initialization: Pine Script's RMA usually propagates NaN until valid data.
    # Standard RSI initialization uses SMA for the first value.
    # However, strictly speaking, RMA is: out = alpha * src + (1 - alpha) * nz(out[1])
    # If we strictly follow `nz`, the first value is alpha * src[0].
    # But `ta.rsi` in TradingView waits for `length` bars.
    # We will use the standard Wilder's initialization (SMA of first length values)
    # which matches `ta.rsi` output.
    
    # Find first valid index
    valid_start = 0
    for i in range(len(src)):
        if not np.isnan(src[i]):
            valid_start = i
            break
            
    if len(src) - valid_start < length:
        return out

    # Initialize with SMA
    sum_val = 0.0
    for i in range(valid_start, valid_start + length):
        sum_val += src[i]
    
    out[valid_start + length - 1] = sum_val / length
    
    # Calculate rest
    for i in range(valid_start + length, len(src)):
        prev = out[i-1]
        curr = src[i]
        if np.isnan(curr):
            out[i] = prev # Hold value? Or NaN? Pine RMA with NaN input updates with 0? 
                          # Usually price data isn't NaN.
            continue
            
        out[i] = alpha * curr + (1.0 - alpha) * prev
        
    return out

@njit
def rsi_numba(close, length):
    """
    Calculates RSI using RMA smoothing to match Pine Script `ta.rsi`.
    """
    delta = np.zeros_like(close)
    delta[1:] = close[1:] - close[:-1]
    
    gain = np.where(delta > 0, delta, 0.0)
    loss = np.where(delta < 0, -delta, 0.0)
    
    avg_gain = rma_numba(gain, length)
    avg_loss = rma_numba(loss, length)
    
    rsi = np.full_like(close, np.nan)
    
    for i in range(len(close)):
        if np.isnan(avg_loss[i]):
            continue
            
        if avg_loss[i] == 0:
            if avg_gain[i] == 0:
                rsi[i] = 50.0 # No move
            else:
                rsi[i] = 100.0
        else:
            rs = avg_gain[i] / avg_loss[i]
            rsi[i] = 100.0 - (100.0 / (1.0 + rs))
            
    return rsi

@njit
def stoch_numba(src, high, low, length):
    """
    Calculates Stochastic Oscillator.
    100 * (src - lowest(low, length)) / (highest(high, length) - lowest(low, length))
    """
    stoch = np.full_like(src, np.nan)
    
    for i in range(length - 1, len(src)):
        # Window for min/max
        window_low = low[i - length + 1 : i + 1]
        window_high = high[i - length + 1 : i + 1]
        
        min_val = np.min(window_low)
        max_val = np.max(window_high)
        
        if max_val == min_val:
            stoch[i] = 0.0 # Avoid div by zero
        else:
            stoch[i] = 100.0 * (src[i] - min_val) / (max_val - min_val)
            
    return stoch

@njit
def sma_numba(src, length):
    """Simple Moving Average"""
    out = np.full_like(src, np.nan)
    for i in range(length - 1, len(src)):
        out[i] = np.mean(src[i - length + 1 : i + 1])
    return out

@njit
def calculate_positions(entries, exits):
    """
    Determines executed trades based on position state.
    Replicates TradingView strategy behavior: Entry only if flat, Exit only if long.
    """
    n = len(entries)
    in_pos = np.zeros(n, dtype=np.bool_)
    exec_entry = np.zeros(n, dtype=np.bool_)
    exec_exit = np.zeros(n, dtype=np.bool_)
    
    currently_in_pos = False
    
    for i in range(n):
        if currently_in_pos:
            if exits[i]:
                exec_exit[i] = True
                currently_in_pos = False
        else:
            if entries[i]:
                exec_entry[i] = True
                currently_in_pos = True
        
        in_pos[i] = currently_in_pos
        
    return in_pos, exec_entry, exec_exit

def get_stoch_rsi(close, lengthRSI, lengthStoch, smoothK):
    """
    Calculates Stochastic RSI.
    rsi1 = ta.rsi(src_stoch, lengthRSI)
    k_val = ta.sma(ta.stoch(rsi1, rsi1, rsi1, lengthStoch), smoothK)
    """
    # 1. RSI
    rsi1 = rsi_numba(close, lengthRSI)
    
    # 2. Stoch of RSI
    # ta.stoch(rsi1, rsi1, rsi1, lengthStoch) means src=rsi1, high=rsi1, low=rsi1
    stoch_raw = stoch_numba(rsi1, rsi1, rsi1, lengthStoch)
    
    # 3. Smooth K (SMA)
    k_val = sma_numba(stoch_raw, smoothK)
    
    return k_val

def generate_signals(df_price, gaussian_per=144, gaussian_mult=1.414, stoch_lower=20, stoch_upper=80, stoch_threshold=None):
    """
    Generates Entry and Exit signals based on the strategy logic.
    """
    # Handle optional threshold logic for symmetry
    if stoch_threshold is not None:
        stoch_lower = stoch_threshold
        stoch_upper = 100 - stoch_threshold

    # Extract arrays
    open_arr = df_price['open'].values.astype(np.float64)
    high_arr = df_price['high'].values.astype(np.float64)
    low_arr = df_price['low'].values.astype(np.float64)
    close_arr = df_price['close'].values.astype(np.float64)
    time_arr = df_price['time'].values
    
    # --- Gaussian Channel ---
    # Defaults: N=4, per=144, mult=1.414, modeLag=False, modeFast=False
    filt, hband, lband = gaussian_logic.get_gaussian_channel(
        open_arr, high_arr, low_arr, close_arr, 
        N=4, per=gaussian_per, mult=gaussian_mult, modeLag=False, modeFast=False
    )
    
    # --- Stochastic RSI ---
    # Defaults: smoothK=3, lengthRSI=14, lengthStoch=14
    stoch_value = get_stoch_rsi(close_arr, lengthRSI=14, lengthStoch=14, smoothK=3)
    
    # --- Trading Logic ---
    
    # 1. Gaussian Green: filt > filt[1]
    # Shift filt by 1 to compare
    filt_prev = np.roll(filt, 1)
    filt_prev[0] = np.nan
    gaussian_green = filt > filt_prev
    
    # 2. Price Above Hband
    price_above_hband = close_arr > hband
    
    # 3. Stoch Condition: > 80 or < 20
    stoch_condition = (stoch_value > stoch_upper) | (stoch_value < stoch_lower)
    
    # 4. Long Condition (Entry)
    # All must be true
    long_condition = gaussian_green & price_above_hband & stoch_condition
    
    # 5. Close All Condition (Exit)
    # ta.crossunder(close, hband) -> close < hband AND close[1] >= hband[1]
    close_prev = np.roll(close_arr, 1)
    hband_prev = np.roll(hband, 1)
    # Handle first element
    close_prev[0] = np.nan 
    
    cross_under = (close_arr < hband) & (close_prev >= hband_prev)
    close_all_condition = cross_under
    
    # --- Position Management (Stateful) ---
    in_pos, exec_entry, exec_exit = calculate_positions(long_condition, close_all_condition)
    
    # Create result DataFrame
    df_signals = pd.DataFrame({
        'time': time_arr,
        'close': close_arr,
        'low': low_arr,
        'filt': filt,
        'hband': hband,
        'stoch_rsi': stoch_value,
        'entry_long': long_condition,      # Raw Signal
        'exit_long': close_all_condition,  # Raw Signal
        'in_trade': in_pos,                # State
        'execute_entry': exec_entry,       # Executed Trade
        'execute_exit': exec_exit          # Executed Trade
    })
    
    return df_signals