"""
自定义指标库 — MQL5 转换后的 Python 指标

所有函数遵循统一规范：
  - 输入: np.ndarray (float64), 参数用关键字参数带默认值
  - 输出: np.ndarray (float64), 长度 = 输入长度, warmup 期填 np.nan
  - 多值输出: 返回 tuple of np.ndarray

使用方式:
  from my_indicators import cci, williams_r
  cci_values = cci(high, low, close, period=20)
"""

import numpy as np
import raptorbt


def typical_price(high: np.ndarray, low: np.ndarray, close: np.ndarray) -> np.ndarray:
    return (high + low + close) / 3.0


def cci(high: np.ndarray, low: np.ndarray, close: np.ndarray, period: int = 20) -> np.ndarray:
    tp = typical_price(high, low, close)
    result = np.full(len(close), np.nan)
    for i in range(period - 1, len(close)):
        window = tp[i - period + 1 : i + 1]
        tp_sma = np.mean(window)
        mean_dev = np.mean(np.abs(window - tp_sma))
        result[i] = (tp[i] - tp_sma) / (0.015 * mean_dev) if mean_dev > 1e-10 else 0.0
    return result


def williams_r(high: np.ndarray, low: np.ndarray, close: np.ndarray, period: int = 14) -> np.ndarray:
    result = np.full(len(close), np.nan)
    for i in range(period - 1, len(close)):
        hh = np.max(high[i - period + 1 : i + 1])
        ll = np.min(low[i - period + 1 : i + 1])
        result[i] = (hh - close[i]) / (hh - ll) * -100.0 if (hh - ll) > 1e-10 else -50.0
    return result


def roc(data: np.ndarray, period: int = 12) -> np.ndarray:
    result = np.full(len(data), np.nan)
    for i in range(period, len(data)):
        if abs(data[i - period]) > 1e-10:
            result[i] = (data[i] - data[i - period]) / data[i - period] * 100.0
    return result


def trix(data: np.ndarray, period: int = 15) -> np.ndarray:
    ema1 = raptorbt.ema(data, period)
    ema2 = raptorbt.ema(ema1, period)
    ema3 = raptorbt.ema(ema2, period)
    result = np.full(len(data), np.nan)
    for i in range(1, len(data)):
        if not np.isnan(ema3[i]) and not np.isnan(ema3[i - 1]) and abs(ema3[i - 1]) > 1e-10:
            result[i] = (ema3[i] - ema3[i - 1]) / ema3[i - 1] * 10000.0
    return result


def dmi(high: np.ndarray, low: np.ndarray, close: np.ndarray, period: int = 14):
    n = len(close)
    plus_dm = np.zeros(n)
    minus_dm = np.zeros(n)
    tr = np.zeros(n)

    for i in range(1, n):
        up_move = high[i] - high[i - 1]
        down_move = low[i - 1] - low[i]
        plus_dm[i] = up_move if (up_move > down_move and up_move > 0) else 0
        minus_dm[i] = down_move if (down_move > up_move and down_move > 0) else 0

        hl = high[i] - low[i]
        hc = abs(high[i] - close[i - 1])
        lc = abs(low[i] - close[i - 1])
        tr[i] = max(hl, hc, lc)

    atr_vals = raptorbt.atr(high, low, close, period)
    plus_di = np.full(n, np.nan)
    minus_di = np.full(n, np.nan)
    adx_vals = np.full(n, np.nan)

    smooth_plus_dm = np.zeros(n)
    smooth_minus_dm = np.zeros(n)
    smooth_tr = np.zeros(n)

    if n > period:
        smooth_plus_dm[period] = np.sum(plus_dm[1 : period + 1])
        smooth_minus_dm[period] = np.sum(minus_dm[1 : period + 1])
        smooth_tr[period] = np.sum(tr[1 : period + 1])

        for i in range(period + 1, n):
            smooth_plus_dm[i] = smooth_plus_dm[i - 1] - smooth_plus_dm[i - 1] / period + plus_dm[i]
            smooth_minus_dm[i] = smooth_minus_dm[i - 1] - smooth_minus_dm[i - 1] / period + minus_dm[i]
            smooth_tr[i] = smooth_tr[i - 1] - smooth_tr[i - 1] / period + tr[i]

        for i in range(period, n):
            if smooth_tr[i] > 1e-10:
                plus_di[i] = smooth_plus_dm[i] / smooth_tr[i] * 100.0
                minus_di[i] = smooth_minus_dm[i] / smooth_tr[i] * 100.0

    dx = np.full(n, np.nan)
    for i in range(period, n):
        di_sum = plus_di[i] + minus_di[i]
        if di_sum > 1e-10:
            dx[i] = abs(plus_di[i] - minus_di[i]) / di_sum * 100.0

    if n > 2 * period - 1:
        adx_vals[2 * period - 1] = np.mean(dx[period : 2 * period])
        for i in range(2 * period, n):
            if not np.isnan(dx[i]) and not np.isnan(adx_vals[i - 1]):
                adx_vals[i] = (adx_vals[i - 1] * (period - 1) + dx[i]) / period

    return plus_di, minus_di, adx_vals


def ichimoku(high: np.ndarray, low: np.ndarray, close: np.ndarray,
             tenkan_period: int = 9, kijun_period: int = 26, senkou_b_period: int = 52):
    n = len(close)
    tenkan = np.full(n, np.nan)
    kijun = np.full(n, np.nan)
    senkou_a = np.full(n, np.nan)
    senkou_b = np.full(n, np.nan)

    for i in range(tenkan_period - 1, n):
        tenkan[i] = (np.max(high[i - tenkan_period + 1 : i + 1]) + np.min(low[i - tenkan_period + 1 : i + 1])) / 2.0

    for i in range(kijun_period - 1, n):
        kijun[i] = (np.max(high[i - kijun_period + 1 : i + 1]) + np.min(low[i - kijun_period + 1 : i + 1])) / 2.0

    for i in range(kijun_period - 1, n):
        if not np.isnan(tenkan[i]) and not np.isnan(kijun[i]):
            senkou_a[i] = (tenkan[i] + kijun[i]) / 2.0

    for i in range(senkou_b_period - 1, n):
        senkou_b[i] = (np.max(high[i - senkou_b_period + 1 : i + 1]) + np.min(low[i - senkou_b_period + 1 : i + 1])) / 2.0

    chikou = close.copy()

    return tenkan, kijun, senkou_a, senkou_b, chikou


def parabolic_sar(high: np.ndarray, low: np.ndarray, close: np.ndarray,
                  step: float = 0.02, max_step: float = 0.2) -> np.ndarray:
    n = len(close)
    if n < 2:
        return np.full(n, np.nan)

    sar = np.full(n, np.nan)
    af = step
    is_long = close[1] > close[0]
    ep = high[1] if is_long else low[1]
    sar[1] = low[0] if is_long else high[0]

    for i in range(2, n):
        sar[i] = sar[i - 1] + af * (ep - sar[i - 1])

        if is_long:
            sar[i] = min(sar[i], low[i - 1], low[i - 2] if i >= 2 else low[i - 1])
            if low[i] < sar[i]:
                is_long = False
                sar[i] = ep
                af = step
                ep = low[i]
            else:
                if high[i] > ep:
                    ep = high[i]
                    af = min(af + step, max_step)
        else:
            sar[i] = max(sar[i], high[i - 1], high[i - 2] if i >= 2 else high[i - 1])
            if high[i] > sar[i]:
                is_long = True
                sar[i] = ep
                af = step
                ep = high[i]
            else:
                if low[i] < ep:
                    ep = low[i]
                    af = min(af + step, max_step)

    return sar


def awesome_oscillator(high: np.ndarray, low: np.ndarray,
                       fast_period: int = 5, slow_period: int = 34) -> np.ndarray:
    midpoint = (high + low) / 2.0
    sma_fast = raptorbt.sma(midpoint, fast_period)
    sma_slow = raptorbt.sma(midpoint, slow_period)
    result = np.full(len(high), np.nan)
    for i in range(len(high)):
        if not np.isnan(sma_fast[i]) and not np.isnan(sma_slow[i]):
            result[i] = sma_fast[i] - sma_slow[i]
    return result


def mfi(high: np.ndarray, low: np.ndarray, close: np.ndarray,
        volume: np.ndarray, period: int = 14) -> np.ndarray:
    tp = typical_price(high, low, close)
    result = np.full(len(close), np.nan)
    for i in range(period, len(close)):
        pos_flow = 0.0
        neg_flow = 0.0
        for j in range(i - period + 1, i + 1):
            mf = tp[j] * volume[j]
            if tp[j] > tp[j - 1]:
                pos_flow += mf
            elif tp[j] < tp[j - 1]:
                neg_flow += mf
        if neg_flow > 1e-10:
            result[i] = 100.0 - 100.0 / (1.0 + pos_flow / neg_flow)
        else:
            result[i] = 100.0
    return result


def obv(close: np.ndarray, volume: np.ndarray) -> np.ndarray:
    result = np.zeros(len(close))
    result[0] = volume[0]
    for i in range(1, len(close)):
        if close[i] > close[i - 1]:
            result[i] = result[i - 1] + volume[i]
        elif close[i] < close[i - 1]:
            result[i] = result[i - 1] - volume[i]
        else:
            result[i] = result[i - 1]
    return result