import numpy as np
from mesh_to_sdf import mesh_to_sdf
import sys
import trimesh                            
import warnings                           
#
# Function to read the geometry data from CaNS
#
def readGrid(filenamei,iskip,r0,non_uniform_grid):
    ''' 
    This function reads the grid generated by CaNS. 
        Originally written by Pedro Costa.
    INPUT
        filenamei:          [string] Input name of the file
        iskip:              [3x1 array] that has data skip information for the written array
        r0:                 [3x1 array] that defines the origin of the domain
        non_uniform_grid:   [Boolean] Is the grid non-uniform?
    RETURNS
        xp, yp, zp:         Grid location for the scalar
        xu, yv, zw:         Grid location for the face values
    '''
    #
    # setting up some parameters
    #
    precision  = 'float64'    # precision of float
    #
    # read geometry file
    #
    geofile  = filenamei + "geometry.out"
    geo = np.loadtxt(geofile, comments = "!", max_rows = 2)
    ng = geo[0,:].astype('int')
    l  = geo[1,:]
    dl = l/(1.*ng)
    #
    # read and generate grid
    #
    xp = np.arange(r0[0]+dl[0]/2.,r0[0]+l[0],dl[0]) # centered  x grid
    yp = np.arange(r0[1]+dl[1]/2.,r0[1]+l[1],dl[1]) # centered  y grid
    zp = np.arange(r0[2]+dl[2]/2.,r0[2]+l[2],dl[2]) # centered  z grid
    xu = xp + dl[0]/2.                              # staggered x grid
    yv = yp + dl[1]/2.                              # staggered y grid
    zw = zp + dl[2]/2.                              # staggered z grid
    if(non_uniform_grid):
        gridfile=filenamei + "grid.bin"
        f   = open(gridfile,'rb')
        grid_z = np.fromfile(f,dtype=precision)
        f.close()
        grid_z = np.reshape(grid_z,(ng[2],4),order='F')
        zp = r0[2] + np.transpose(grid_z[:,2]) # centered  z grid
        zw = r0[2] + np.transpose(grid_z[:,3]) # staggered z grid
        #
        # reshape grid
        #
        xp = xp[0:ng[0]:iskip[0]]
        yp = yp[0:ng[1]:iskip[1]]
        zp = zp[0:ng[2]:iskip[2]]
        xu = xu[0:ng[0]:iskip[0]]
        yv = yv[0:ng[1]:iskip[1]]
        zw = zw[0:ng[2]:iskip[2]]
    return xp, yp, zp, xu, yv, zw
#
# Load Dopamine grid
#
def loaddopaminegrid(infile):
    '''
        This function loads the grid information supplied by the user
    INPUT
        infile:     [string] Name and location of the input file
    OUTPUT
        x,y,z:      [numpy arrays] X, Y, and Z face locations
        xm,ym,zm:   [numpy arrays] Xm, Ym, and Zm cell-center locations
    '''
    x = np.loadtxt(infile,max_rows=1)
    y = np.loadtxt(infile,skiprows=1,max_rows=1)
    z = np.loadtxt(infile,skiprows=2,max_rows=1)
    xm = np.loadtxt(infile,skiprows=3,max_rows=1)
    ym = np.loadtxt(infile,skiprows=4,max_rows=1)
    zm = np.loadtxt(infile,skiprows=5,max_rows=1)

    return x,y,z,xm,ym,zm
#
# Function definition for computing the signed-distance field
#
def computeSDF(meshObj,xin,yin,zin,nSamples):
    '''
        This function computes the SDF for the entire grid
    INPUT
        meshObj:        [mesh object] Mesh object of the stl/obj file
        xin:            [numpy array] x array corresponding to the streamwise coordinate
        yin:            [numpy array] y array corresponding to the spanwise coordinate
        zin:            [numpy array] z array corresponding to the vertical coordinate
        nSamples:       [float .or. integer] Number of sample points used in the scan method
    OUTPUT
        sdf:            [float64] Signed distance field        
    '''
    #
    # Compute the signed-distance field
    #
    dummy = np.vstack(np.meshgrid(xin,yin,zin,indexing='ij')).reshape(3,-1).T
    sdf = mesh_to_sdf(meshObj,dummy,surface_point_method='sample', sign_method='normal', bounding_radius=None, scan_count=100, scan_resolution=100, sample_point_count=nSamples, normal_sample_count=11)

    return sdf
#
# Function to convert the numpy array to CaNS compatible binary files
#
# This file has the functions to convert arrays from numpy to CaNS compatible binary
def n2carray(filein,N,fileout):
    '''
        This function converts the filein to binary array that is CaNS compatible
    INPUT
        filein: [string] Location and name of the numpy array
        N:      [3 x 1 numpy array] Number of grid points in x,y, and z
    RETURNS
        fileout: [string] Location and name of the output binary array that is CaNS compatible
    '''
    #
    # Load the file
    #
    data = np.load(filein)
    data = np.reshape(data,N)		        # First reshape using row major array
    data = np.asfortranarray(data)	        # Cast as column major array
    data = data.flatten(order='F')	        # Flatten it to make FORTRAN binary compatible
    print("%s has size %d"%(filein,sys.getsizeof(data)))
    #
    # Write file to binary array
    #
    # write the array to a binary file
    with open(fileout, 'wb') as f:
        data.tofile(f)
    f.close()
#
# Translate and rotate geometry
#     
def trasRot(saveMesh=False,inMesh='dummy.stl',outMesh='output.stl',tarr=[0.0,0.0,0.0],rotang=0.0,rotax=[0.0,0.0,0.0],scalegeo=0.0):    
    '''
        This function translates and rotates the geometry based on user specification
    INPUT
        saveMesh:   [Boolean] Do you wish to save the mesh?  
        inMesh:     [string] Name of the mesh to be read        
        tarr:       [3 x 1 list] Translation array in x, y, and z
        rotang:     [float] Rotation angle in degrees
        rotax:      [3 x 1 list] Rotation axis
        scalegeo:   [float] scale the input geometry by a factor
    OUTPUT:
        outMesh:    [string] Name of the output mesh, default = output.stl
    '''
    #
    # Exit if inputfile is not defined
    #
    if(inMesh=='dummy.stl'):
        sys.exit("ERROR: Please specify the input filename and location....")
    #
    # Preliminary setup
    #
    rotAngRad = np.radians(rotang)                                      # Convert to radians        
    rotMat = trimesh.transformations.rotation_matrix(rotAngRad, rotax)  # Make the rotation matrix
    #
    # Load the mesh
    #
    mesh = trimesh.load(inMesh)
    #
    # Check if the mesh if watertight
    #
    waterTight = mesh.is_watertight     # mesh watertight? [Important for generating SDF]
    #
    # Print message
    #
    if(waterTight == 0):
        # Print message
        warnings.warn("Input file is not watertight, this may cause issues for generating the SDF!")
        print("------------------------------------------------------------------------------------")
    else:
        print("Input file is watertight")
        print("------------------------------------------------------------------------------------")
    #
    # Bounding box
    #
    print("Oriented Bounding box [before operations]")
    print("xmin, ymin, zmin --",mesh.bounds[0])
    print("xmax, ymax, zmax --",mesh.bounds[1])
    print("------------------------------------------------------------------------------------")
    #
    # Rotate the mesh
    #
    mesh.apply_transform(rotMat)
    #
    # Scale the mesh to user prompted scale
    #
    mesh.vertices *= scalegeo 
    #
    # Translate mesh
    #
    mesh.vertices += tarr
    #
    # Bounding box
    #
    print("Oriented Bounding box [after operations]")
    print("xmin, ymin, zmin --",mesh.bounds[0])
    print("xmax, ymax, zmax --",mesh.bounds[1])
    print("------------------------------------------------------------------------------------")
    #
    # Save Mesh now
    #
    if saveMesh:
        mesh.export(outMesh)
    #
    # Exit
    #
    print("Done.......")
