"""
=========================================
CV6 AI Trading OS
Average True Range (ATR)
=========================================
"""

import pandas as pd


class ATRIndicator:

    @staticmethod
    def calculate(
        data: pd.DataFrame,
        period: int = 14
    ):

        high = data["High"]
        low = data["Low"]
        close = data["Close"]

        tr1 = high - low
        tr2 = (high - close.shift()).abs()
        tr3 = (low - close.shift()).abs()

        tr = pd.concat(
            [tr1, tr2, tr3],
            axis=1
        ).max(axis=1)

        atr = tr.rolling(period).mean()

        return atr

    @staticmethod
    def latest(
        data: pd.DataFrame,
        period: int = 14
    ):

        atr = ATRIndicator.calculate(
            data,
            period
        )

        return atr.iloc[-1]