from gc import collect
from pathlib import Path
from traceback import format_exc

import matplotlib.pyplot as plt
from PySide6.QtWidgets import QMessageBox, QMainWindow, QApplication
from matplotlib import gridspec
from matplotlib.backends.backend_qt import NavigationToolbar2QT
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg
from overrides import overrides
from pandas import read_csv, DataFrame
from yaml import dump, load, FullLoader

from func.utils.ConfigParams import Filename, Params
from func.utils.PublicFunc import PublicFunc
from func.utils.Constants import Constants
from func.utils.Result import Result
from func.utils.detect_Jpeak import preprocess, Jpeak_Detection

from ui.MainWindow.MainWindow_detect_Jpeak import Ui_MainWindow_detect_Jpeak
from ui.setting.detect_Jpeak_input_setting import Ui_MainWindow_detect_Jpeak_input_setting


Config = {

}

ButtonState = {
    "Default": {
        "pushButton_input_setting": True,
        "pushButton_input": True,
        "pushButton_view": False,
        "pushButton_save": False
    },
    "Current": {
        "pushButton_input_setting": True,
        "pushButton_input": True,
        "pushButton_view": False,
        "pushButton_save": False
    }
}


class SettingWindow(QMainWindow):

    def __init__(self, root_path, sampID):
        super(SettingWindow, self).__init__()
        self.ui = Ui_MainWindow_detect_Jpeak_input_setting()
        self.ui.setupUi(self)

        self.root_path = root_path
        self.sampID = sampID

        self.msgBox = QMessageBox()
        self.msgBox.setWindowTitle(Constants.MAINWINDOW_MSGBOX_TITLE)

        self.config = None
        self.__read_config__()
        self.__examine_freq__()

        self.ui.spinBox_input_freq.valueChanged.connect(self.__update_ui__)
        self.ui.pushButton_confirm.clicked.connect(self.__write_config__)
        self.ui.pushButton_cancel.clicked.connect(self.__rollback_config__)
        self.ui.pushButton_cancel.clicked.connect(self.close)

    def __read_config__(self):
        if not Path(Params.DETECT_JPEAK_CONFIG_FILE_PATH).exists():
            with open(Params.DETECT_JPEAK_CONFIG_FILE_PATH, "w") as f:
                dump(Params.DETECT_JPEAK_CONFIG_NEW_CONTENT, f)

        with open(Params.DETECT_JPEAK_CONFIG_FILE_PATH, "r") as f:
            file_config = load(f.read(), Loader=FullLoader)
            Config.update(file_config)
            self.config = file_config

        Config.update({
            "Path": {
                "Input": str((Path(self.root_path) / Filename.PATH_ORGBCG_TEXT /
                              Path(str(self.sampID)))),
                "Save": str((Path(self.root_path) / Filename.PATH_ORGBCG_TEXT /
                             Path(str(self.sampID))))
            }
        })

        # 数据回显
        self.ui.spinBox_input_freq.setValue(Config["InputConfig"]["Freq"])
        self.ui.plainTextEdit_file_path_input.setPlainText(Config["Path"]["Input"])
        self.ui.plainTextEdit_deepmodel_path.setPlainText(Config["ModelFolderPath"])
        self.ui.plainTextEdit_file_path_save.setPlainText(Config["Path"]["Save"])

    def __write_config__(self):
        # 从界面写入配置
        Config["InputConfig"]["Freq"] = self.ui.spinBox_input_freq.value()
        Config["Path"]["Input"] = self.ui.plainTextEdit_file_path_input.toPlainText()
        Config["Path"]["Save"] = self.ui.plainTextEdit_file_path_save.toPlainText()
        Config["ModelFolderPath"] = self.ui.plainTextEdit_deepmodel_path.toPlainText()

        # 保存配置到文件
        self.config["InputConfig"]["Freq"] = self.ui.spinBox_input_freq.value()
        self.config["ModelFolderPath"] = self.ui.plainTextEdit_deepmodel_path.toPlainText()

        with open(Params.DETECT_JPEAK_CONFIG_FILE_PATH, "w") as f:
            dump(self.config, f)

        self.close()

    def __rollback_config__(self):
        self.__read_config__()

    def __update_ui__(self):
        self.ui.plainTextEdit_file_path_input.setPlainText(
            str((Path(self.root_path) /
                 Filename.PATH_ORGBCG_TEXT /
                 Path(str(self.sampID)) /
                 Path(Filename.BCG_FILTER +
                      str(self.ui.spinBox_input_freq.value()) +
                      Params.ENDSWITH_TXT))))

    def __examine_freq__(self):
        if Path(Config["Path"]["Input"]).is_file():
            Config["Path"]["Input"] = str(Path(Config["Path"]["Input"]).parent)

        result = PublicFunc.examine_file(Config["Path"]["Input"], Filename.BCG_FILTER, Params.ENDSWITH_TXT)
        if result.status:
            Config["InputConfig"]["Freq"] = result.data["freq"]
        else:
            PublicFunc.msgbox_output(self, Filename.BCG_FILTER + Constants.FAILURE_REASON["Get_Freq_Not_Correct"], Constants.MSGBOX_TYPE_ERROR)

        # 数据回显
        self.ui.spinBox_input_freq.setValue(Config["InputConfig"]["Freq"])


class MainWindow_detect_Jpeak(QMainWindow):

    def __init__(self):
        super(MainWindow_detect_Jpeak, self).__init__()
        self.ui = Ui_MainWindow_detect_Jpeak()
        self.ui.setupUi(self)

        self.root_path = None
        self.sampID = None

        self.data = None
        self.model = None

        self.setting = None

        # 初始化进度条
        self.progressbar = None
        PublicFunc.add_progressbar(self)

        #初始化画框
        self.fig = None
        self.canvas = None
        self.figToolbar = None
        self.gs = None
        self.ax0 = None

        self.ui.textBrowser_info.setStyleSheet("QTextBrowser { background-color: rgb(255, 255, 200); }")
        PublicFunc.__styleAllButton__(self, ButtonState)

        self.msgBox = QMessageBox()
        self.msgBox.setWindowTitle(Constants.MAINWINDOW_MSGBOX_TITLE)

    @overrides
    def show(self, root_path, sampID):
        super().show()
        self.root_path = root_path
        self.sampID = sampID

        self.setting = SettingWindow(root_path, sampID)

        # 初始化画框
        self.fig = plt.figure(figsize=(12, 9), dpi=100)
        self.canvas = FigureCanvasQTAgg(self.fig)
        self.figToolbar = NavigationToolbar2QT(self.canvas)
        for action in self.figToolbar.actions():
            if action.text() == "Subplots" or action.text() == "Customize":
                self.figToolbar.removeAction(action)
        self.ui.verticalLayout_canvas.addWidget(self.canvas)
        self.ui.verticalLayout_canvas.addWidget(self.figToolbar)
        self.gs = gridspec.GridSpec(1, 1, height_ratios=[1])
        self.fig.subplots_adjust(top=0.98, bottom=0.05, right=0.98, left=0.1, hspace=0, wspace=0)
        self.ax0 = self.fig.add_subplot(self.gs[0])
        self.ax0.grid(True)
        self.ax0.xaxis.set_major_formatter(Params.FORMATTER)

        PublicFunc.__resetAllButton__(self, ButtonState)

        self.ui.doubleSpinBox_bandPassLow.setValue(Config["Filter"]["BandPassLow"])
        self.ui.doubleSpinBox_bandPassHigh.setValue(Config["Filter"]["BandPassHigh"])
        self.ui.spinBox_peaksValue.setValue(Config["PeaksValue"])
        self.ui.doubleSpinBox_ampValue.setValue(Config["AmpValue"])
        self.ui.spinBox_intervalLow.setValue(Config["IntervalLow"])
        self.ui.spinBox_intervalHigh.setValue(Config["IntervalHigh"])
        self.ui.checkBox_useCPU.setChecked(Config["UseCPU"])

        self.ui.pushButton_input.clicked.connect(self.__slot_btn_input__)
        self.ui.pushButton_input_setting.clicked.connect(self.setting.show)
        self.ui.pushButton_view.clicked.connect(self.__slot_btn_view__)
        self.ui.pushButton_save.clicked.connect(self.__slot_btn_save__)
        self.ui.doubleSpinBox_bandPassLow.editingFinished.connect(self.__update_config__)
        self.ui.doubleSpinBox_bandPassHigh.editingFinished.connect(self.__update_config__)
        self.ui.spinBox_peaksValue.editingFinished.connect(self.__update_config__)
        self.ui.doubleSpinBox_ampValue.editingFinished.connect(self.__update_config__)
        self.ui.spinBox_intervalLow.editingFinished.connect(self.__update_config__)
        self.ui.spinBox_intervalHigh.editingFinished.connect(self.__update_config__)
        self.ui.checkBox_useCPU.stateChanged.connect(self.__update_config__)
        self.ui.comboBox_model.currentTextChanged.connect(self.__update_config__)
        self.ui.radioButton_skip_process.toggled.connect(self.__update_config__)
        self.ui.radioButton_bandpass_process.toggled.connect(self.__update_config__)

    @overrides
    def closeEvent(self, event):
        reply = QMessageBox.question(self, '确认', '确认退出吗？', QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if reply == QMessageBox.Yes:
            PublicFunc.__disableAllButton__(self, ButtonState)

            PublicFunc.statusbar_show_msg(self, PublicFunc.format_status_msg(Constants.SHUTTING_DOWN))
            QApplication.processEvents()

            # 清空画框
            if self.ax0 is not None:
                self.ax0.clear()

            # 释放资源
            self.setting.close()
            del self.data
            del self.model
            self.fig.clf()
            plt.close(self.fig)
            self.deleteLater()
            collect()
            self.canvas = None
            event.accept()
        else:
            event.ignore()

    def __reset__(self):
        ButtonState["Current"].update(ButtonState["Default"].copy())

    def __plot__(self):
        # 清空画框
        self.reset_axes()

        sender = self.sender()

        if sender == self.ui.pushButton_view:
            self.ax0.plot(self.data.processed_data,
                                                color=Constants.PLOT_COLOR_BLUE,
                                                label=Constants.DETECT_JPEAK_PLOT_LABEL_BCG)
            self.ax0.plot(self.data.peak, self.data.processed_data[self.data.peak],
                                             'r.',
                                             label=Constants.DETECT_JPEAK_PLOT_LABEL_J_PEAKS)
            self.ax0.plot(self.data.interval,
                                                color=Constants.PLOT_COLOR_ORANGE,
                                                label=Constants.DETECT_JPEAK_PLOT_LABEL_INTERVAL)
            self.ax0.legend(loc=Constants.PLOT_UPPER_RIGHT)

            self.canvas.draw()
            return Result().success(info=Constants.DRAW_FINISHED)
        else:
            self.canvas.draw()
            return Result().failure(info=Constants.DRAW_FAILURE)

    def __update_config__(self):
        sender = self.sender()
        Config["Filter"]["BandPassLow"] = self.ui.doubleSpinBox_bandPassLow.value()
        Config["Filter"]["BandPassHigh"] = self.ui.doubleSpinBox_bandPassHigh.value()
        Config["PeaksValue"] = self.ui.spinBox_peaksValue.value()
        Config["AmpValue"] = self.ui.doubleSpinBox_ampValue.value()
        Config["IntervalLow"] = self.ui.spinBox_intervalLow.value()
        Config["IntervalHigh"] = self.ui.spinBox_intervalHigh.value()
        Config["UseCPU"] = self.ui.checkBox_useCPU.isChecked()
        Config["DetectMethod"] = self.ui.comboBox_model.currentText()
        if sender == self.ui.radioButton_skip_process:
            Config["Filter"]["Mode"] = "skip"
        elif sender == self.ui.radioButton_bandpass_process:
            Config["Filter"]["Mode"] = "bandpass"

    def __slot_btn_input__(self):
        PublicFunc.__disableAllButton__(self, ButtonState)

        # 清空画框
        self.reset_axes()
        self.canvas.draw()

        # 清空模型列表
        self.ui.comboBox_model.clear()

        self.data = Data()
        self.model = Model()

        # 寻找模型
        PublicFunc.progressbar_update(self, 1, 2, Constants.DETECT_JPEAK_LOADING_MODEL, 0)
        result = self.model.seek_model()
        if not result.status:
            PublicFunc.text_output(self.ui, "(1/2)" + result.info, Constants.TIPS_TYPE_ERROR)
            PublicFunc.msgbox_output(self, result.info, Constants.MSGBOX_TYPE_ERROR)
            PublicFunc.finish_operation(self, ButtonState)
            return
        else:
            PublicFunc.text_output(self.ui, "(1/2)" + result.info, Constants.TIPS_TYPE_INFO)
            self.update_ui_comboBox_model(self.model.model_list)

        # 导入数据
        PublicFunc.progressbar_update(self, 2, 2, Constants.INPUTTING_DATA, 10)
        result = self.data.open_file()
        if not result.status:
            PublicFunc.text_output(self.ui, "(2/2)" + result.info, Constants.TIPS_TYPE_ERROR)
            PublicFunc.msgbox_output(self, result.info, Constants.MSGBOX_TYPE_ERROR)
            PublicFunc.finish_operation(self, ButtonState)
            return
        else:
            PublicFunc.text_output(self.ui, "(2/2)" + result.info, Constants.TIPS_TYPE_INFO)

        self.setting.close()
        ButtonState["Current"]["pushButton_input_setting"] = False
        ButtonState["Current"]["pushButton_input"] = False
        ButtonState["Current"]["pushButton_view"] = True
        ButtonState["Current"]["pushButton_save"] = False
        PublicFunc.finish_operation(self, ButtonState)

    def __slot_btn_view__(self):
        PublicFunc.__disableAllButton__(self, ButtonState)

        # 数据预处理
        PublicFunc.progressbar_update(self, 1, 3, Constants.PREPROCESSING_DATA, 0)
        result = self.data.preprocess()
        if not result.status:
            PublicFunc.text_output(self.ui, "(1/3)" + result.info, Constants.TIPS_TYPE_ERROR)
            PublicFunc.msgbox_output(self, result.info, Constants.MSGBOX_TYPE_ERROR)
            PublicFunc.finish_operation(self, ButtonState)
            return
        else:
            PublicFunc.text_output(self.ui, "(1/3)" + result.info, Constants.TIPS_TYPE_INFO)

        # 预测峰值
        PublicFunc.progressbar_update(self, 2, 3, Constants.DETECT_JPEAK_PREDICTING_PEAK, 10)
        self.model.selected_model = Config["DetectMethod"]
        scale = self.ui.spinBox_scaleValue.value() if self.ui.checkBox_scaleEnable.isChecked() else 0
        result = self.data.predict_Jpeak(self.model, batch_size=int(self.ui.comboBox_batchSize.currentText()), scale=scale)
        if not result.status:
            PublicFunc.text_output(self.ui, "(2/3)" + result.info, Constants.TIPS_TYPE_ERROR)
            PublicFunc.msgbox_output(self, result.info, Constants.MSGBOX_TYPE_ERROR)
            PublicFunc.finish_operation(self, ButtonState)
            return
        else:
            PublicFunc.text_output(self.ui, "(2/3)" + result.info, Constants.TIPS_TYPE_INFO)
            PublicFunc.text_output(self.ui, Constants.DETECT_JPEAK_DATA_LENGTH_POINTS + str(len(self.data.raw_data)),
                                   Constants.TIPS_TYPE_INFO)
            PublicFunc.text_output(self.ui, Constants.DETECT_JPEAK_DURATION_MIN +
                                   str((len(self.data.raw_data) / Config["InputConfig"]["Freq"] / 60)),
                                   Constants.TIPS_TYPE_INFO)
            PublicFunc.text_output(self.ui, Constants.DETECT_JPEAK_PEAK_AMOUNT + str(len(self.data.peak)),
                                   Constants.TIPS_TYPE_INFO)

        # 绘图
        PublicFunc.progressbar_update(self, 3, 3, Constants.DRAWING_DATA, 70)
        result = self.__plot__()
        if not result.status:
            PublicFunc.text_output(self.ui, "(3/3)" + result.info, Constants.TIPS_TYPE_ERROR)
            PublicFunc.msgbox_output(self, result.info, Constants.MSGBOX_TYPE_ERROR)
            PublicFunc.finish_operation(self, ButtonState)
            return
        else:
            PublicFunc.text_output(self.ui, "(3/3)" + result.info, Constants.TIPS_TYPE_INFO)

        ButtonState["Current"]["pushButton_save"] = True
        PublicFunc.finish_operation(self, ButtonState)

    def __slot_btn_save__(self):
        reply = QMessageBox.question(self, Constants.QUESTION_TITLE,
                                     Constants.QUESTION_CONTENT + Config["Path"]["Save"],
                                     QMessageBox.Yes | QMessageBox.No,
                                     QMessageBox.Yes)
        if reply == QMessageBox.Yes:
            PublicFunc.__disableAllButton__(self, ButtonState)

            # 保存
            PublicFunc.progressbar_update(self, 1, 1, Constants.SAVING_DATA, 0)

            total_rows = len(DataFrame(self.data.peak.reshape(-1)))
            chunk_size = Params.DETECT_JPEAK_SAVE_CHUNK_SIZE
            try:
                with open(Config["Path"]["Save"], 'w') as f:
                    for start in range(0, total_rows, chunk_size):
                        end = min(start + chunk_size, total_rows)
                        chunk = DataFrame(self.data.peak.reshape(-1)).iloc[start:end]
                        result = self.data.save(chunk)
                        progress = int((end / total_rows) * 100)
                        self.progressbar.setValue(progress)
                        QApplication.processEvents()
            except FileNotFoundError as e:
                result = Result().failure(info=Constants.SAVE_FAILURE + Constants.FAILURE_REASON["Save_File_Not_Found"])

            if not result.status:
                PublicFunc.text_output(self.ui, "(1/1)" + result.info, Constants.TIPS_TYPE_ERROR)
                PublicFunc.msgbox_output(self, result.info, Constants.MSGBOX_TYPE_ERROR)
                PublicFunc.finish_operation(self, ButtonState)
                return
            else:
                PublicFunc.text_output(self.ui, "(1/1)" + result.info, Constants.TIPS_TYPE_INFO)

            PublicFunc.msgbox_output(self, result.info, Constants.TIPS_TYPE_INFO)
            PublicFunc.finish_operation(self, ButtonState)

    def reset_axes(self):
        self.ax0.clear()
        self.ax0.grid(True)
        self.ax0.xaxis.set_major_formatter(Params.FORMATTER)

    def update_ui_comboBox_model(self, model_list):
        self.ui.comboBox_model.clear()
        self.ui.comboBox_model.addItems(model_list)


class Data:

    def __init__(self):
        self.raw_data = None
        self.processed_data = None
        self.peak = None
        self.interval = None

    def open_file(self):
        if Path(Config["Path"]["Input"]).is_file():
            Config["Path"]["Input"] = str(Path(Config["Path"]["Input"]).parent)

        result = PublicFunc.examine_file(Config["Path"]["Input"], Filename.BCG_FILTER, Params.ENDSWITH_TXT)
        if result.status:
            Config["Path"]["Input"] = result.data["path"]
            Config["InputConfig"]["Freq"] = result.data["freq"]
            Config["Path"]["Save"] = str(
                Path(Config["Path"]["Save"]) / Path(Filename.JPEAK_REVISE + str(Config["InputConfig"]["Freq"]) + Params.ENDSWITH_TXT))
        else:
            return result

        try:
            self.raw_data = read_csv(Config["Path"]["Input"],
                                     encoding=Params.UTF8_ENCODING,
                                     header=None).to_numpy().reshape(-1)
        except Exception as e:
            return Result().failure(info=Constants.INPUT_FAILURE +
                                         Constants.FAILURE_REASON["Open_Data_Exception"] + "\n" + format_exc())

        return Result().success(info=Constants.INPUT_FINISHED)

    def preprocess(self):
        if self.raw_data is None:
            return Result().failure(info=Constants.PREPROCESS_FAILURE +
                                         Constants.FAILURE_REASON["Data_Not_Exist"])

        try:
            if Config["Filter"]["Mode"] == "skip":
                self.processed_data = self.raw_data[:len(self.raw_data) // (Config["InputConfig"]["Freq"] * 10) * Config["InputConfig"]["Freq"] * 10]
                return Result().success(info=Constants.PREPROCESS_NO_NEED)
            elif Config["Filter"]["Mode"] == "bandpass":
                if ((Config["Filter"]["BandPassLow"] >= Config["Filter"]["BandPassHigh"]) or
                        (Config["Filter"]["BandPassLow"] <= 0) or (Config["Filter"]["BandPassHigh"] <= 0)):
                    return Result().failure(
                        info=Constants.PREPROCESS_FAILURE + Constants.FAILURE_REASON["Filter_Args_Not_Correct"])
                self.processed_data = preprocess(self.raw_data,
                                                 Config["InputConfig"]["Freq"],
                                                 Config["Filter"]["BandPassLow"],
                                                 Config["Filter"]["BandPassHigh"],
                                                 Config["AmpValue"])
                return Result().success(info=Constants.PREPROCESS_FINISHED)
        except Exception as e:
            return Result().failure(info=Constants.PREPROCESS_FAILURE +
                                         Constants.FAILURE_REASON["Preprocess_Exception"] + "\n" + format_exc())

    def predict_Jpeak(self, model, batch_size=0, scale=0):
        if not (Path(model.model_folder_path) / Path(model.selected_model)).exists():
            return Result().failure(info=Constants.DETECT_JPEAK_PREDICT_FAILURE +
                                         Constants.FAILURE_REASON["Model_File_Not_Exist"])

        if self.processed_data is None:
            return Result().failure(info=Constants.DETECT_JPEAK_PREDICT_FAILURE +
                                         Constants.FAILURE_REASON["Data_Not_Exist"])

        try:
            if Config["IntervalLow"] >= Config["IntervalHigh"]:
                return Result().failure(
                    info=Constants.DETECT_JPEAK_PREDICT_FAILURE + Constants.FAILURE_REASON["Interval_Args_Not_Correct"])
            self.peak, self.interval = Jpeak_Detection(model.selected_model,
                                                       Path(model.model_folder_path) / Path(model.selected_model),
                                                       self.processed_data,
                                                       Config["InputConfig"]["Freq"],
                                                       Config["IntervalHigh"],
                                                       Config["IntervalLow"],
                                                       Config["PeaksValue"],
                                                       Config["UseCPU"],
                                                       batch_size,
                                                       scale)
        except Exception as e:
            return Result().failure(info=Constants.DETECT_JPEAK_PREDICT_FAILURE +
                                         Constants.FAILURE_REASON["Predict_Exception"] + "\n" + format_exc())

        return Result().success(info=Constants.DETECT_JPEAK_PREDICT_FINISHED)

    def save(self, chunk):
        if self.peak is None:
            return Result().failure(info=Constants.SAVE_FAILURE + Constants.FAILURE_REASON["Data_Not_Exist"])

        try:
            chunk.to_csv(Config["Path"]["Save"], mode='a', index=False, header=False)
        except PermissionError as e:
            return Result().failure(info=Constants.SAVE_FAILURE + Constants.FAILURE_REASON["Save_Permission_Denied"])
        except Exception as e:
            return Result().failure(info=Constants.SAVE_FAILURE +
                                         Constants.FAILURE_REASON["Save_Exception"] + "\n" + format_exc())

        return Result().success(info=Constants.SAVE_FINISHED)


class Model:

    def __init__(self):
        self.model_folder_path = Config["ModelFolderPath"]
        self.model_list = None
        self.selected_model_path = None
        self.selected_model = None

    def seek_model(self):
        if not Path(Config["ModelFolderPath"]).exists():
            return Result().failure(info=Constants.DETECT_JPEAK_LOAD_FAILURE + Constants.FAILURE_REASON["Model_Path_Not_Exist"])

        try:
            self.model_list = [file.name for file in Path(Config["ModelFolderPath"]).iterdir() if file.is_file()]
            if len(self.model_list) == 0:
                return Result().failure(info=Constants.FAILURE_REASON["Model_File_Not_Exist"])
        except Exception as e:
            return Result().failure(info=Constants.DETECT_JPEAK_LOAD_FAILURE + Constants.FAILURE_REASON["Read_Model_Exception"] + "\n" + format_exc())

        return Result().success(info=Constants.DETECT_JPEAK_LOAD_FINISHED)