import numpy as np
from math import nan
import pandas as pd
import plotly.graph_objects as go
from plotly import subplots
import os
from datetime import datetime

from hts.HTS import HTS
from hts.DailyStatus import DailyStatus
from hts.BuySellChecker import BuySellChecker

class Simulation (HTS):
    buySellChecker = None
    stockPredictor = None
    dailyStatus = None

    def __init__(self, RESOURCE_PATH):
        super().__init__(RESOURCE_PATH)

        self.RESOURCE_PATH = RESOURCE_PATH

        self.buySellChecker = BuySellChecker()

        self.dailyStatus = DailyStatus(RESOURCE_PATH)
        return

    def draw(self, stock_code, given_day, data, bsLine):
        if bsLine is None:
            return

        # 어제 데이터는 지운다.
        buy_line = bsLine['buy']
        buy_weight_line = bsLine['buy_weight']
        sell_line = bsLine['sell']

        # 그래프 설정을 위한 변수를 생성한다.
        data = data.astype(
            {
                'open': 'int',
                'high': 'int',
                'low': 'int',
                'close': 'int',
                'slow_k': 'float',
                'slow_d': 'float',
                'macd': 'float',
                'macds': 'float',
                'envelope_upper': 'float',
                'envelope_lower': 'float',
                'envelope_middle': 'float',
                'rsi': 'float',
                'rsis': 'float'
            }
        )

        buy_size = []
        buy_colors = []
        for i in range(len(buy_line)):
            if buy_line[i] < 0:
                buy_colors.append("#ffffff")
                buy_line[i] = nan
                buy_size.append(0)
            else:
                buy_colors.append("#B2028C")
                buy_size.append(10 + (5 * buy_weight_line[i]))

        sell_colors = []
        for i in range(len(sell_line)):
            if sell_line[i] < 0:
                sell_colors.append("#ffffff")
                sell_line[i] = nan
            else:
                sell_colors.append("#00ced1")

        # 그래프를 설정한다.
        buy_check = go.Scatter(x=data['ymd'], y=buy_line, mode='markers', name="buy", marker=dict(size=buy_size, color=buy_colors, line_width=0))
        sell_check = go.Scatter(x=data['ymd'], y=sell_line, mode='markers', name="sell", marker=dict(size=14, color=sell_colors, line_width=0))
        envelope_upper = go.Scatter(x=data['ymd'], y=data["envelope_upper"], name="upper", line_color='#000000')
        envelope_middle = go.Scatter(x=data['ymd'], y=data["envelope_middle"], name="upper", line_color='#927786')
        envelope_lower = go.Scatter(x=data['ymd'], y=data["envelope_lower"], name="lower", line_color='#000000')

        candle_stick = go.Candlestick(x=data['ymd'], open=data['open'], high=data['high'], low=data['low'], close=data['close'], increasing_line_color='red', decreasing_line_color='blue', showlegend=False)

        macd_line = go.Scatter(x=data['ymd'], y=data["macd"], line=dict(color='red', width=2), name='macd')
        macd_s_line = go.Scatter(x=data['ymd'], y=data["macds"], line=dict(dash='dashdot', color='black', width=2), name='macds')

        # fast_k_line = go.Scatter(x=hts['date'], y=hts["fast_k"], mode='lines', name='fast_k')
        slow_k_line = go.Scatter(x=data['ymd'], y=data["slow_k"], line=dict(color='red', width=2), name='slow_k')
        slow_d_line = go.Scatter(x=data['ymd'], y=data["slow_d"], line=dict(dash='dashdot', color='black', width=2), name='slow_d')

        rsi_line = go.Scatter(x=data['ymd'], y=data["rsi"], line=dict(color='red', width=2), name='rsi')
        rsis_line = go.Scatter(x=data['ymd'], y=data["rsis"], line=dict(dash='dashdot', color='black', width=2), name='rsis')

        candle_data = [candle_stick, envelope_upper, envelope_middle, envelope_lower, buy_check, sell_check]
        macd_data = [macd_line, macd_s_line]
        stochastic_data = [slow_k_line, slow_d_line]
        rsi_data = [rsi_line, rsis_line]

        # 그래프를 그린다.
        """
        fig = go.Figure(data=candle_data)
        fig.update_layout(title=stock_code + "_" + given_day)
        fig.show()
        """

        fig = subplots.make_subplots(
            rows=4, cols=1,
            subplot_titles=("MACD", "RSI", "스토캐스틱", '캔들'),
            #specs=[[{}], [{}], [{}], [{}], [{}], [{}]],
            shared_xaxes=True, horizontal_spacing=0.03, vertical_spacing=0.01,
            row_heights=[200, 200, 200, 750]
        )
        for trace in macd_data:
            fig.append_trace(trace, 1, 1)
        for trace in rsi_data:
            fig.append_trace(trace, 2, 1)
        for trace in stochastic_data:
            fig.append_trace(trace, 3, 1)
        for trace in candle_data:
            fig.append_trace(trace, 4, 1)


        #fig.update_xaxes(nticks=5)
        #fig.update_layout(height=1800, title=stock_code + "_" + given_day, xaxis_rangeslider_visible=False)

        df = pd.DataFrame(bsLine)
        df = df.fillna(-1)
        buy_count = len(df.loc[df["buy"] > 0])
        sell_count = len(df.loc[df["sell"] > 0])

        fig.update_layout(height=1700, title=stock_code + "_" + given_day + "_" + str(buy_count)+","+str(sell_count))
        #fig.update_layout(title=stock_code + "_" + given_day + "_" + str(buy_count) + "," + str(sell_count))
        fig.show()

        return


    def simulate(self, stock_code, n=100):
        today = datetime.today().strftime('%Y%m%d')
        data = self.dailyStatus.getLastData(stock_code, today, n)

        # 사야 할 시점과 팔아야 할 시점을 체크한다.
        bsLine, data = self.buySellChecker.checkEnvelopeTiming(data, stock_code, isRealTime=False)

        # 그래프를 그린다.
        self.draw(stock_code, today, data, bsLine)

        return

if __name__ == "__main__":

    PROJECT_HOME = os.getcwd()
    RESOURCE_PATH = os.path.join(PROJECT_HOME, "resources")

    # to check bying
    stock_codes = ["252670", "122630"]
    #stock_codes = ["252670"]
    #stock_codes = ["122630"]

    for stock_code in stock_codes:
        simulation = Simulation(RESOURCE_PATH)
        simulation.simulate(stock_code, 2000)

    print ("done...")