"""
@
@
"""
import os
import sys
from PyQt5.QtCore import Qt, QUrl
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QFileDialog, QVBoxLayout, QLineEdit, QHBoxLayout, \
    QComboBox, QRadioButton, QLabel, QFrame, QListWidget, QListWidgetItem, QCheckBox
from obspy.core import read
import glob
from PyQt5.QtGui import QFont, QColor
from PyQt5.QtWebEngineWidgets import *
from pyecharts.charts import Line, Grid
from pyecharts import options as opts


class PlotSgy:
    def __init__(self, checked_list):
        self.timestamp_count = 0  # 用于多子图绘制时的位置索引
        self.trace_ = None  # 存储某一组对应的道序号
        self.checked_list = checked_list

    def plot(self, fix_timestamp):
        """
        idx：单组全部时间戳还是单个时间戳
        辅助道占3道,从第一个采集道(第四道)开始
            len(mw.time_list)+3---->组数+3---->实际道数
            遍历每一道匹配位置
        first_timestamp, last_timestamp, len(tr.data)--->1688533350.0 1688533365.0 60001
        mark_time---->采样点index
        """
        smooth_state = mw.check_box.isChecked()
        for x, y in self.time_file_dic.items():
            if int(y[0]) < int(fix_timestamp) < int(y[1]):
                st = read(x, format='SEGY')
                #   只有一道
                if str(st).split(' ', 1)[0] == '1':
                    print('单道输出')
                    self.trace_ = 0
                else:
                    #   包含多道
                    if self.trace_ != 0:
                        for t in range(3, int(str(st).split(' ', 1)[0])):
                            tr = st[t]
                            mark_time = int((len(tr.data) - 1) * ((fix_timestamp - y[0]) / (y[1] - y[0])))
                            if max(tr.data[(mark_time - mw.offset_num):(mark_time + mw.offset_num)]) > mw.pick_num:
                                self.trace_ = t
                                break

                if self.trace_ is not None:
                    tr = st[self.trace_]
                    mark_time = int((len(tr.data) - 1) * ((fix_timestamp - y[0]) / (y[1] - y[0])))
                    self.timestamp_count += 1
                    timestamp_list = list(
                        range((int(fix_timestamp) - mw.offset_num), (int(fix_timestamp) + mw.offset_num)))
                    signal_list = tr.data[(mark_time - mw.offset_num):(mark_time + mw.offset_num)].tolist()
                    if mw.radio_y_draw.isChecked():
                        #   纵向
                        line = (
                            Line()
                                .set_global_opts(
                                title_opts=opts.TitleOpts(title=str(fix_timestamp),
                                                          pos_left=(str((self.timestamp_count - 1) * 18)) + "%", ),
                                tooltip_opts=opts.TooltipOpts(is_show=True,
                                                              trigger_on="mousemove | click",
                                                              axis_pointer_type='cross'
                                                              ),
                                xaxis_opts=opts.AxisOpts(type_="value",
                                                         ),
                                yaxis_opts=opts.AxisOpts(
                                    type_="value",
                                    min_=timestamp_list[0],
                                    max_=timestamp_list[-1],
                                    interval=1,
                                ),
                            )
                                .add_xaxis(
                                xaxis_data=signal_list)
                                .add_yaxis(
                                series_name="",
                                y_axis=timestamp_list,
                                symbol="emptyCircle",
                                is_symbol_show=True,
                                label_opts=opts.LabelOpts(is_show=False),
                                is_smooth=smooth_state,
                                areastyle_opts=opts.AreaStyleOpts(opacity=0.5),
                            )
                        )
                        horizontal_line_data = [fix_timestamp for _ in signal_list]
                        line.add_yaxis("MRAK LINE", horizontal_line_data,
                                       is_symbol_show=False,
                                       is_connect_nones=True,
                                       linestyle_opts=opts.LineStyleOpts(  # 设置线条样式
                                           color="teal",  # 线条颜色
                                           width=2,  # 线条宽度
                                           type_="dotted",  # 线条类型，可以是 solid, dashed, dotted
                                           opacity=0.8  # 线条透明度
                                       ),
                                       )
                        if self.timestamp_count == 1:
                            left_position = "3%"
                        else:
                            left_position = (str((self.timestamp_count - 1) * 18 + 3)) + "%"
                        self.grid.add(line,
                                      grid_opts=opts.GridOpts(pos_left=left_position,
                                                              height="95%",
                                                              width="15%"))
                    else:
                        timestamp_list = [str(x) for x in timestamp_list]
                        x_point = str(fix_timestamp)
                        y_point = signal_list[timestamp_list.index(str(fix_timestamp))]
                        if self.timestamp_count == 1:
                            top_position = "3%"
                        else:
                            top_position = (str((self.timestamp_count - 1) * 18 + 3)) + "%"
                        line = (
                            Line()
                                .set_global_opts(
                                title_opts=opts.TitleOpts(title=str(fix_timestamp),
                                                          pos_top=(str((self.timestamp_count - 1) * 18)) + "%", ),
                                tooltip_opts=opts.TooltipOpts(is_show=True,
                                                              trigger_on="mousemove | click",
                                                              axis_pointer_type='cross'
                                                              ),
                                xaxis_opts=opts.AxisOpts(
                                    interval=1,
                                    splitline_opts=opts.SplitLineOpts(is_show=True), ),
                                yaxis_opts=opts.AxisOpts(
                                    type_="value",
                                    axistick_opts=opts.AxisTickOpts(is_show=True),
                                    splitline_opts=opts.SplitLineOpts(is_show=True),
                                ),
                            )
                                .add_xaxis(xaxis_data=timestamp_list)
                                .add_yaxis(
                                series_name="",
                                y_axis=signal_list,
                                symbol="emptyCircle",
                                is_symbol_show=True,
                                label_opts=opts.LabelOpts(is_show=False),
                                is_smooth=smooth_state,
                                areastyle_opts=opts.AreaStyleOpts(opacity=0.5),
                                markpoint_opts=opts.MarkPointOpts(
                                    data=[opts.MarkPointItem(name='标记点',
                                                             coord=[x_point, y_point],
                                                             value=x_point,
                                                             symbol_size=[200, 20],
                                                             symbol='rect',
                                                             itemstyle_opts={'color': 'rgba(0, 0, 0, 0)'},

                                                             )],
                                    label_opts=opts.LabelOpts(position="outside", color="teal", font_size=18),
                                )
                            )
                        )
                        self.grid.add(line,
                                      grid_opts=opts.GridOpts(pos_top=top_position,
                                                              height="15%",
                                                              width="95%"))
                    break
        self.grid.render("line_chart.html")
        url4 = os.getcwd() + '\line_chart.html'
        mw.browser.load(QUrl.fromLocalFile(url4))

    def work(self):
        self.grid = Grid(init_opts=opts.InitOpts(width="2800px", height="1800px"))
        mw.offset_num = int(mw.offset_value_in.text())
        mw.pick_num = float(mw.pick_value_in.text())
        self.time_file_dic = {}
        try:
            for sgy_file in mw.file_list:  # 遍历每个sgy文件
                st = read(sgy_file, format='SEGY')
                self.time_file_dic[sgy_file] = [st[0].times('timestamp')[0], st[0].times('timestamp')[-1]]
            for tp in self.checked_list:  # 遍历每个时间戳
                tp = int(str(tp)[:10])  # 截取时间戳
                self.plot(tp)
        except Exception as e:
            print('run plot fuc',e)


class MyWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.path1 = None
        self.path2 = None
        self.time_list = None  # 时间戳列表
        self.file_list = None  # sgy文件列表
        self.browser = QWebEngineView()
        self.init_ui()

    def init_ui(self):
        # 创建按钮
        btn_choose_folder = QPushButton('选择SEGY文件夹', self)
        btn_choose_folder.setFixedWidth(300)
        self.btn_folder_path = QPushButton('', self)
        self.btn_folder_path.setFixedWidth(300)
        btn_choose_file = QPushButton('选择时间戳文件', self)
        btn_choose_file.setFixedWidth(300)
        self.btn_file_path = QPushButton('', self)
        self.btn_file_path.setFixedWidth(300)
        self.btn_init = QPushButton('初始化', self)
        self.btn_init.setFixedWidth(300)
        self.radio_x_draw = QRadioButton('横向绘图', self)
        self.radio_y_draw = QRadioButton('纵向绘图', self)
        self.check_box = QCheckBox('平滑绘图')
        self.radio_y_draw.setChecked(True)
        self.check_box.setChecked(True)
        btn_work = QPushButton('Start Plot', self)
        btn_work.setFixedWidth(300)
        button_style = '''
            QPushButton {
                background-color: #98F5FF;
                border: 2px solid #4CAF50;
                border-radius: 5px;
                padding: 10px 20px;
            }

            QPushButton:hover {
                background-color: #FFECB5;
            }

            QPushButton:pressed {
                background-color: #00FF7F;
                border: 2px solid #45a049;
            }
        '''
        # button_style = "QPushButton { border: 3px solid teal; border-radius: 5px; padding: 5px; }"
        btn_work.setStyleSheet(button_style)
        self.btn_init.setStyleSheet(button_style)

        offset_btn = QPushButton('道数', self)
        offset_btn.setFixedWidth(300)
        pick_value_btn = QPushButton('极大值', self)
        pick_value_btn.setFixedWidth(300)

        self.group_comb = QComboBox(self)
        self.group_comb.setFixedWidth(300)
        # time_list_widget
        self.time_list_wid = QListWidget(self)
        self.time_list_wid.setFixedWidth(300)
        self.time_list_wid.setMinimumHeight(450)

        self.offset_value_in = QLineEdit(self)
        self.offset_value_in.setFixedWidth(300)
        self.offset_value_in.setAlignment(Qt.AlignCenter)
        self.offset_value_in.setText('10')
        self.pick_value_in = QLineEdit(self)
        self.pick_value_in.setFixedWidth(300)
        self.pick_value_in.setAlignment(Qt.AlignCenter)
        self.pick_value_in.setText('0.1')

        # 创建文本输入框
        self.text_input = QLabel(self)
        self.text_input.setFixedHeight(24)

        # 设置布局
        all_box = QHBoxLayout()

        # 左侧布局
        self.left_box = QVBoxLayout()
        vbox1 = QVBoxLayout()
        vbox2 = QVBoxLayout()
        vbox3 = QVBoxLayout()
        vbox4 = QVBoxLayout()
        # 右侧布局
        self.frame = QFrame(self)
        self.right_box = QVBoxLayout(self.frame)
        self.right_box.addWidget(self.browser)

        vbox1.addWidget(btn_choose_folder)
        vbox1.addWidget(self.btn_folder_path)
        vbox2.addWidget(btn_choose_file)
        vbox2.addWidget(self.btn_file_path)
        vbox3.addWidget(self.group_comb)
        vbox3.addWidget(self.time_list_wid)
        vbox4.addWidget(offset_btn)
        vbox4.addWidget(self.offset_value_in)
        vbox4.addWidget(pick_value_btn)
        vbox4.addWidget(self.pick_value_in)
        vbox4.addWidget(self.radio_x_draw)
        vbox4.addWidget(self.radio_y_draw)
        vbox4.addWidget(self.check_box)

        self.left_box.addLayout(vbox1)
        self.left_box.addLayout(vbox2)
        self.left_box.addWidget(self.btn_init)
        self.left_box.addLayout(vbox3)
        self.left_box.addLayout(vbox4)
        self.left_box.addWidget(btn_work)
        self.left_box.addStretch()

        all_box.addLayout(self.left_box, 1)
        all_box.addWidget(self.frame, 6)

        # 设置主窗口布局
        self.setLayout(all_box)

        # 按钮点击事件连接到相应的槽函数
        btn_choose_folder.clicked.connect(self.show_folder_dialog)
        btn_choose_file.clicked.connect(self.show_file_dialog)
        self.group_comb.currentIndexChanged.connect(self.on_combobox_changed)
        btn_work.clicked.connect(self.do)
        self.time_list_wid.itemClicked.connect(self.check_count)

        self.btn_init.clicked.connect(self.data_init)
        # 设置窗口属性
        self.setGeometry(300, 300, 1200, 600)
        self.show()

    def show_folder_dialog(self):
        folder_dialog = QFileDialog.getExistingDirectory(self, '选择文件夹', r'C:\Users\demo\Desktop\工具')
        if folder_dialog:
            self.btn_folder_path.setText(folder_dialog)
        self.path1 = folder_dialog

    def show_file_dialog(self):
        file_dialog, _ = QFileDialog.getOpenFileName(self, '选择文件', r'C:\Users\demo\Desktop\工具',
                                                     'All Files (*);;Text Files (*.txt)')
        if file_dialog:
            self.btn_file_path.setText(file_dialog)
        self.path2 = file_dialog

    def data_init(self):
        self.group_comb.clear()
        try:
            if self.path1 is None:
                self.path1 = r'C:\Users\demo\Desktop\工具\0705'
            else:
                self.path1 = self.btn_folder_path.text()
            if self.path2 is None:
                self.path2 = r'C:\Users\demo\Desktop\工具\0705\sjc.csv'
            else:
                self.path2 = self.btn_file_path.text()
            if self.path1 and self.path2:
                self.path1 = self.path1 + '/*.sgy'
                ls = []
                with open(self.path2, 'r', ) as f:
                    for row in f:
                        row = row.strip()
                        ls.append(row.split(','))
                ls0 = list(zip(*ls))
                self.time_list = list(map(list, ls0))
                for i in range(len(self.time_list)):
                    self.group_comb.addItem(f'第{i + 1}组')
                self.file_list = glob.glob(self.path1)
        except Exception as e:
            print(e)

    def on_combobox_changed(self, index):
        # self.time_list_wid.setSelectionMode(QAbstractItemView.MultiSelection)
        self.time_list_wid.clear()
        for x in self.time_list[index]:
            item = QListWidgetItem()
            item.setText(str(x))
            item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
            item.setCheckState(Qt.Unchecked)
            self.time_list_wid.addItem(item)

    def check_count(self):
        limit = 0
        for index in range(self.time_list_wid.count()):
            item = self.time_list_wid.item(index)
            if item.checkState() == Qt.Checked:
                item.setBackground(QColor("teal"))
                limit += 1
            else:
                item.setBackground(QColor("white"))
        if limit > 5:
            print('limit error')

    def do(self):
        checked_time_list = []
        for index in range(self.time_list_wid.count()):
            item = self.time_list_wid.item(index)
            if item.checkState() == Qt.Checked:
                checked_time_list.append(int(item.text()))
        if self.time_list:
            ps = PlotSgy(checked_time_list)
            try:
                ps.work()
            except Exception as e:
                print(e)


if __name__ == '__main__':
    font = QFont('Arial', 12)
    app = QApplication(sys.argv)
    app.setFont(font)
    mw = MyWindow()
    mw.setWindowTitle('Plot_Pulse')
    mw.show()
    sys.exit(app.exec_())