from .surface import Surface
import numpy as np
from scipy import interpolate
from skimage import measure

__author__ = "Pedram Tavadze"
__maintainer__ = "Pedram Tavadze"
__email__ = "petavazohi@mail.wvu.edu"
__date__ = "March 31, 2020"


class Isosurface(Surface):
    def __init__(self,
             XYZ=None,
             V=None,
             isovalue=None,
             V_matrix=None,
             algorithm='lewiner',
             interpolation_factor=1,
             padding=None,
             transform_matrix=None,
             boundaries=None,):
        
        """
        Parameters
        ----------
        XYZ : TYPE, optional
            DESCRIPTION. The default is None.
        V : TYPE, optional
            DESCRIPTION. The default is None.
        isovalue : TYPE, optional
            DESCRIPTION. The default is None.
        V_matrix : TYPE, optional
            DESCRIPTION. The default is None.
        dxyz : TYPE, optional
            DESCRIPTION. The default is None.
        algorithm : TYPE, optional
            DESCRIPTION. The default is 'lewiner'.
        interpolation_factor : TYPE, optional
            DESCRIPTION. The default is 1.
        padding : TYPE, optional
            DESCRIPTION. The default is [0,0,0].

        Returns
        -------
        None.

        """
        
        """
        This class contains a surface that finds all the poins correcponding 
        to the following equation
        V(X,Y,Z) = f
        :param XYZ: a list of coordinates [[x1,y1,z1],[x2,y2,z2],...] 
            corresponding V
        :param V: a list of values [V1,V2,...] corresponding to XYZ
        XYZ[0] >>> V[0]
        XYZ[1] >>> V[1]
        ...
        :param isovalue: The constant value of the surface (f)
        :param V_matrix: one can present V_matrix instead of XYZ and V. 
            V_matrix is a matrix representation of XYZ and V together. This 
            matrix is generated if XYZ and V are provided. 
        :param algorithm: The algorithm used to find the isosurface, This 
            function used scikit-image to find the isosurface. possibilities 
            ['classic','lewiner']
        :param interpolation_factor: This module uses Fourier Transform 
            interpolation. interpolation factor will increase the grid points
            in each direction by a this factor, the dafault is set to 1
        :param padding: padding is used for periodic datasets such as bands in
            a solid state calculation. e.g. The 1st BZ is not covered fully so
            one might want to pad the matrix with wrap(look at padding in 
            numpy for wrap), afterwards one has to clip the surface to the 
            first BZ. easily doable using pyvista of trimesh
            padding goes as follows np.pad(self.eigen_matrix,
                              ((padding[0]/2, padding[0]/2),
                              (padding[1]/2, padding[1]/2)
                              (padding[2]/2, padding[2])),
                              "wrap")
            In other words it creates a super cell withpadding
        """
    
        self.XYZ = np.array(XYZ)
        self.V = V
        self.isovalue = isovalue
        self.V_matrix = V_matrix
        self.algorithm = algorithm
        self.padding = padding
        self.interpolation_factor = interpolation_factor
        self.transform_matrix = transform_matrix
        self.boundaries = boundaries
        if self.algorithm not in ['classic','lewiner']:
            print("The algorithm chose has to be from ['classic','lewiner'], automtically choosing 'lewiner'")
            self.algorithm = 'lewiner'
    
        if self.V_matrix is None:
            self.V_matrix = map2matrix(self.XYZ,self.V)
        
        if self.padding is None:
            self.padding = [self.nX//2,self.nY//2,self.nZ//2]

        verts, faces, normals, values = self.get_isosurface(interpolation_factor)
        if verts is not None and faces is not None:
            if transform_matrix is not None:
                verts = np.dot(verts,transform_matrix)
            """
            Python, unlike statically typed languages such as Java, allows complete
            freedom when calling methods during object initialization. However, 
            standard object-oriented principles apply to Python classes using deep 
            inheritance hierarchies. Therefore the developer has responsibility for 
            ensuring that objects are properly initialized when there are multiple 
            __init__ methods that need to be called.
            For this reason I will make one temporary surface and from there I will
            using the other surface provided.
            """
            if boundaries is not None:
                suprecell_surface = Surface(verts=verts,faces=faces)
                verts, faces = self.clip(suprecell_surface,boundaries)

        Surface.__init__(self,verts=verts,faces=faces)

    

    def clip(self,S1,S2):
        """
        This function clips S1 using the boundaries of S2 and returns 

        Parameters
        ----------
        S1 : TYPE
            DESCRIPTION.
        S2 : TYPE
            DESCRIPTION.

        Returns
        -------
        verts,faces

        """
        
        for iface in range(len(S2.faces)):
            normal = S2.face_normals[iface]
            center = np.average(S2.verts[S2.faces[iface]],axis=0)
            S1.pyvista_obj.clip(origin=center,
                                          normal=normal,inplace=True)
        faces = []
        courser = 0
        for i in range(S1.pyvista_obj.n_faces):
            npoints = S1.pyvista_obj.faces[courser]
            face = []
            courser += 1
            for ipoint in range(npoints):
                face.append(S1.pyvista_obj.faces[courser])
                courser +=1
            faces.append(face)
            
        return S1.pyvista_obj.points,faces


    @property
    def X(self):
        return np.unique(self.XYZ[:, 0])

    @property
    def Y(self):
        return np.unique(self.XYZ[:, 1])

    @property
    def Z(self):
        return np.unique(self.XYZ[:, 2])
        
    @property
    def dxyz(self):
        dx = np.abs(self.X[-1] - self.X[-2])
        dy = np.abs(self.Y[-1] - self.Y[-2])
        dz = np.abs(self.Z[-1] - self.Z[-2])
        return [dx,dy,dz]

    @property
    def nX(self):
        return len(self.X)

    @property
    def nY(self):
        return len(self.Y)

    @property
    def nZ(self):
        return len(self.Z)

    @property
    def surface_boundaries(self):
        
        padding_x = self.padding[0]
        padding_y = self.padding[1]
        padding_z = self.padding[2]

        eigen_matrix = np.pad(self.V_matrix,
                              ((padding_x, padding_x), (padding_y, padding_y), 
                                (padding_z, padding_z)), "wrap")
        try:
            verts, faces, normals, values = measure.marching_cubes_lewiner(
                eigen_matrix, self.fermi)
            for ix in range(3):
                verts[:, ix] -= verts[:, ix].min()
                verts[:, ix] -= (verts[:, ix].max() - verts[:, ix].min()) / 2 #+self.origin[ix]
                verts[:, ix] *= self.dxyz[ix] 
            mins = verts.min(axis=0)
            maxs = verts.max(axis=0)
            
            return [(mins[0],maxs[0]),(mins[1],maxs[1]),(mins[2],maxs[2])]
        except:
            return None
        
                
        

    def get_isosurface(self, interp_factor=1):

        # Amount of kpoints needed to add on to fully sample 1st BZ
        
        padding_x = self.padding[0]
        padding_y = self.padding[1]
        padding_z = self.padding[2]

            
        eigen_matrix = np.pad(self.V_matrix,
                              ((padding_x, padding_x), (padding_y, padding_y), 
                               (padding_z, padding_z)), "wrap")

        bnd = self.surface_boundaries
        
        if interp_factor != 1:
            # Fourier interpolate the mapped function E(x,y,z)
            
            eigen_matrix = fft_interpolate(eigen_matrix, interp_factor)

            # after the FFT we loose the center of the BZ, using numpy roll we
            # bring back the center of the BZ
            # eigen_matrix = np.roll(eigen_matrix,4  ,
            #     axis=[0, 1, 2])
            
        try:
            verts, faces, normals, values = measure.marching_cubes_lewiner(
                eigen_matrix, self.fermi)
        except BaseException:
            print("No isosurface for this band")
            return None, None, None, None
        # recenter
        
        for ix in range(3):
            verts[:, ix] -= verts[:, ix].min()
            verts[:, ix] -= (verts[:, ix].max() - verts[:, ix].min()) / 2 #+self.origin[ix]
            verts[:, ix] *= self.dxyz[ix] / interp_factor
            if bnd is not None and interp_factor != 1:
                verts[:,ix] -= (verts[:,ix].min() - bnd[ix][0])
            #     x_shift = verts[:,0].min() - bnd[0]
            # y_shift = verts[:,1].min() - bnd[1]
            # z_shift = verts[:,2].min() - bnd[2]
            
        # transfare from fraction to cartesian
        # verts = np.dot(verts, self.reciprocal_)
        # new_faces = np.zeros(shape=(len(faces), 4))
        # new_faces[:, 0] = 3
        # new_faces[:, 1:] = faces
        # faces = new_faces
        return verts, faces, normals, values    
        
def map2matrix(XYZ,V):
    XYZ = XYZ
    V = V
    X = np.unique(XYZ[:, 0])
    Y = np.unique(XYZ[:, 1])
    Z = np.unique(XYZ[:, 2])
    mapped_func = np.zeros(shape=(len(X), len(Y), len(Z)))
    #kpoint_matrix = np.zeros(shape=(len(kx), len(ky), len(kz), 3)) This was added to check if the mesh grid is working
    for ix in range(len(X)):
        condition1 = XYZ[:, 0] == X[ix]
        for iy in range(len(Y)):
            condition2 = XYZ[:, 1] == Y[iy]
            for iz in range(len(Z)):
                condition3 = XYZ[:, 2] == Z[iz]
                tot_cond = np.all([condition1, condition2, condition3], axis=0)
                if len(V[tot_cond]) != 0:
                    mapped_func[ix, iy, iz] = V[tot_cond][0]
                    # kpoint_matrix[ikx, iky, ikz] = [
                    #     kx[ikx], ky[iky], kz[ikz]]
                else:
                    mapped_func[ix, iy, iz] = np.nan
                    # kpoint_matrix[ikx, iky, ikz] = [np.nan, np.nan, np.nan]
    return mapped_func

def fft_interpolate(function, interpolation_factor=2):
    """
    if I = interpolation_factor
    This function withh recieve f(x,y,z) with dimensions of (nx,ny,nz)
    and returns f(x,y,z) with dimensions of (nx*I,ny*I,nz*I)
    """
    eigen_fft = np.fft.fftn(function)
    shifted_fft = np.fft.fftshift(eigen_fft)
    nx, ny, nz = np.array(shifted_fft.shape)
    pad_x = nx * (interpolation_factor - 1) // 2
    pad_y = ny * (interpolation_factor - 1) // 2
    pad_z = nz * (interpolation_factor - 1) // 2
    new_matrix = np.pad(
        shifted_fft,
        ((pad_x, pad_x), (pad_y, pad_y), (pad_z, pad_z)),
        "constant",
        constant_values=0,
    )

    new_matrix = np.fft.ifftshift(new_matrix)
    interpolated = np.real(np.fft.ifftn(new_matrix)) * (interpolation_factor ** 3)
    return interpolated