import pandas as pd
import numpy as np


def rsi_tradingview(ohlc: pd.DataFrame, period: int = 14, round_rsi: bool = True):
    """ Implements the RSI indicator as defined by TradingView on March 15, 2021.
        The TradingView code is as follows:
        //@version=4
        study(title="Relative Strength Index", shorttitle="RSI", format=format.price, precision=2, resolution="")
        len = input(14, minval=1, title="Length")
        src = input(close, "Source", type = input.source)
        up = rma(max(change(src), 0), len)
        down = rma(-min(change(src), 0), len)
        rsi = down == 0 ? 100 : up == 0 ? 0 : 100 - (100 / (1 + up / down))
        plot(rsi, "RSI", color=#8E1599)
        band1 = hline(70, "Upper Band", color=#C0C0C0)
        band0 = hline(30, "Lower Band", color=#C0C0C0)
        fill(band1, band0, color=#9915FF, transp=90, title="Background")

    :param ohlc:
    :param period:
    :param round_rsi:
    :return: an array with the RSI indicator values
    """

    delta = ohlc["close"].diff()

    up = delta.copy()
    up[up < 0] = 0
    up = pd.Series.ewm(up, alpha=1/period).mean()

    down = delta.copy()
    down[down > 0] = 0
    down *= -1
    down = pd.Series.ewm(down, alpha=1/period).mean()

    rsi = np.where(up == 0, 0, np.where(down == 0, 100, 100 - (100 / (1 + up / down))))

    return np.round(rsi, 2) if round_rsi else rsi

def rr(ohlc: pd.DataFrame, period: int = 14, round_rsi: bool = True):
    lst = [
        [1245,39494.5],
        [1250,39547.5],
        [1255,39522.0],
        [1300,39495.0],
        [1305,39146.0],
        [1310,39280.5],
        [1315,39380.0],
        [1320,39304.5],
        [1325,39257.0],
        [1330,39316.0],
        [1335,39241.5],
        [1340,39200.0],
        [1345,39235.5],
        [1350,39157.5],
    ]
    N = 14
    close_price0 = float(lst[0][1])
    gain_avg0 = loss_avg0 = close_price0
    for kline in lst[1:]:
        close_price = float(kline[1])
        if close_price > close_price0:
            gain = close_price - close_price0
            loss = 0
        else:
            gain = 0
            loss = close_price0 - close_price
        close_price0 = close_price
        gain_avg = (gain_avg0 * (N - 1) + gain) / N
        loss_avg = (loss_avg0 * (N - 1) + loss) / N
        rsi = 100 - 100 / (1 + gain_avg / loss_avg)
        gain_avg0 = gain_avg
        loss_avg0 = loss_avg
        print (rsi)

def rsi_tradingview2(ohlc: pd.DataFrame, period: int = 14) -> pd.Series:
    """See source https://github.com/peerchemist/finta
    and fix https://www.tradingview.com/wiki/Talk:Relative_Strength_Index_(RSI)
    Relative Strength Index (RSI) is a momentum oscillator that measures the speed and change of price movements.
    RSI oscillates between zero and 100. Traditionally, and according to Wilder, RSI is considered overbought when above 70 and oversold when below 30.
    Signals can also be generated by looking for divergences, failure swings and centerline crossovers.
    RSI can also be used to identify the general trend."""

    delta = ohlc["close"].diff()

    up, down = delta.copy(), delta.copy()
    up[up < 0] = 0
    down[down > 0] = 0

    _gain = up.ewm(com=(period - 1), min_periods=period).mean()
    _loss = down.abs().ewm(com=(period - 1), min_periods=period).mean()

    RS = _gain / _loss
    return pd.Series(100 - (100 / (1 + RS)), name="RSI")