import os
import csv
from math import nan
import pandas as pd
import plotly.graph_objects as go
from plotly import subplots

from hts.HTS import HTS
from hts.BuySellChecker import BuySellChecker
from stock.util.Stock2Vector import Stock2Vector

class LabelChecker (HTS):
    buySellChecker = None
    stock2Vector = None

    def __init__(self, RESOURCE_PATH):
        super().__init__(RESOURCE_PATH)
        self.buySellChecker = BuySellChecker()
        self.stock2Vector = Stock2Vector(RESOURCE_PATH)
        return

    def checkTransaction(self, data, duration=60):
        bsLine = {}
        size = len(data["close"])

        # Type=False, 시뮬레이션 적용
        bsLine['buy'] = [-1 for i in range(size)]
        bsLine['buy_weight'] = [-1 for i in range(size)]
        bsLine['sell'] = [-1 for i in range(size)]
        bsLine['sell_weight'] = [-1 for i in range(size)]

        # 매수 처리
        p_min_price = 999999999
        # - 10을 한 것은 3시 11분까지 매수를 하기 위함임
        for i in range(size-duration-10):
            min_price, min_price_c, = 9999999, -1
            for c in range(i, i+duration):
                if min(data["close"][c], data["close"][c]) < min_price:
                    min_price = min(data["close"][c], data["close"][c])
                    min_price_c = c

            # 매수 처리
            if min_price_c > 0:
                isValid = True
                for c in range(2, 10):
                    if bsLine['buy'][min_price_c - c] > 0:
                        isValid = False
                        break
                if isValid:
                    if min_price < p_min_price:
                        bsLine['buy'][min_price_c] = min_price
                        bsLine['buy_weight'][min_price_c] = 1
                    p_min_price = min_price

        # 매도 처리
        buy_i = [c for c in range(size) if bsLine['buy'][c] > 0]
        for i in range(len(buy_i)):
            if i < len(buy_i)-1:
                boundry = data["high"][buy_i[i]+1: buy_i[i+1]]
                if len(boundry) > 1:
                    max_price = max(boundry)
                    for c in range(buy_i[i], buy_i[i+1]):
                        if data["high"][c] == max_price:
                            isValid = False
                            for d in range(c - 1, 0, -1):
                                if bsLine['sell'][d] > 0 and bsLine['buy'][d] < 0:
                                    break
                                if bsLine['sell'][d] < 0 and bsLine['buy'][d] > 0:
                                    isValid = True
                                    break
                            if isValid:
                                bsLine['sell'][c] = max_price
                                bsLine['sell_weight'][c] = 1
            else:
                # - 5를 한 것은 3시 16분까지 매도를 하기 위함임
                boundry = data["high"][buy_i[i] + 1: len(data["high"])-5]
                if len(boundry) > 1:
                    max_price = max(boundry)
                    for c in range(buy_i[i], len(data["high"])):
                        if data["high"][c] == max_price:
                            isValid = False
                            for d in range(c-1, 0, -1):
                                if bsLine['sell'][d] > 0 and bsLine['buy'][d] < 0 :
                                    break
                                if bsLine['sell'][d] < 0 and bsLine['buy'][d] > 0 :
                                    isValid = True
                                    break
                            if isValid:
                                bsLine['sell'][c] = max_price
                                bsLine['sell_weight'][c] = 1

        return bsLine, data

    def draw_simple(self, stock_code, given_day, data, bsLine):
        data = data.loc[pd.DatetimeIndex(data.index).day == int(given_day[6:])]
        buy_line = bsLine['buy'][381:]
        sell_line = bsLine['sell'][381:]

        buy_colors = []
        for i in range(len(buy_line)):
            if buy_line[i] < 0:
                buy_colors.append("#ffffff")
                buy_line[i] = nan
            else:
                buy_colors.append("#ff00ff")
        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['date'], y=buy_line, mode='markers', name="buy", marker=dict(size=14, color=buy_colors, line_width=0))
        sell_check = go.Scatter(x=data['date'], y=sell_line, mode='markers', name="sell", marker=dict(size=14, color=sell_colors, line_width=0))
        candle_stick = go.Candlestick(x=data['date'], open=data['open'], high=data['high'], low=data['low'], close=data['close'], increasing_line_color='red', decreasing_line_color='blue')
        candle_data = [candle_stick, buy_check, sell_check]

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

        self.fig = go.FigureWidget(data=candle_data)
        self.buy_data = self.fig.data[1]
        self.fig.update_layout(height=1000, title=stock_code + "_" + given_day + "_" + str(buy_count) + "," + str(sell_count))
        self.fig.show()

        return

    def draw_detail(self, stock_code, given_day, data, bsLine):
        buy_line = bsLine['buy']
        sell_line = bsLine['sell']

        # 그래프 설정을 위한 변수를 생성한다.
        data = data.astype({'open': 'int',
                            'high': 'int',
                            'low': 'int',
                            'close': 'int',
                            'volume': 'int',
                            'avg3': 'float',
                            'avg6': 'float',
                            'avg9': 'float',
                            'avg12': 'float',
                            'avg27': 'float',
                            'avg54': 'float',
                            'fast_k': 'float',
                            'slow_k': 'float',
                            'slow_d': '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("#D87D0F")
                buy_size.append(10 + (5 * 1))

        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['date'], y=buy_line, mode='markers', name="buy", marker=dict(size=buy_size, color=buy_colors, line_width=0))
        sell_check = go.Scatter(x=data['date'], y=sell_line, mode='markers', name="sell", marker=dict(size=14, color=sell_colors, line_width=0))
        upper = go.Scatter(x=data['date'], y=data["upper"], name="upper", line_color='#000000')
        lower = go.Scatter(x=data['date'], y=data["lower"], name="lower", line_color='#000000')
        avg3 = go.Scatter(x=data['date'], y=data["avg3"], name="avg3", line_color='#8F8203')
        avg6 = go.Scatter(x=data['date'], y=data["avg6"], name="avg6", line_color='#089B5B')
        avg9 = go.Scatter(x=data['date'], y=data["avg9"], name="avg9", line_color='#ff00ff')
        avg12 = go.Scatter(x=data['date'], y=data["avg12"], name="avg12", line_color='#1469F4')
        avg27 = go.Scatter(x=data['date'], y=data["avg27"], name="avg27", line_color='#000000')
        avg54 = go.Scatter(x=data['date'], y=data["avg54"], name="avg54", line_color='#008000')

        candle_stick = go.Candlestick(x=data['date'], open=data['open'], high=data['high'], low=data['low'], close=data['close'], increasing_line_color='red', decreasing_line_color='blue')
        volume_line = go.Scatter(x=data['date'], y=data["volume"], mode='lines', name='volume')
        # fast_k_line = go.Scatter(x=hts['date'], y=hts["fast_k"], mode='lines', name='fast_k')

        macd_line = go.Scatter(x=data['date'], y=data["macd"], mode='lines', name='macd')
        macd_s_line = go.Scatter(x=data['date'], y=data["macds"], mode='lines', name='macds')
        macd_o_line = go.Scatter(x=data['date'], y=data["macdo"], mode='lines', name='macdo')

        slow_k_line = go.Scatter(x=data['date'], y=data["slow_k"], mode='lines', name='slow_k')
        slow_d_line = go.Scatter(x=data['date'], y=data["slow_d"], mode='lines', name='slow_d')

        rsi_line = go.Scatter(x=data['date'], y=data["rsi"], mode='lines', name='rsi')
        rsis_line = go.Scatter(x=data['date'], y=data["rsis"], mode='lines', name='rsis')

        candle_data = [candle_stick, upper, lower, avg3, avg6, avg9, avg12, avg27, avg54, buy_check, sell_check]
        volume_data = [volume_line]
        macd_data = [macd_line, macd_s_line, macd_o_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=5, cols=1, subplot_titles=('캔들', "거래량", "MACD", "스토캐스틱", "RSI"))
        for trace in candle_data:
            fig.append_trace(trace, 1, 1)
        for trace in volume_data:
            fig.append_trace(trace, 2, 1)
        for trace in macd_data:
            fig.append_trace(trace, 3, 1)
        for trace in stochastic_data:
            fig.append_trace(trace, 4, 1)
        for trace in rsi_data:
            fig.append_trace(trace, 5, 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=5000, title=stock_code + "_" + given_day + "_" + str(buy_count) + "," + str(sell_count))
        fig.show()

        return

    def writeLabelFile(self, stock_code, bsLine, data, ymd):
        if not os.path.isdir(os.path.join(self.RESOURCE_PATH, "tmp")):
            os.mkdir(os.path.join(self.RESOURCE_PATH, "tmp"))

        outFileName = os.path.join(self.RESOURCE_PATH, "tmp", stock_code+"."+ymd+".sell.csv")
        with open(outFileName, "w", encoding="utf-8") as outFp:
            writer = csv.writer(outFp)
            for i, price in enumerate(bsLine["sell"]):
                if price != -1:
                    writer.writerow([data['date'][i], bsLine["sell"][i]])

        outFileName = os.path.join(self.RESOURCE_PATH, "tmp", stock_code+"."+ymd + ".buy.csv")
        with open(outFileName, "w", encoding="utf-8") as outFp:
            writer = csv.writer(outFp)
            for i, price in enumerate(bsLine["buy"]):
                if price != -1:
                    writer.writerow([data['date'][i], bsLine["buy"][i]])
        return

    def makeCandidate(self, stock_code, ymd, view=False):
        result = {"check": set(), "time": [], "open": [], "close": [], "high": [], "low": [], "vol": [], "label": []}
        self.getDBData(stock_code, ymd, result)

        data = self.buySellChecker.analyze(result)
        bsLine, data = self.checkTransaction(data)

        if view:
            self.draw_simple(stock_code, ymd, data, bsLine)
            #self.draw_detail(stock_code, ymd, data, bsLine)

        return bsLine, data
    def showLabels(self, stock_code, ymd, outFp=None):
        LAST_DATA = self.stock2Vector.getLastData(stock_code, ymd)
        result = self.stock2Vector.getRealTime(stock_code, ymd, LAST_DATA)

        # 이동평균, RSI, MACD, 일목균형, 볼린저밴드 상/하단을 계산한다.
        data = self.buySellChecker.analyze(result)

        # 4일치 중에서 앞에 3일은 제거한다.
        date = data['date'].dt.date.unique().tolist()
        data = data[data['date'].dt.date != date[0]]
        data = data[data['date'].dt.date != date[1]]
        data = data[data['date'].dt.date != date[2]]

        bsLine = {}
        lows = data["low"].tolist()
        closes = data["close"].tolist()
        labels = data["label"].tolist()
        size = len(lows)

        # Type=False, 시뮬레이션 적용
        bsLine['buy'] = [-1 for i in range(size)]
        bsLine['sell'] = [-1 for i in range(size)]
        for i in range(size):
            if outFp is not None:
                self.buySellChecker.write(outFp, data, i)
                outFp.flush()
            if labels[i] == 1:
                bsLine['sell'][i] = closes[i]
            if labels[i] == 2:
                bsLine['buy'][i] = lows[i]

        if outFp is None:
            #self.draw_simple(stock_code, ymd, data, bsLine)
            self.draw_detail(stock_code, ymd, data, bsLine)
        return

    def getDate(self, stock_code):
        ymds = self.getYMD(stock_code)
        return ymds

if __name__ == "__main__":
    PROJECT_HOME = os.path.join(os.path.dirname(os.path.join(os.path.dirname(os.path.join(os.path.dirname(__file__))))))
    RESOURCE_PATH = os.path.join(PROJECT_HOME, "resources")
    db_filename = os.path.join(RESOURCE_PATH, "hts.db")

    labelChecker = LabelChecker(RESOURCE_PATH)

    stock_codes = {
        "252670": ['20220701', '20220704', '20220705', '20220706', '20220707',
               '20220708', '20220711', '20220712', '20220713', '20220714',
               '20220715', '20220718', '20220719', '20220720', '20220721',
               '20220722', '20220725', '20220726', '20220727', '20220728',
               '20220729', '20220801', '20220802', '20220803', '20220804',
               '20220805', '20220808', '20220809', '20220810', '20220811',
               '20220812', '20220816'],
        #"122630": ['20220801', '20220802', '20220803', '20220804', '20220805', '20220808', '20220809', '20220810', '20220811', '20220812'],
    }

    for stock_code in stock_codes:
        for ymd in stock_codes[stock_code]:
            labelChecker.showLabels(stock_code, ymd)
