import numpy as np
from numba import njit
import math

@njit
def calculate_alpha_beta(N, per, mult):
    """
    Calculates alpha and beta for the Gaussian filter exactly as in Pine Script.
    """
    # beta  = (1 - math.cos(4 * math.asin(1) / per)) / (math.pow(1.414, 2 / N) - 1)
    # 4 * asin(1) is 2 * pi
    beta = (1.0 - math.cos(2.0 * math.pi / per)) / (math.pow(1.414, 2.0 / N) - 1.0)
    alpha = -beta + math.sqrt(math.pow(beta, 2.0) + 2.0 * beta)
    return alpha

@njit
def binomial_coeff(n, k):
    """Calculates binomial coefficient (n choose k)."""
    if k < 0 or k > n:
        return 0
    if k == 0 or k == n:
        return 1
    k = min(k, n - k)  # Take advantage of symmetry
    c = 1
    for i in range(k):
        c = c * (n - i) // (i + 1)
    return c

@njit
def gaussian_filter_numba(src, N, alpha):
    """
    Implements the N-pole Gaussian filter recursive logic.
    Replicates f_filt9x logic:
    _f := math.pow(_a, _i) * nz(_s) +
          _i * (1 - _a) * nz(_f[1]) -
          (_i >= 2 ? _m2 * math.pow((1 - _a), 2) * nz(_f[2]) : 0) + ...
    """
    length = len(src)
    output = np.zeros(length, dtype=np.float64)
    
    # Pre-calculate coefficients
    # The formula is: y[t] = alpha^N * x[t] + sum_{k=1}^{N} ( (-1)^(k-1) * binom(N, k) * (1-alpha)^k * y[t-k] )
    
    alpha_n = math.pow(alpha, N)
    one_minus_alpha = 1.0 - alpha
    
    coeffs = np.zeros(N + 1, dtype=np.float64)
    for k in range(1, N + 1):
        # Calculate term: (-1)^(k-1) * binom(N, k) * (1-alpha)^k
        sign = 1.0 if (k - 1) % 2 == 0 else -1.0
        binom = binomial_coeff(N, k)
        term = sign * binom * math.pow(one_minus_alpha, k)
        coeffs[k] = term

    # Iterate through the array
    for i in range(length):
        # Source term: alpha^N * src[i]
        # nz(src) check: if src is nan, treat as 0 (Pine Script nz behavior)
        s_val = src[i]
        if np.isnan(s_val):
            s_val = 0.0
            
        val = alpha_n * s_val
        
        # Recursive terms
        for k in range(1, N + 1):
            prev_idx = i - k
            if prev_idx >= 0:
                prev_val = output[prev_idx]
                # nz(_f[k]) check is implicit since we initialized output with 0.0
                val += coeffs[k] * prev_val
            # else: prev_val is 0.0 (boundary condition)
            
        output[i] = val
        
    return output

@njit
def calculate_true_range(high, low, close):
    length = len(close)
    tr = np.zeros(length, dtype=np.float64)
    tr[0] = high[0] - low[0] # First TR is usually H-L
    for i in range(1, length):
        hl = high[i] - low[i]
        hc = abs(high[i] - close[i-1])
        lc = abs(low[i] - close[i-1])
        tr[i] = max(hl, max(hc, lc))
    return tr

def get_gaussian_channel(open_arr, high_arr, low_arr, close_arr, N=4, per=144, mult=1.414, modeLag=False, modeFast=False):
    """
    Main entry point to calculate Gaussian Channel bands.
    Returns: filt, hband, lband
    """
    # 1. Calculate Source (hlc3)
    src = (high_arr + low_arr + close_arr) / 3.0
    
    # 2. Calculate True Range
    tr = calculate_true_range(high_arr, low_arr, close_arr)
    
    # 3. Calculate Alpha
    alpha = calculate_alpha_beta(N, per, mult)
    
    # 4. Handle Lag Mode
    lag = int((per - 1) / (2 * N))
    
    if modeLag:
        # srcdata = src + (src - src[lag])
        # trdata = tr + (tr - tr[lag])
        # We need to shift arrays. src[lag] means src shifted right by lag.
        src_shifted = np.roll(src, lag)
        src_shifted[:lag] = np.nan # Pine Script behavior for history lookup out of bounds is usually NaN
        
        tr_shifted = np.roll(tr, lag)
        tr_shifted[:lag] = np.nan
        
        # Pine Script arithmetic with NaN results in NaN, but nz() is used inside the filter.
        # However, the calculation `src - src[lag]` happens *before* the filter.
        # If src[lag] is NaN, the result is NaN.
        srcdata = src + (src - src_shifted)
        trdata = tr + (tr - tr_shifted)
        
        # Replace NaNs with 0 for the filter input, as Pine's nz() would handle it inside,
        # but we want clean input arrays.
        # Actually, Pine's `srcdata` variable would hold NaNs at the start.
        # The filter function `f_filt9x` calls `nz(_s)` on the input.
        # So we leave NaNs in `srcdata` and let `gaussian_filter_numba` handle them.
    else:
        srcdata = src
        trdata = tr

    # 5. Calculate Filters
    # f_pole returns [_fn, _f1]. We need both N-pole and 1-pole outputs.
    
    # Filter for Source
    filtn = gaussian_filter_numba(srcdata, N, alpha)
    filt1 = gaussian_filter_numba(srcdata, 1, alpha)
    
    # Filter for True Range
    filtntr = gaussian_filter_numba(trdata, N, alpha)
    filt1tr = gaussian_filter_numba(trdata, 1, alpha)
    
    # 6. Combine based on modeFast
    if modeFast:
        filt = (filtn + filt1) / 2.0
        filttr = (filtntr + filt1tr) / 2.0
    else:
        filt = filtn
        filttr = filtntr
        
    # 7. Calculate Bands
    hband = filt + filttr * mult
    lband = filt - filttr * mult
    
    return filt, hband, lband