from itertools import combinations

import numpy as np
from scipy.spatial import KDTree

from interp.grid import grid, contains

THREE_NODE_TRIANGLE = 2
FOUR_NODE_TET = 4
MAX_SEARCH_COUNT = 256


class ggrid(grid):
    def __init__(self, filename, values=None, dimension=3):
        """
            construct an interp.grid.grid-compliant grid
            object out of a {2,3}D gmsh file
        """
        self.dim = dimension
        self.values = None

        gmsh_file = open(filename, 'r')

        gmsh_file.readline()  # $MeshFormat
        fmt = gmsh_file.readline()
        self.version, self.file_type, self.data_size = fmt.split()
        gmsh_file.readline()  # $EndMeshFormat

        gmsh_file.readline()  # $Nodes

        node_count = int(gmsh_file.readline())

        self.points = np.empty((node_count, dimension), dtype=np.float64)
        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.points[i][0] = float(x)
            self.points[i][1] = float(y)
            if self.dim == 3:
                self.points[i][2] = float(z)

        self.tree = KDTree(self.points)

        gmsh_file.readline()  # $EndNodes
        gmsh_file.readline()  # $Elements

        neighbors = {}
        simplices = []
        self.point_to_simplex = [[] for i in xrange(len(self.points))]
        simplex_counter = 0

        element_count = int(gmsh_file.readline())
        for simplex_id in xrange(element_count):
            cur_line = gmsh_file.readline()
            cur_line = cur_line.split()
            simplex_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 and self.dim == 2) \
            or (node_type == FOUR_NODE_TET and self.dim == 3):
                points_for_simplex = [i - 1 for i in rest[rest[0] + 1:]]
                for point in points_for_simplex:
                    self.point_to_simplex[point].append(simplex_counter)
                simplices.append(points_for_simplex)
                simplex_counter += 1

        self.simplex_to_simplex = [[] for i in xrange(len(simplices))]

        for simplex in xrange(len(simplices)):
            edges = [tuple(sorted(i)) for i in \
                        combinations(simplices[simplex], self.dim)]
            for edge in edges:
                if edge in neighbors:
                    neighbors[edge].append(simplex)
                else:
                    neighbors[edge] = [simplex]

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

        self.simplices = np.array(simplices, dtype=np.int32)

    def find_simplex(self, X, max_search_count=MAX_SEARCH_COUNT):
        # get closest point
        (dist, indicies) = self.tree.query(X, 2)
        closest_point = indicies[0]

        simplex = None
        checked_cells = []
        cells_to_check = list(self.point_to_simplex[closest_point])

        attempts = 0
        while not simplex and cells_to_check:
            cur_cell = cells_to_check.pop(0)
            checked_cells.append(cur_cell)

            R = self.points[self.simplices[cur_cell]]
            if contains(X, R):
                simplex = cur_cell
                continue

            attempts += 1
            if attempts >= max_search_count:
                raise Exception("Is the search becoming exhaustive?"
                                    " (%dth attempt)" % attempts)

            for neighbor in self.simplex_to_simplex[cur_cell]:
                if (neighbor not in checked_cells) \
                        and (neighbor not in cells_to_check):
                    cells_to_check.append(neighbor)

        if not simplex:
            raise Exception('no containing simplex found')

        return simplex