import os
import sys
import json
import pandas as pd
from PyQt5 import QtCore, QtGui
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import pyqtSlot
import datetime
import tushare as ts  # 金融接口

from model.pandasModel import PandasModel
from view.Ui_mainWindow import Ui_MainWindow
from controller.sys_config import SysConfig
from controller.plate_config import PlateConfig
from controller.kline_plot import KLineWidget

from predict.SVM.paramPredict import ParamPredicter
from predict.ARMA.arma_trainer import ArmaTrainer


# 通过ts金融接口获取股票数据df
def getStockData():
    pro = ts.pro_api()
    data = pro.stock_basic(exchange='', list_status='L', fields='ts_code,symbol,name,market,cnspell')
    today = datetime.datetime.today().strftime('%Y%m%d')  # 今天日期
    df = pro.daily(trade_date=today)  # 今日数据
    # 自动往前获取数据
    for i in range(10):
        if not df.empty:
            break
        day = (datetime.date.today() + datetime.timedelta(-i)).strftime('%Y%m%d')
        df = pro.daily(trade_date=day)

    # 初始化表格数据
    stockData = pd.merge(data, df)[
        ['ts_code', 'symbol', 'name', 'pct_chg', 'close', 'change', 'open', 'high', 'low', 'vol', 'amount', 'market', 'cnspell']
    ].sort_index()
    stockData['cnspell'] = stockData['cnspell'].str.upper()  # 将拼音缩写变为大写
    return stockData


# 从文件中读取数据,如果文件不存在或是非最新数据,则会更新数据,并写入文件
def initData(filePath):
    if os.path.exists(filePath):  # 文件存在,则判断是否是最新数据
        mtime = datetime.datetime.fromtimestamp(os.path.getmtime(filePath)).strftime("%Y%m%d")  # 文件上次修改日期
        today = datetime.datetime.today().strftime('%Y%m%d')  # 今天日期
        if mtime == today:  # 是最新数据
            # 如果不设置dtype=str,会将数字字符串将数字读,自动删除前面的0, 比如'000001'变成1
            # 但是如果设置了又会有别的BUG,此处用另一个解决方案
            stockData = pd.read_csv(filePath)
            # 重新单独读取symbol那一列,并指定dtype=str
            stockData['symbol'] = pd.read_csv(filePath, usecols=[1], dtype=str)
        else:  # 非最新数据,需要更新数据,并写入文件
            stockData = getStockData()
            stockData.to_csv(filePath, encoding="utf_8_sig", index=False)
    else:  # 文件不存在或今天还没更新,则生成数据,并写入文件
        stockData = getStockData()
        stockData.to_csv(filePath, encoding="utf_8_sig", index=False)
    return stockData


# 通过股票代码获取公司名称
def getName(stoke_code):
    data = ts.pro_api().query('stock_basic', fields='symbol,name')
    company_name = list(data.loc[data['symbol'] == stoke_code].name)[0]
    return company_name


# 通过股票代码获取ts代码
def getTsCode(stoke_code):
    data = ts.pro_api().query('stock_basic', fields='ts_code,symbol,name')
    # print(data)
    ts_code = list(data.loc[data['symbol'] == stoke_code].ts_code)[0]
    return ts_code


class MainWindow(QMainWindow, Ui_MainWindow):  # 多重继承QMainWindow和Ui_MainWindow
    def __init__(self):
        super(MainWindow, self).__init__()  # 先调用父类QMainWindow的初始化方法
        self.setupUi(self)  # 再调用setupUi方法

        self.setMinimumSize(1250, 600)

        # 设置图标
        icon = QtGui.QIcon()
        icon.addPixmap(QtGui.QPixmap("resources/stock_icon.ico"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
        self.setWindowIcon(icon)

        # palette = QtGui.QPalette()
        # palette.setColor(self.backgroundRole(), QColor(192, 253, 123))   # 设置背景颜色
        # self.setPalette(palette)

        # 表格样式
        self.stockView.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)  # 列宽自适应充满表格
        font = QFont('微软雅黑', 11)
        font.setBold(True)
        self.stockView.horizontalHeader().setFont(font)  # 设置表头字体
        self.stockView.setFont(QFont('微软雅黑', 8))  # 设置表身字体
        self.stockView.setEditTriggers(QAbstractItemView.NoEditTriggers)  # 设置表格不可更改
        self.stockView.setSelectionBehavior(QAbstractItemView.SelectRows)  # 设置不可选择单个单元格,只可选择一行
        self.stockView.setSelectionMode(QAbstractItemView.SingleSelection)  # 设置为只能选择单行,不多同时选择多行
        self.stockView.verticalHeader().setDefaultSectionSize(25)  # 设置默认行高
        # self.stockView.horizontalHeader().setDefaultSectionSize(200)  # 设置默认列宽
        self.stockView.setShowGrid(True)  # 设置有网格

        self.stockData = initData("data/stockData.csv")  # 初始化数据
        # 初始化常规板块数据
        self.plate1 = self.stockData[self.stockData['market'] == "主板"]
        self.plate2 = self.stockData[self.stockData['market'] == "中小板"]
        self.plate3 = self.stockData[self.stockData['market'] == "创业板"]
        self.plate4 = self.stockData[self.stockData['market'] == "科创板"]
        self.plate5 = pd.DataFrame()  # 自选股的初始化比较复杂,点击时候再写入数据

        # 表格初始化显示股票列表
        self.headers = ['TS代码', '股市代码', '名称', '涨幅%', '收盘', '涨跌', '开盘', '最高价', '最低价', '成交量(手)', '成交额', '市场', '拼音缩写']
        self.model = PandasModel(self.plate1, self.headers)  # 默认展示主板界面
        self.stockView.setModel(self.model)  # 关联QTableView控件和Model
        self.setBtnStyle(self.plate1Btn)
        self.stockView.setColumnHidden(self.stockData.columns.get_loc('ts_code'), True)  # 隐藏 ts_code 列
        self.stockView.setColumnHidden(self.stockData.columns.get_loc('market'), True)  # 隐藏 market 列
        self.stockView.setColumnHidden(self.stockData.columns.get_loc('cnspell'), True)  # 隐藏 cnspell 列

        # 设置表格双击事件
        self.stockView.doubleClicked.connect(self.stockView_double_clicked)

        # 隐藏TabWidget的标签标题
        self.stockWidget.tabBar().hide()  # 隐藏stockWidget的标签标题
        self.trainWidget.tabBar().hide()  # 隐藏trainWidget的标签标题
        self.trainShowWidget.tabBar().hide()
        # 为【股票】的【技术分析】创建一个新的Tab页
        # self.dataGetter = GetData()  # 用来获取数据的对象
        self.wg = KLineWidget(self)
        self.stockWidget.addTab(self.wg, 'K线图')  # 将self.wg作为新的Tab页追加到stockWidget中
        # 为【训练】的【获取数据】创建一个新的Tab页
        self.wg2 = KLineWidget(self)
        self.wg2.btnReture.setVisible(False)  # 【训练】无返回按钮
        self.trainShowWidget.addTab(self.wg2, 'K线图')

        # 【训练】界面的参数
        self.paramPredicter = None
        self.ppData = None
        self.test_size = 0.05

        self.armaTrainer = None

        self.trainDayNum = 0

        self.welcomeLabel.setStyleSheet("background-color:black;color:white")  # "欢迎使用训练功能。" 的样式

        # 默认进入【股票】页面
        self.mainWidget.setCurrentIndex(0)

    # 【菜单栏】的事件 ######################################################
    @pyqtSlot()
    def on_actSysConfig_triggered(self):
        QMessageBox.information(self, '提示信息', '系统配置功能。')

    @pyqtSlot()
    def on_actSysConfig_triggered(self):
        w = SysConfig()
        w.exec_()

    @pyqtSlot()
    def on_actPlateConfig_triggered(self):
        w = PlateConfig()
        w.exec_()

    @pyqtSlot()
    def on_actExit_triggered(self):
        self.close()

    # 【股票】标签页面的事件 ######################################################
    @pyqtSlot()
    def on_btnTech_clicked(self):  # 技术分析
        index = self.stockView.currentIndex()  # 当前选中的行的索引
        table_row = index.row()  # 根据索引获取行序号
        if table_row == -1:
            return
        stock_code = index.model().index(table_row, 1).data()  # 股票代码
        stock_name = index.model().index(table_row, 2).data()  # 股票名
        self.wg.stock_code, self.wg.stock_name = stock_code, stock_name  # 给K线图对象设置属性
        self.wg.stockLabel.setText("[%s - %s]" % (stock_code, stock_name))  # 设置K线图界面显示的股票信息
        self.wg.setCrosshairInfo(stock_name, 'D')  # 默认展示每日行情数据,所以为D
        self.wg.loadData(self.wg.dataGetter.getData(stock_code, 'D'))
        self.wg.refreshAll()
        self.stockWidget.setCurrentIndex(1)  # 切换标签

    # 表格行的双击事件
    def stockView_double_clicked(self, index):
        table_row = index.row()  # 点击的表格行
        if table_row == -1:
            return
        stock_code = index.model().index(table_row, 1).data()  # 股票代码
        stock_name = index.model().index(table_row, 2).data()  # 股票名
        self.wg.stock_code, self.wg.stock_name = stock_code, stock_name
        self.wg.stockLabel.setText("[%s - %s]" % (stock_code, stock_name))  # 设置K线图界面显示的股票信息
        self.wg.setCrosshairInfo(stock_name, 'D')  # 默认展示每日行情数据,所以为D
        self.wg.loadData(self.wg.dataGetter.getData(stock_code, 'D'))
        self.wg.refreshAll()
        self.stockWidget.setCurrentIndex(1)  # 切换标签

    # 搜索功能
    def on_search(self, str):
        searchText = self.searchInput.text()
        if self.find(searchText, self.stockData.columns.get_loc('ts_code')):
            pass
        elif self.find(searchText, self.stockData.columns.get_loc('name')):
            pass
        elif self.find(searchText, self.stockData.columns.get_loc('cnspell')):
            pass

    # 在column列搜索text文本
    def find(self, text, column=0):
        model = self.stockView.model()
        start = model.index(0, column)
        matches = model.match(
            start, QtCore.Qt.DisplayRole,
            text, 1, QtCore.Qt.MatchContains  # 搜索模式为包含
        )
        if matches:
            index = matches[0]  # index.row(), index.column()
            self.stockView.selectRow(index.row())
            return True
        return False

    # 设置选中某个板块时,按钮的样式
    def setBtnStyle(self, plateBtn):
        self.plate1Btn.setEnabled(True)
        self.plate2Btn.setEnabled(True)
        self.plate3Btn.setEnabled(True)
        self.plate4Btn.setEnabled(True)
        self.plate5Btn.setEnabled(True)
        self.plate6Btn.setEnabled(True)
        plateBtn.setEnabled(False)

    @pyqtSlot()
    def on_plate1Btn_clicked(self):  # 主板
        self.model = PandasModel(self.plate1, self.headers)
        self.stockView.setModel(self.model)
        self.setBtnStyle(self.plate1Btn)

    @pyqtSlot()
    def on_plate2Btn_clicked(self):  # 中小板
        self.model = PandasModel(self.plate2, self.headers)
        self.stockView.setModel(self.model)
        self.setBtnStyle(self.plate2Btn)

    @pyqtSlot()
    def on_plate3Btn_clicked(self):  # 创业板
        self.model = PandasModel(self.plate3, self.headers)
        self.stockView.setModel(self.model)
        self.setBtnStyle(self.plate3Btn)

    @pyqtSlot()
    def on_plate4Btn_clicked(self):  # 科创板
        self.model = PandasModel(self.plate4, self.headers)
        self.stockView.setModel(self.model)
        self.setBtnStyle(self.plate4Btn)

    @pyqtSlot()
    def on_plate5Btn_clicked(self):
        self.plate5.drop(self.plate5.index, inplace=True)  # 清空数据
        # 初始化自选股板块
        jsonFilePath = 'data/plate.json'
        with open(jsonFilePath, 'r', encoding='utf-8') as f:
            json_data = json.loads(f.read())
        for item in json_data[0]['stocks']:
            self.plate5 = self.plate5.append(self.stockData[self.stockData['symbol'] == item[:6]])
        self.model = PandasModel(self.plate5, self.headers)
        self.stockView.setModel(self.model)
        self.setBtnStyle(self.plate5Btn)

    # 【训练】标签页面的事件 ######################################################
    @pyqtSlot()
    def on_dataButton_clicked(self):  # 训练页面, 获取数据
        stock_code = self.stockCodeInput.text()  # 股票代码
        try:
            if stock_code != None:
                stock_name = getName(stock_code)
                self.wg2.stockLabel.setText("[%s - %s]" % (stock_code, stock_name))  # 设置K线图界面显示的股票信息
                self.wg2.stock_code, self.wg2.stock_name = stock_code, stock_name
                self.wg2.setCrosshairInfo(stock_name, 'D')  # 默认展示每日行情数据,所以为D
                self.wg2.loadData(self.wg2.dataGetter.getData(stock_code, 'D'))
                self.wg2.refreshAll()
                self.trainShowWidget.setCurrentIndex(1)  # 切换标签

                self.trainConfig.setEnabled(True)
        except:
            QMessageBox.information(self, '提示信息', '获取股票信息失败。')

    # 开始训练按钮,创建训练对象,训练相关面板可用
    @pyqtSlot()
    def on_beginTrainBtn_clicked(self):
        self.beginTrainBtn.setEnabled(False)  # 开始训练不可用
        self.endTrainBtn.setEnabled(True)  # 停止训练可用
        # 训练配置面板不可用
        self.startDateEdit.setEnabled(False)
        self.endDateEdit.setEnabled(False)
        self.splitInput.setEnabled(False)
        # 其余面板可用,并设置初始值
        if self.splitInput.text() == "":
            self.splitInput.setText("0.05")
        self.trainInfo.setEnabled(True)  # 训练信息面板可用
        self.predictOption.setEnabled(True)  # 预测选项面板可用
        self.predictOpen.setEnabled(True)
        self.predictClose.setEnabled(True)
        self.predictHigh.setEnabled(True)
        self.predictLow.setEnabled(True)
        self.paramValue.setEnabled(True)  # 参数评估面板可用

        # 获取界面信息
        begin_date = self.startDateEdit.text().replace("-", "")
        # print(begin_date)
        end_date = self.endDateEdit.text().replace("-", "")
        # print(end_date)
        self.test_size = float(self.splitInput.text())

        # 创建训练对象, 生成训练数据
        self.paramPredicter = ParamPredicter()
        self.ppData = self.paramPredicter.initData(
            code=getTsCode(self.stockCodeInput.text()),  # 此处要处理一下
            start=begin_date,
            end=end_date)

    # 停止训练按钮
    @pyqtSlot()
    def on_endTrainBtn_clicked(self):
        self.beginTrainBtn.setEnabled(True)  # 开始训练可用
        self.endTrainBtn.setEnabled(False)  # 结束训练不可用
        # 训练配置面板可用
        self.startDateEdit.setEnabled(True)
        self.endDateEdit.setEnabled(True)
        self.splitInput.setEnabled(True)
        # 其余面板不可用
        self.trainInfo.setEnabled(False)
        self.predictOption.setEnabled(False)
        self.paramValue.setEnabled(False)
        self.splitInput.setText("")

        begin_date = self.startDateEdit.text().replace("-", "")
        end_date = self.endDateEdit.text().replace("-", "")

        self.trainDayNum = 0
        self.trainDay.setText(str(self.trainDayNum))

    # 下一K线按钮
    @pyqtSlot()
    def on_nextKLineBtn_clicked(self):
        self.armaTrainer = ArmaTrainer()
        code = getTsCode(self.stockCodeInput.text())
        data = self.armaTrainer.initData(code, start="20180901")
        nextData = self.armaTrainer.predictNext(data)
        print(nextData)
        self.wg2.loadData(self.wg2.dataGetter.getData(self.wg2.stock_code, 'D', data=nextData))
        self.wg2.refreshAll()

        self.trainDayNum = self.trainDayNum + 1
        self.trainDay.setText(str(self.trainDayNum))


    # 预测开盘价
    @pyqtSlot()
    def on_predictOpen_clicked(self):
        self.paramPredicter.predict(
            data=self.ppData,
            param="open",
            test_size=self.test_size)

    # 预测收盘价
    @pyqtSlot()
    def on_predictClose_clicked(self):
        self.paramPredicter.predict(
            data=self.ppData,
            param="close",
            test_size=self.test_size)

    # 预测最高价
    @pyqtSlot()
    def on_predictHigh_clicked(self):
        self.paramPredicter.predict(
            data=self.ppData,
            param="high",
            test_size=self.test_size)

    # 预测最低价
    @pyqtSlot()
    def on_predictLow_clicked(self):
        self.paramPredicter.predict(
            data=self.ppData,
            param="low",
            test_size=self.test_size)

    # 预测股票走势
    @pyqtSlot()
    def on_predictTrend_clicked(self):
        self.paramPredicter.predictTrend(data=self.ppData)

    # 公共方法 ######################################################
    def returnToList(self):
        self.stockWidget.setCurrentIndex(0)

    # 处理关闭时的提示  
    def closeEvent(self, event):
        reply = QMessageBox.question(self, '退出', "确定要退出吗?", QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if reply == QMessageBox.Yes:
            event.accept()
        else:
            event.ignore()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    # app.setStyle(QStyleFactory.create("Fusion"))
    ui = MainWindow()
    ui.show()
    sys.exit(app.exec_())