from datetime import datetime, timedelta

import pandas as pd
from pandas import DataFrame
from ta.momentum import rsi, awesome_oscillator
from tinkoff.invest import Client, RequestError, CandleInterval, HistoricCandle
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

account_token = 't.P9y0_FG94o0Qj1esRJRdEP6iZG1I6q0yYQLWiEWdYPPk7O_NiLboZctkgUOghra1vSd4-ual9gEmbm9jTTD1gw'
commission = 0.04/100
begin_time = '2021-01-01 00:00:00+00:00'


def AllCandles(figi_name, num_days):
    df = DataFrame([{
        'time': c.time,
        'volume': c.volume,
        'open': MoneyValueToFloat(c.open),
        'close': MoneyValueToFloat(c.close),
        'high': MoneyValueToFloat(c.high),
        'low': MoneyValueToFloat(c.low),
    } for c in client.get_all_candles(
        figi=figi_name,
        from_=datetime.utcnow() - timedelta(days=num_days),
        to=datetime.utcnow(),
        interval=CandleInterval.CANDLE_INTERVAL_5_MIN
    )])
    return df


def MoneyValueToFloat(MonVal):
    return MonVal.units + MonVal.nano * 1e-9


def candleToDF(candles: [HistoricCandle]):
    dfCandles: DataFrame = DataFrame([{
        'time': c.time,
        'volume': c.volume,
        'open': MoneyValueToFloat(c.open),
        'close': MoneyValueToFloat(c.close),
        'high': MoneyValueToFloat(c.high),
        'low': MoneyValueToFloat(c.low),
    } for c in candles])
    return dfCandles


def RSI_OSCILATOR(df: DataFrame):
    df['rsi'] = rsi(close=df['close'], window=14)
    df.plot(x='time', y='rsi')
    # plt.show()
    return df


def AO_OSCILATOR(df: DataFrame):
    df['ao'] = awesome_oscillator(high=df['high'], low=df['low'], window1=5, window2=25)
    fig, axs = plt.subplots(2, 1)
    ax1 = df.plot(x='time', y='close', ax=axs[0])
    ax2 = df.plot(x='time', y='ao', ax=axs[1])
    axs[1].axhline(0, linestyle='--', alpha=0.5)
    ax1.xaxis.set_major_locator(mdates.WeekdayLocator())
    ax2.xaxis.set_major_locator(mdates.WeekdayLocator())
    # for i in range(len(df) - 1):
    #     if (df['ao'][i] < 0) and (df['ao'][i + 1] > 0):
    #         a = df['time'][i]
    #         b = df['open'][i]
    #         c = df['close'][i]
    #         print(a, b, c, 'buy')
    #     if (df['ao'][i] > 0) and (df['ao'][i + 1] < 0):
    #         a = df['time'][i]
    #         b = df['open'][i]
    #         c = df['close'][i]
    #         print(a, b, c, 'sell')
    # plt.show()
    return df

def make_signal(DatFr):
    rsi_low = 30.0
    rsi_high = 70.0
    avg_buy = 0.0
    num_r = 0
    num_a = 0
    num = 0
    balance = 10000.0
    RSI = False
    AO = False
    to_sold = 0
    print('start', balance)
    for i in range(len(DatFr) - 1):
        if (df['ao'][i] < 0) and (df['ao'][i + 1] > 0) and not (avg_buy > 0) and (balance - DatFr['close'][i+1] > 0) and str(DatFr['time'][i+1]) > begin_time:
            RSI = True
            if num_r == 0:
                num_r = 1
        elif (df['ao'][i] < 0) and (df['ao'][i + 1] > 0)  and (balance - DatFr['close'][i+1] > 0) and (avg_buy > DatFr['close'][i+1]):
            RSI = True
            if num_r == 0:
                num_r = 1
        elif (df['ao'][i] > 0) and (df['ao'][i + 1] < 0) and  0 < avg_buy < (1 + commission) * DatFr['close'][i+1]:
            RSI = False
            to_sold = 1
            if to_sold == 1:
                to_sold = 2

        if not (rsi_high > DatFr['rsi'][i] > rsi_low) and (balance - DatFr['close'][i + 1] > 0) and str(
            DatFr['time'][i + 1]) > begin_time:
            if (DatFr['rsi'][i] < rsi_low) and (DatFr['rsi'][i + 1] >= rsi_low) and not (avg_buy > 0.0):
                AO = True
                if num_a == 0:
                    num_a = 1
            elif (DatFr['rsi'][i] < rsi_low) and (DatFr['rsi'][i + 1] >= rsi_low) and (DatFr['close'][i] < avg_buy) and (
                    avg_buy > 0.0):
                AO = True
                if num_a == 0:
                    num_a = 1
            elif (DatFr['rsi'][i] > rsi_high) and (DatFr['rsi'][i + 1] <= rsi_high) and (
                    DatFr['close'][i + 1] > (1 + 2.3 * commission) * avg_buy) and (avg_buy > 0.0):
                AO = False
                to_sold = 1
                if to_sold == 1:
                    to_sold = 2

    if num_r > 0:
        num_r = num_r + 1
    if num_a > 0:
        num_a = num_a + 1

    if num_r > 8:
        num_r = 0
        RSI = False
    if num_a > 8:
        num_a = 0
        AO = False

    if RSI == True and AO == True:
        if not (avg_buy > 0):
            balance = balance - (1 + commission) * DatFr['close'][i + 1]
            avg_buy = DatFr['close'][i + 1]
            num = num + 1
            print('bought', balance, DatFr['close'][i + 1])
        else:
            balance = balance - (1 + commission) * DatFr['close'][i + 1]
            avg_buy = (num * avg_buy + DatFr['close'][i + 1]) / (num + 1)
            num = num + 1
            print('rebought', balance, DatFr['close'][i + 1])

    if to_sold == 2:
        balance = balance + num * (1 - commission) * DatFr['close'][i + 1]
        avg_buy = 0.0
        num = 0
        print('sold', balance, DatFr['close'][i + 1])


def balance_ao(DatFr):
    avg_buy = 0.0
    num = 0
    balance = 10000.0
    print('start', balance)
    for i in range(len(DatFr) - 1):
        if (df['ao'][i] < 0) and (df['ao'][i + 1] > 0) and not (avg_buy > 0) and (balance - DatFr['close'][i+1] > 0) and str(DatFr['time'][i+1]) > begin_time:
            balance = balance - (1 + commission) * DatFr['close'][i+1]
            avg_buy = DatFr['close'][i+1]
            num = num + 1
            print('bought', balance, DatFr['close'][i+1])
        elif (df['ao'][i] < 0) and (df['ao'][i + 1] > 0)  and (balance - DatFr['close'][i+1] > 0) and (avg_buy > DatFr['close'][i+1]):
            balance = balance - (1 + commission) * DatFr['close'][i + 1]
            avg_buy = (num * avg_buy + DatFr['close'][i + 1]) / (num + 1)
            num = num + 1
            print('rebought', balance, DatFr['close'][i + 1])
        elif (df['ao'][i] > 0) and (df['ao'][i + 1] < 0) and  0 < avg_buy < (1 + commission) * DatFr['close'][i+1]:
            balance = balance + num * (1 - commission) * DatFr['close'][i+1]
            avg_buy = 0.0
            num = 0
            print('sold', balance, DatFr['close'][i+1])


    print(balance, avg_buy)


def balance_rsi(DatFr):
    rsi_low = 30.0
    rsi_high = 70.0
    avg_buy = 0.0
    num = 0
    balance = 10000.0
    print('start', balance)
    for i in range(len(DatFr)-1):
        if not(rsi_high > DatFr['rsi'][i] > rsi_low) and (balance - DatFr['close'][i+1] > 0) and str(DatFr['time'][i+1]) > begin_time:
            # print(DatFr['time'][i+1])
            if(DatFr['rsi'][i] < rsi_low) and (DatFr['rsi'][i+1] >= rsi_low) and not (avg_buy > 0.0):
                balance = balance - (1 + commission) * DatFr['close'][i+1]
                num = 1
                avg_buy = DatFr['close'][i+1]
                print('bought', balance, DatFr['close'][i+1])
            elif (DatFr['rsi'][i] < rsi_low) and (DatFr['rsi'][i+1] >= rsi_low) and (DatFr['close'][i] < avg_buy) and (avg_buy > 0.0) :
                balance = balance - (1 + commission) * DatFr['close'][i+1]
                avg_buy = (num * avg_buy + DatFr['close'][i+1]) / (num + 1)
                num = num + 1
                print('rebought', balance, DatFr['close'][i+1])
            elif (DatFr['rsi'][i] > rsi_high) and (DatFr['rsi'][i + 1] <= rsi_high) and (
                    DatFr['close'][i + 1] > (1 + 2.3 * commission) * avg_buy) and (avg_buy > 0.0):
                balance = balance + num * (1 - commission) * DatFr['close'][i + 1]
                num = 0
                avg_buy = 0
                print('sold', balance, DatFr['close'][i + 1])
            # elif (avg_buy > 0) and ((1 - DatFr['close'][i + 1] / avg_buy) * 100 > 8):
            #     balance = balance - (1 + commission) * DatFr['close'][i + 1]
            #     num = num + 1
            #     avg_buy = (avg_buy + DatFr['close'][i + 1]) / 2
            #     print('rerebought', balance, DatFr['close'][i])
            # elif (avg_buy > 0) and ((1 - avg_buy/DatFr['close'][i + 1]) * 100 > 10):
            #     balance = balance + num * (1 - commission) * DatFr['close'][i + 1]
            #     num = 0
            #     avg_buy = 0
            #     print('resold', balance, DatFr['close'][i+1])
    print(balance)


if __name__ == '__main__':
    yandex = 'BBG006L8G4H1'
    gazprom = 'BBG004730RP0'
    sgng = 'BBG0047315D0'
    rosneft = 'BBG004731354'
    pik = 'BBG004S68BH6'
    with Client(account_token) as client:
        df: DataFrame = AllCandles(pik, 90)
        RSI_OSCILATOR(df)
        AO_OSCILATOR(df)
        # balance_rsi(df)
        # print('rsi')
        # balance_ao(df)
        # print('ao')
        make_signal(df)
    # print(df)
    #
    # r = client.users.get_accounts()
    # account_id = r.accounts[0].id
    # r: PortfolioResponse = client.operations.get_portfolio(account_id=account_id)
    # # keys = ['total_amount_bonds', 'total_amount_etf', 'total_amount_currencies']
    #
    # # print(pd.DataFrame([MoneyValueToFloat(getattr(r, k)) for k in keys], index=keys))
    # # print(r)
    #
    # data = pd.DataFrame([{
    #     'figi': p.figi,
    #     # 'instrument_type': p.instrument_type,
    #     'num': MoneyValueToFloat(p.quantity),
    #     'expected_yield': MoneyValueToFloat(p.expected_yield),
    #     'average_buy_price': MoneyValueToFloat(p.average_position_price_fifo),
    #     'current_price': MoneyValueToFloat(p.current_price)
    #     # 'currency': p.average_position_price.currency,
    #     # 'nkd': MoneyValueToFloat(p.current_nkd)
    #
    # } for p in r.positions])
    # print(data)