'''
Utilities for simplifying FEM modelling, e.g. Meshes.
'''

from dolfin import *
import meshio
import numpy as np

class Spheres():
    '''
    Generate non-overlapping spheres in a box to use in meshing.
    '''
    def __init__(self, n=100, r=0.2, ls=[2, 2, 2]):
        '''
        Creates n randomly distributed non-overlapping spheres
        with radius r in in a bounding box of dimensions ls.
        '''
        self.n_points = n # number of points
        self.r_sphere = r # radius of spheres
        self.l_box = np.array(ls) # dimensions of box
        self.point_list = []

    def create_coords(self):
        '''
        Create random coordinates scaled to the right dimensions
        '''
        return np.random.rand(len(self.l_box))*self.l_box

    def check_rad(self, point):
        '''
        Check whether coordinates are within radius r
        '''
        for p in self.point_list:
            if np.sqrt(np.sum((p-point)**2)) < self.r_sphere:
                return 0
        return 1

    def create_list(self):
        '''
        Make the actual list of points.
        '''
        self.point_list.append(self.create_coords())
        while len(self.point_list) < self.n_points:
            point = self.create_coords()
            if self.check_rad(point):
                self.point_list.append(point)

    def write_to_txt(self):
        '''
        Write the indices and positions of the points
        created in Spheres to .txt files named numbers.txt and
        points.txt
        '''
        num_str = ''
        for i in range(len(self.point_list)):
            num_str+=str(i+5)+','
        text_file = open("Meshes/numbers.txt", "w")
        text_file.write(num_str)
        text_file.close()

        point_str = ''
        for (i, p) in enumerate(self.point_list):
            point_str += ('Point('+str(i+5)+')={'+str(p[0])+','
                          +str(p[1])+','+str(p[2])+',0.1};'+'\n')
        text_file = open("Meshes/points.txt", "w")
        text_file.write(point_str)
        text_file.close()
        np.savetxt('Meshes/points_clean.txt', self.point_list, delimiter=',',
                    fmt='%.4f')


def convert_msh_to_xdmf(input_path, output_path):
    '''
    Convert a .msh file (e.g. obtained via gmsh -3 file -format msh2)
    to xdmf format, which can be parallelised.
    '''
    msh = meshio.read(input_path)

    line_cells = []
    for cell in msh.cells:
        if cell.type == "tetra":
            tetra_cells = cell.data
        elif  cell.type == "line":
            if len(line_cells) == 0:
                line_cells = cell.data
            else:
                line_cells = np.vstack([line_cells, cell.data])
        elif cell.type == 'tetra':
            tetra_cells = cell.data

    line_data = []
    for key in msh.cell_data_dict["gmsh:physical"].keys():
        if key == "line":
            if len(line_data) == 0:
                line_data = msh.cell_data_dict["gmsh:physical"][key]
            else:
                line_data = np.vstack([line_data,
                                       msh.cell_data_dict["gmsh:physical"][key]])
    #     elif key == "tetra":
    #         triangle_data = msh.cell_data_dict["gmsh:physical"][key]
        elif key == "tetra":
            tetra_data = msh.cell_data_dict["gmsh:physical"][key]

    # triangle_mesh = meshio.Mesh(points=msh.points,
                                # cells={"triangle": triangle_cells})
    tetra_mesh = meshio.Mesh(points=msh.points, cells={"tetra": tetra_cells})
    line_mesh = meshio.Mesh(points=msh.points,
                            cells=[("line", line_cells)],
                            cell_data={"name_to_read":[line_data]})

    meshio.write(output_path, tetra_mesh)


def save_time_point(f, name):
    ''' Save Fenics function to h5 file.'''
    fFile = HDF5File(MPI.comm_world,name,"w")
    fFile.write(f,"/f")
    fFile.close()

def load_time_point(name, mesh):
    '''Load Fenics function from h5 file.'''
    V = FunctionSpace(mesh, 'CG', 1)
    f = Function(V)
    fFile = HDF5File(MPI.comm_world, name, "r")
    fFile.read(f, "/f")
    fFile.close()
    return f