from __future__ import annotations

import numpy as np
import pandas as pd


def rma(series: pd.Series, length: int) -> pd.Series:
    alpha = 1.0 / length
    result = series.astype(float).copy()
    for i in range(1, len(series)):
        prev = result.iat[i - 1]
        curr = series.iat[i]
        result.iat[i] = alpha * curr + (1 - alpha) * prev
    return result


def calculate_atr(df: pd.DataFrame, length: int) -> pd.Series:
    high, low, close = df["high"], df["low"], df["close"]
    tr = pd.concat([
        (high - low),
        (high - close.shift(1)).abs(),
        (low - close.shift(1)).abs(),
    ], axis=1).max(axis=1)
    tr.iloc[0] = (high.iloc[0] - low.iloc[0])
    return rma(tr, length)


def calculate_adx(df: pd.DataFrame, length: int) -> pd.Series:
    high, low, close = df["high"], df["low"], df["close"]
    up_move = high.diff()
    down_move = -low.diff()
    plus_dm = np.where((up_move > down_move) & (up_move > 0), up_move, 0.0)
    minus_dm = np.where((down_move > up_move) & (down_move > 0), down_move, 0.0)
    tr = pd.concat([
        (high - low),
        (high - close.shift(1)).abs(),
        (low - close.shift(1)).abs(),
    ], axis=1).max(axis=1)
    tr.iloc[0] = (high.iloc[0] - low.iloc[0])

    atr = rma(tr, length)
    plus_di = 100.0 * rma(pd.Series(plus_dm), length) / atr
    minus_di = 100.0 * rma(pd.Series(minus_dm), length) / atr
    dx = 100.0 * (plus_di - minus_di).abs() / (plus_di + minus_di).replace(0, np.nan)
    dx = dx.fillna(0.0)
    return rma(dx, length)