import pickle

from   itertools   import combinations
from   collections import defaultdict

import numpy as np
from   scipy.spatial import KDTree

from   interp.grid         import grid
from   interp.grid.simplex import cell
from   interp.tools        import exact_func, exact_func_3D



THREE_NODE_TRIANGLE = 2
FOUR_NODE_TET       = 4

EDGES_FOR_FACE_CONNECTIVITY   = 2
EDGES_FOR_VOLUME_CONNECTIVITY = 3



class gmsh_grid(grid):

  def __init__(self, filename):
    """
      construct an interp.grid.grid-compliant grid
      object out of a 2D gmsh file
    """

    gmsh_file = open(filename, 'r')


    gmsh_file.readline() # $MeshFormat
    format = gmsh_file.readline()
    gmsh_file.readline() # $EndMeshFormat

    gmsh_file.readline() # $Nodes

    node_count = int(gmsh_file.readline())

    # for dim = 2, see note in next for loop
    self.verts = np.empty((node_count, 2))
    self.q     = np.empty(node_count)

    for i in xrange(node_count):
      cur_line = gmsh_file.readline()
      (index, x,y,z) = cur_line.split()
      index = int(index) - 1

      self.verts[i][0] = float(x)
      self.verts[i][1] = float(y)

      # for the general baker method to work, it must have 2
      # components if it it a 2D mesh, so I removed:
      # self.verts[i][2] = float(z)


    grid.__init__(self)
    self.tree = KDTree(self.verts)
    gmsh_file.readline() # $EndNodes
    gmsh_file.readline() # $Elements

    # temporary dict used to compute cell connectivity
    neighbors = {}

    element_count = int(gmsh_file.readline())
    for i in xrange(element_count):
      cur_line = gmsh_file.readline()
      cur_line = cur_line.split()
      cur_cell_index, node_type, rest = (int(cur_line[0]),
                                         int(cur_line[1]),
                                         [int(j) for j in  cur_line[2:]])

      if(node_type == THREE_NODE_TRIANGLE):
        points_for_cur_cell = [i-1 for i in rest[rest[0]+1:]]

        cur_cell = cell(cur_cell_index)

        for cur_point in points_for_cur_cell:
          self.cells_for_vert[cur_point].append(cur_cell)

        cur_cell.verts       = points_for_cur_cell

        self.cells[cur_cell_index] = cur_cell
        edges = [tuple(sorted(i)) for i in combinations(points_for_cur_cell, EDGES_FOR_FACE_CONNECTIVITY)]

        # edge is two verts
        for edge in edges:
          if edge in neighbors:
            neighbors[edge].append(cur_cell_index)
          else:
            neighbors[edge] = [cur_cell_index]

    for k,v in neighbors.iteritems():
      if len(v) > 1:
        self.cells[v[0]].add_neighbor(self.cells[v[1]])
        self.cells[v[1]].add_neighbor(self.cells[v[0]])

  def dump_to_blender_files(self, pfile = '/tmp/points.p', ffile = '/tmp/cells.p'):
    pickle.dump([(p[0], p[1], p[2]) for p in self.verts],   open(pfile, 'w'))
    pickle.dump([f.verts for f in self.cells.itervalues()], open(ffile, 'w'))

class gmsh_grid3D(grid):

  def __init__(self, filename):
    """
      construct an interp.grid.grid-compliant grid
      object out of a 3D gmsh file
    """

    gmsh_file = open(filename, 'r')


    gmsh_file.readline() # $MeshFormat
    format = gmsh_file.readline()
    gmsh_file.readline() # $EndMeshFormat

    gmsh_file.readline() # $Nodes

    node_count = int(gmsh_file.readline())

    self.verts = np.empty((node_count, 3))
    self.q     = np.empty(node_count)

    for i in xrange(node_count):
      cur_line = gmsh_file.readline()
      (index, x,y,z) = cur_line.split()
      index = int(index) - 1

      self.verts[i][0] = float(x)
      self.verts[i][1] = float(y)
      self.verts[i][2] = float(z)
      # self.q[i] = exact_func_3D(self.verts[i])


    grid.__init__(self)
    self.tree = KDTree(self.verts)
    gmsh_file.readline() # $EndNodes
    gmsh_file.readline() # $Elements

    # temporary dict used to compute cell connectivity
    neighbors = {}

    element_count = int(gmsh_file.readline())
    for i in xrange(element_count):
      cur_line = gmsh_file.readline()
      cur_line = cur_line.split()
      cur_cell_index, node_type, rest = (int(cur_line[0]),
                                         int(cur_line[1]),
                                         [int(j) for j in  cur_line[2:]])

      if(node_type == FOUR_NODE_TET):
        points_for_cur_cell = [i-1 for i in rest[rest[0]+1:]]

        cur_cell = cell(cur_cell_index)

        for cur_point in points_for_cur_cell:
          self.cells_for_vert[cur_point].append(cur_cell)

        cur_cell.verts       = points_for_cur_cell

        self.cells[cur_cell_index] = cur_cell
        edges = [tuple(sorted(i)) for i in combinations(points_for_cur_cell, EDGES_FOR_VOLUME_CONNECTIVITY)]

        for edge in edges:
          if edge in neighbors:
            neighbors[edge].append(cur_cell_index)
          else:
            neighbors[edge] = [cur_cell_index]

    for k,v in neighbors.iteritems():
      if len(v) > 1:
        self.cells[v[0]].add_neighbor(self.cells[v[1]])
        self.cells[v[1]].add_neighbor(self.cells[v[0]])

  def dump_to_blender_files(self, pfile = '/tmp/points.p', ffile = '/tmp/cells.p'):
    pickle.dump([(p[0], p[1], p[2]) for p in self.verts],   open(pfile, 'w'))
    pickle.dump([f.verts for f in self.cells.itervalues()], open(ffile, 'w'))