import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde
from scipy.spatial.distance import cdist
import copy
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from matplotlib.colors import Normalize
import numpy as np
from skimage import measure
import matplotlib.cm as cm

class create():
    def __init__(self):
        """
        data is the x and y coordinate of data.pos
        contours is the edge points
        f is the kernel density value
        levels is the levels for contours
        x_range  is the range of x label
        y_range is the range of y label

        """
        self.data = []
        self.vertices = []
        self.faces = []
        self.levels = []
        self.x_range = []
        self.y_range = []

    def data_pre(self, data_name):
        with open(data_name, 'r') as f:
            lines = f.readlines()
        L_en=len(lines)
        lines= lines[1:L_en]
        data = []
        for line in lines:
            x, y,z,t = line.strip().split("\t")
            data.append(list(map(float, [x,y,z])))
        data = np.array(data)
        self.data = data
        self.x_range = [min(data[:, 0]), max(data[:,0])]
        self.y_range = [min(data[:, 1]), max(data[:,1])]
        self.z_range = [min(data[:, 2]), max(data[:,2])]

    def contours_pre(self, level):

        x = self.data[:, 0]
        y = self.data[:, 1]
        z = self.data[:, 2]
        # 使用scipy库中的gaussian_kde函数计算密度估计
        k = gaussian_kde(self.data.T)
        xi, yi, zi = np.mgrid[x.min()*1.5:x.max()*1.5:30j, y.min()*1.5:y.max()*1.5:30j, z.min()-50:z.max()+50:50j]
        density = k(np.vstack([xi.flatten(), yi.flatten(), zi.flatten()]))
        self.density =density
        level = density.max()*level/100

        # 使用 marching_cubes 生成等值面顶点和面
        verts, faces, _, _ = measure.marching_cubes(density.reshape(xi.shape), level=level)
        a_0 = (x.max() - x.min()) / (verts[:, 0].max() - verts[:, 0].min())
        vertices_0 = (verts[:, 0] - verts[:, 0].min()) * a_0 + x.min()-5
        a_1 = (y.max() - y.min()+10) / (verts[:, 1].max() - verts[:, 1].min())
        vertices_1 = (verts[:, 1] - verts[:, 1].min()) * a_1 + y.min()-5
        a_2 = (z.max() - z.min()+10) / (verts[:, 2].max() - verts[:, 2].min())
        vertices_2 = (verts[:, 2] - verts[:, 2].min()) * a_2 + z.min()-5
        vertices = np.array([vertices_0, vertices_1, vertices_2]).T
        self.vertices = vertices
        self.faces = faces


class well_to_edge():
    def __init__(self):
        """
        name is used to store the well names
        type is the types of the wells
        position is the coordinates of wells
        min_distance is the minimum distances between wells and edge
        welltoedge_points is the points responding to the min_distance
        angle is the angles between the shortest distance direction vector from the well to the edge and the positive direction of the y-axis during clockwise rotation;
        wells_num: the number of wells
        """
        self.name = []
        self.type = []
        self.position = []
        self.min_distance = []
        self.welltoedge_points = []
        self.angle = []
        self.wells_num = 0

    def wells_name_and_position(self, wells_name):
        # 读取井位信息
        with open(wells_name, 'r') as f_j:
            j_ing = f_j.readlines()
        points = []
        typee = []
        namee = []
        for line in j_ing:
            if ('0' or '1') in line:
                name, x, y, z, type = line.strip().split("\t")
                if name != 'name':
                    points.append(list(map(float, [x, y, z])))
                    typee.append(list(map(int, [type])))
                    namee.append(name.split('\n'))
        self.position = points
        self.name = namee
        self.type = typee
        self.wells_num = len(points)