import sys
from PySide6.QtWidgets import QApplication, QMainWindow, QGraphicsView
from PySide6.QtUiTools import QUiLoader
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
from PySide6.QtWidgets import QVBoxLayout
import CREAT
import numpy as np
import matplotlib as mpl


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.toolbar = None
        self.canvas = None
        # 加载UI文件
        self.ui = QUiLoader().load('QT/PLT2.ui')
        # 更新图窗
        self.ui.graphicsView.repaint()
        self.ui.graphicsView.update()
        self.graphics_layout = QVBoxLayout(self.ui.graphicsView)
        # 绘制Matplotlib图形
        self.ui.pushButton_4.clicked.connect(lambda: self.plot_density())

    def plot_density(self):
        fig, ax = self.scene_fig()
        data_name = self.ui.lineEdit.text().split(',')
        level_nums = self.ui.spinBox_2.value()
        gWeight = self.ui.spinBox.value()
        data = CREAT.create()
        zmin = 1
        zmax = 0
        for i in range(len(data_name)):
            data.data_pre(data_name[i])
            data.contours_pre(level_nums, gWeight)

            """画出密度等高线"""
            f = data.f
            x_min, x_max = data.x_range[0], data.x_range[1]
            y_min, y_max = data.y_range[0], data.y_range[1]
            levels = data.levels

            xx, yy = np.mgrid[x_min:x_max:200j, y_min:y_max:200j]
            # 填充等高线图中间的区域
            cmap = mpl.colormaps.get_cmap('jet')
            colors = cmap(np.linspace(0, 1, level_nums))
            level = [levels[0], levels[-1]]
            ff = ax.contourf(xx, yy, f, levels=level, colors=colors[level_nums - i - 1:level_nums - i], alpha=1,
                             zorder=len(data_name) - i)
            zmin = min(zmin, ff.zmin)
            zmax = max(zmax, ff.zmax)
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_title('Density Distribution')

        # 生成colorbar
        sm = mpl.cm.ScalarMappable(cmap=cmap, norm=mpl.colors.Normalize(vmin=zmin, vmax=zmax))
        sm.set_array([])
        fig.colorbar(sm, ax=ax)

        wells_name = self.ui.lineEdit_4.text()
        if wells_name == '':
            self.canvas_adjust(fig)
        else:
            wells = CREAT.well_to_edge()
            wells.wells_name_and_position(wells_name)
            wells.welltoedge_distance(data.contours)

            """画井位信息"""
            typee = wells.type
            points = wells.position
            namee = wells.name
            min_distance = wells.min_distance
            contours_p = wells.welltoedge_points
            for i in range(len(points)):
                if typee[i][0] == 0:
                    ax.scatter(points[i][0], points[i][1], marker='o', edgecolors='black', facecolors='none', s=100)
                    ax.scatter(points[i][0], points[i][1], marker='o', edgecolors='black', facecolors='none', s=50,
                               linewidths=1)
                    ax.scatter(points[i][0], points[i][1], marker='o', edgecolors='black', facecolors='none', s=30)
                    ax.text(points[i][0] + min_distance[i] / 20, points[i][1] - min_distance[i] / 20, f'{namee[i][0]}',
                            fontdict={'size': '10', 'color': 'b'}, zorder=len(data_name) + 2)  # 井名信息
                else:

                    ax.scatter(points[i][0], points[i][1], marker='o', edgecolors='black', facecolors='black', s=100, zorder=len(data_name)+1)
                    ax.text(points[i][0] + min_distance[i] / 20, points[i][1] - min_distance[i] / 20, f'{namee[i][0]}',
                            fontdict={'size': '10', 'color': 'b'}, zorder=len(data_name) + 2)  # 井名信息
                    continue
                # 绘制油井边缘最短距离点的箭头并标出距离
                ax.quiver(points[i][0], points[i][1], contours_p[i][0] - points[i][0], contours_p[i][1] - points[i][1],
                          angles='xy', scale=1.03,
                          scale_units='xy', width=0.002, zorder=len(data_name)+1)  # 绘制箭头
                ax.text(points[i][0] * 1 / 2 + contours_p[i][0] * 1 / 2 + min_distance[i] / 20,
                        points[i][1] * 1 / 2 + contours_p[i][1] * 1 / 2 - min_distance[i] / 20, f'{round(min_distance[i], 2)}',
                        fontdict={'size': '10', 'color': 'g'})  # 标出距离
                # ax.text(points[i][0] * 1 / 2 + contours_p[i][0] * 1 / 2 + min_distance[i] / 20,
                #         points[i][1] * 1 / 2 + contours_p[i][1] * 1 / 2 - min_distance[i] / 20,
                #         r'$\theta$=' f'{round(wells.angle[i], 2)}', fontdict={'size': '10', 'color': 'y'})  # 标出角度
            self.canvas_adjust(fig)

    def scene_fig(self):
        fig = Figure()
        # 在Figure对象中添加子图
        ax = fig.add_subplot(111)
        return fig, ax

    def canvas_adjust(self, fig):
        if self.canvas is not None:
            # 从布局中删除旧的 canvas 和 toolbar
            item = self.graphics_layout.takeAt(0)
            while item:
                widget = item.widget()
                if widget:
                    widget.setParent(None)
                item = self.graphics_layout.takeAt(0)
        self.graphics_layout.removeWidget(self.canvas)
        self.canvas = FigureCanvas(fig)
        self.toolbar = NavigationToolbar(self.canvas, self.ui.graphicsView)
        self.toolbar.setParent(self.canvas)
        self.graphics_layout.addWidget(self.canvas)
        self.adjustSize()


if __name__ == '__main__':
    app = QApplication([])
    stats = MainWindow()
    stats.ui.show()
    app.exec()