from numba import njit
import pandas as pd
import numpy as np

def calc_rsi(close):
    from ta.momentum import RSIIndicator
    return ('rsi', RSIIndicator(close, window=14).rsi())

def calc_macd(close):
    from ta.trend import MACD
    return ('macd', MACD(close).macd())

def calc_bollinger(close):
    from ta.volatility import BollingerBands
    bb = BollingerBands(close=close, window=20, window_dev=2)
    return [
        ('bb_bbm', bb.bollinger_mavg()),
        ('bb_bbh', bb.bollinger_hband()),
        ('bb_bbl', bb.bollinger_lband()),
        ('bb_bb_width', bb.bollinger_hband() - bb.bollinger_lband())
    ]

def calc_stochastic(high, low, close):
    from ta.momentum import StochasticOscillator
    stoch = StochasticOscillator(high=high, low=low, close=close, window=14, smooth_window=3)
    return [
        ('stoch_k', stoch.stoch()),
        ('stoch_d', stoch.stoch_signal())
    ]

def calc_atr(high, low, close):
    from ta.volatility import AverageTrueRange
    atr = AverageTrueRange(high=high, low=low, close=close, window=14)
    return ('atr', atr.average_true_range())

def calc_cci(high, low, close):
    from ta.trend import CCIIndicator
    cci = CCIIndicator(high=high, low=low, close=close, window=20)
    return ('cci', cci.cci())

def calc_williamsr(high, low, close):
    from ta.momentum import WilliamsRIndicator
    willr = WilliamsRIndicator(high=high, low=low, close=close, lbp=14)
    return ('williams_r', willr.williams_r())

def calc_ema(close):
    from ta.trend import EMAIndicator
    ema = EMAIndicator(close=close, window=14)
    return ('ema_14', ema.ema_indicator())

def calc_obv(close, volume):
    from ta.volume import OnBalanceVolumeIndicator
    obv = OnBalanceVolumeIndicator(close=close, volume=volume)
    return ('obv', obv.on_balance_volume())

def calc_cmf(high, low, close, volume):
    from ta.volume import ChaikinMoneyFlowIndicator
    cmf = ChaikinMoneyFlowIndicator(high=high, low=low, close=close, volume=volume, window=20)
    return ('cmf', cmf.chaikin_money_flow())

def calc_sma(close):
    from ta.trend import SMAIndicator
    return [
        ('sma_50', SMAIndicator(close, window=50).sma_indicator()),
        ('sma_200', SMAIndicator(close, window=200).sma_indicator())
    ]

def calc_roc(close):
    from ta.momentum import ROCIndicator
    return ('roc_10', ROCIndicator(close, window=10).roc())

def calc_momentum(close):
    return ('momentum_10', close - close.shift(10))

def calc_psar(high, low, close):
    # Use the Numba-accelerated fast_psar function for speed
    psar_values = fast_psar(np.array(high), np.array(low), np.array(close))
    return [('psar', pd.Series(psar_values, index=close.index))]

def calc_donchian(high, low, close):
    from ta.volatility import DonchianChannel
    donchian = DonchianChannel(high, low, close, window=20)
    return [
        ('donchian_hband', donchian.donchian_channel_hband()),
        ('donchian_lband', donchian.donchian_channel_lband()),
        ('donchian_mband', donchian.donchian_channel_mband())
    ]

def calc_keltner(high, low, close):
    from ta.volatility import KeltnerChannel
    keltner = KeltnerChannel(high, low, close, window=20)
    return [
        ('keltner_hband', keltner.keltner_channel_hband()),
        ('keltner_lband', keltner.keltner_channel_lband()),
        ('keltner_mband', keltner.keltner_channel_mband())
    ]

def calc_dpo(close):
    from ta.trend import DPOIndicator
    return ('dpo_20', DPOIndicator(close, window=20).dpo())

def calc_ultimate(high, low, close):
    from ta.momentum import UltimateOscillator
    return ('ultimate_osc', UltimateOscillator(high, low, close).ultimate_oscillator())

def calc_ichimoku(high, low):
    from ta.trend import IchimokuIndicator
    ichimoku = IchimokuIndicator(high, low, window1=9, window2=26, window3=52)
    return [
        ('ichimoku_a', ichimoku.ichimoku_a()),
        ('ichimoku_b', ichimoku.ichimoku_b()),
        ('ichimoku_base_line', ichimoku.ichimoku_base_line()),
        ('ichimoku_conversion_line', ichimoku.ichimoku_conversion_line())
    ]

def calc_elder_ray(close, low, high):
    from ta.trend import EMAIndicator
    ema = EMAIndicator(close, window=13).ema_indicator()
    return [
        ('elder_ray_bull', ema - low),
        ('elder_ray_bear', ema - high)
    ]

def calc_daily_return(close):
    from ta.others import DailyReturnIndicator
    return ('daily_return', DailyReturnIndicator(close).daily_return())

@njit
def fast_psar(high, low, close, af=0.02, max_af=0.2):
    length = len(close)
    psar = np.zeros(length)
    bull = True
    af_step = af
    ep = low[0]
    psar[0] = low[0]
    for i in range(1, length):
        prev_psar = psar[i-1]
        if bull:
            psar[i] = prev_psar + af_step * (ep - prev_psar)
            if low[i] < psar[i]:
                bull = False
                psar[i] = ep
                af_step = af
                ep = low[i]
            else:
                if high[i] > ep:
                    ep = high[i]
                    af_step = min(af_step + af, max_af)
        else:
            psar[i] = prev_psar + af_step * (ep - prev_psar)
            if high[i] > psar[i]:
                bull = True
                psar[i] = ep
                af_step = af
                ep = high[i]
            else:
                if low[i] < ep:
                    ep = low[i]
                    af_step = min(af_step + af, max_af)
    return psar

def compute_lag(df, col, lag):
    return df[col].shift(lag)

def compute_rolling(df, col, stat, window):
    if stat == 'mean':
        return df[col].rolling(window).mean()
    elif stat == 'std':
        return df[col].rolling(window).std()
    elif stat == 'min':
        return df[col].rolling(window).min()
    elif stat == 'max':
        return df[col].rolling(window).max()

def compute_log_return(df, horizon):
    return np.log(df['Close'] / df['Close'].shift(horizon))

def compute_volatility(df, window):
    return df['log_return'].rolling(window).std()

def run_feature_job(job, df):
    feature_name, func, *args = job
    print(f'Computing feature: {feature_name}')
    result = func(df, *args)
    return feature_name, result

def calc_adx(high, low, close):
    from ta.trend import ADXIndicator
    adx = ADXIndicator(high=high, low=low, close=close, window=14)
    return [
        ('adx', adx.adx()),
        ('adx_pos', adx.adx_pos()),
        ('adx_neg', adx.adx_neg())
    ]

def calc_trix(close):
    from ta.trend import TRIXIndicator
    trix = TRIXIndicator(close=close, window=15)
    return ('trix', trix.trix())

def calc_vortex(high, low, close):
    from ta.trend import VortexIndicator
    vortex = VortexIndicator(high=high, low=low, close=close, window=14)
    return [
        ('vortex_pos', vortex.vortex_indicator_pos()),
        ('vortex_neg', vortex.vortex_indicator_neg())
    ]

def calc_kama(close):
    # Simple alternative to KAMA using EMA
    from ta.trend import EMAIndicator
    kama = EMAIndicator(close, window=10).ema_indicator()
    return ('kama', kama)

def calc_force_index(close, volume):
    from ta.volume import ForceIndexIndicator
    fi = ForceIndexIndicator(close=close, volume=volume, window=13)
    return ('force_index', fi.force_index())

def calc_eom(high, low, volume):
    from ta.volume import EaseOfMovementIndicator
    eom = EaseOfMovementIndicator(high=high, low=low, volume=volume, window=14)
    return ('eom', eom.ease_of_movement())

def calc_mfi(high, low, close, volume):
    from ta.volume import MFIIndicator
    mfi = MFIIndicator(high=high, low=low, close=close, volume=volume, window=14)
    return ('mfi', mfi.money_flow_index())

def calc_adi(high, low, close, volume):
    from ta.volume import AccDistIndexIndicator
    adi = AccDistIndexIndicator(high=high, low=low, close=close, volume=volume)
    return ('adi', adi.acc_dist_index())

def calc_tema(close):
    # Simple alternative to TEMA using triple EMA
    from ta.trend import EMAIndicator
    ema1 = EMAIndicator(close, window=10).ema_indicator()
    ema2 = EMAIndicator(ema1, window=10).ema_indicator()
    ema3 = EMAIndicator(ema2, window=10).ema_indicator()
    tema = 3 * ema1 - 3 * ema2 + ema3
    return ('tema', tema)

def calc_stochrsi(close):
    from ta.momentum import StochRSIIndicator
    stochrsi = StochRSIIndicator(close=close, window=14, smooth1=3, smooth2=3)
    return [
        ('stochrsi', stochrsi.stochrsi()),
        ('stochrsi_k', stochrsi.stochrsi_k()),
        ('stochrsi_d', stochrsi.stochrsi_d())
    ]

def calc_awesome_oscillator(high, low):
    from ta.momentum import AwesomeOscillatorIndicator
    ao = AwesomeOscillatorIndicator(high=high, low=low, window1=5, window2=34)
    return ('awesome_osc', ao.awesome_oscillator())