import numpy as np
import pyamg
from scipy import sparse
from scipy.spatial import Delaunay

from linsolver import sparse_solver
from triangulation.delaunay import delaunay


class Element:
    def __init__(self, points, global_indexes, fem):
        self.points = np.array(points)
        self.global_indexes = global_indexes
        self.fem = fem

        self.reference_triangle = np.array([[0, 0], [1., 0], [0, 1.]])
        self.reference_grad = np.array([[-1., -1], [1., 0], [0, 1.]])

    def perform_calculation(self):
        self._calculate_transform()
        self._calculate_stiffness_matrix()
        self._calulate_load_vector()

    def _calculate_transform(self):
        reference_coord = np.array([self.reference_triangle[:, 0], self.reference_triangle[:, 1], [1] * 3])
        transformed_coord = np.array([self.points[:, 0], self.points[:, 1], [1] * 3])

        trans = np.dot(transformed_coord, np.linalg.inv(reference_coord))

        self.transform_matrix = trans[0:-1, 0:-1]

        self.area = abs(np.linalg.det(self.transform_matrix)) / 2

    def _calculate_stiffness_matrix(self):

        transform_matrix_inv = np.linalg.inv(self.transform_matrix)
        self.element_stiffness_matrix = np.zeros((3, 3))

        for row in range(3):
            for col in range(3):
                part_u_left_grad = np.dot(np.dot(self.fem.A, transform_matrix_inv.T), self.reference_grad[row])
                part_u_right_grad = np.dot(transform_matrix_inv.T, self.reference_grad[col])

                part_u_grad = self.area * np.dot(part_u_left_grad, part_u_right_grad)

                part_u = (self.area / 6.0) if row == col else (self.area / 12.0)

                self.element_stiffness_matrix[row, col] = part_u_grad + self.fem.q * part_u

    def _calulate_load_vector(self):

        mean_f = np.mean([self.fem.get_func_value(x) for x in self.points])
        self.element_load_vector = np.array([mean_f * self.area / 3] * 3)


class FiniteElement:
    """
    Finite Element Method to solve the 2D Elliptic Partial Differentiation differential Equation with below form:

        div(A grad(u)) + q u = func

    """

    def __init__(self, points, boundaries, A, q, func, slow_solver=True):
        self.points = np.array(points)
        self.dirichlet_boundaries = np.array(boundaries)
        self.A = A
        self.q = q
        self.f = func
        self.slow_solver = slow_solver

        self.triangles = []
        self.point_num = len(points)

    def solve(self):
        if len(self.triangles) == 0:
            self._get_mesh()
        self._process_each_element()
        self._calculate_global_stiffness_matrix()
        self._calulate_global_load_vector()
        self._deal_with_dirichlet_bound()
        self._solve_linear_equations()

    def update_border_and_func(self, boundaries, func):
        self.dirichlet_boundaries = np.array(boundaries)
        self.f = func

    def get_func_value(self, x):
        if isinstance(self.f, dict):
            return self.f[tuple(x)]
        else:
            return self.f(x)

    def _get_mesh(self):

        if self.slow_solver:
            self.triangles = delaunay(self.points)
        else:
            triangulation = Delaunay(self.points)
            self.triangles = triangulation.simplices

    def _process_each_element(self):
        self.elements = []
        for tri in self.triangles:
            ele = Element(points=[self.points[v] for v in tri], global_indexes=tri, fem=self)
            ele.perform_calculation()
            self.elements.append(ele)

    def _calculate_global_stiffness_matrix(self):

        self.global_stiffness_matrix_row = []
        self.global_stiffness_matrix_col = []
        self.global_stiffness_matrix_data = []

        boundary_indexes = set(self.dirichlet_boundaries[:, 0].astype('int'))

        for ele in self.elements:
            for row in range(3):
                if ele.global_indexes[row] not in boundary_indexes:
                    for col in range(3):
                        self.global_stiffness_matrix_row.append(ele.global_indexes[row])
                        self.global_stiffness_matrix_col.append(ele.global_indexes[col])
                        self.global_stiffness_matrix_data.append(ele.element_stiffness_matrix[row, col])

    def _calulate_global_load_vector(self):

        self.global_load_vector = np.zeros(self.point_num)
        for ele in self.elements:
            for v in range(3):
                self.global_load_vector[ele.global_indexes[v]] += ele.element_load_vector[v]

    def _deal_with_dirichlet_bound(self):
        for index, val in self.dirichlet_boundaries:
            index = int(index)

            self.global_stiffness_matrix_row.append(index)
            self.global_stiffness_matrix_col.append(index)
            self.global_stiffness_matrix_data.append(1)

            self.global_load_vector[index] = val

    def _solve_linear_equations(self):

        if not self.slow_solver:
            self.global_stiffness_matrix_csr = sparse.coo_matrix((self.global_stiffness_matrix_data, (
                self.global_stiffness_matrix_row, self.global_stiffness_matrix_col))).tocsr()
            self.solution = pyamg.solve(self.global_stiffness_matrix_csr, self.global_load_vector, verb=False,
                                        tol=1e-10)
        else:
            global_stiffness_sparse = [np.array(self.global_stiffness_matrix_row),
                                       np.array(self.global_stiffness_matrix_col),
                                       np.array(self.global_stiffness_matrix_data)]
            self.solution = sparse_solver.sparse_gauss_seidel(global_stiffness_sparse, self.global_load_vector,
                                                              sparse_input=True)

            ## these solver methods are for test
            # self.global_stiffness = sparse.coo_matrix((self.global_stiffness_matrix_data, (
            #     self.global_stiffness_matrix_row, self.global_stiffness_matrix_col))).tocsr()
            # self.solution = linsolver.jacobi(self.global_stiffness.toarray(), self.global_load_vector)
            # self.solution = linsolver.gauss_seidel(self.global_stiffness.toarray(), self.global_load_vector)
            # self.solution = sparse_solver.sparse_jacobi(self.global_stiffness.toarray(), self.global_load_vector, sparse_input=False)
            # self.solution = sparse_solver.sparse_gauss_seidel(self.global_stiffness.toarray(), self.global_load_vector, sparse_input=False)

        if isinstance(self.solution, str):
            print("The inputs for linear solver have problems.")

[((12, 22, 12, 38), 'numpy.array', 'np.array', ({(12, 31, 12, 37): 'points'}, {}), '(points)', True, 'import numpy as np\n'), ((16, 34, 16, 70), 'numpy.array', 'np.array', ({(16, 43, 16, 69): '[[0, 0], [1.0, 0], [0, 1.0]]'}, {}), '([[0, 0], [1.0, 0], [0, 1.0]])', True, 'import numpy as np\n'), ((17, 30, 17, 69), 'numpy.array', 'np.array', ({(17, 39, 17, 68): '[[-1.0, -1], [1.0, 0], [0, 1.0]]'}, {}), '([[-1.0, -1], [1.0, 0], [0, 1.0]])', True, 'import numpy as np\n'), ((25, 26, 25, 107), 'numpy.array', 'np.array', ({(25, 35, 25, 106): '[self.reference_triangle[:, (0)], self.reference_triangle[:, (1)], [1] * 3]'}, {}), '([self.reference_triangle[:, (0)], self.reference_triangle[:, (1)],\n    [1] * 3])', True, 'import numpy as np\n'), ((26, 28, 26, 85), 'numpy.array', 'np.array', ({(26, 37, 26, 84): '[self.points[:, (0)], self.points[:, (1)], [1] * 3]'}, {}), '([self.points[:, (0)], self.points[:, (1)], [1] * 3])', True, 'import numpy as np\n'), ((36, 31, 36, 67), 'numpy.linalg.inv', 'np.linalg.inv', ({(36, 45, 36, 66): 'self.transform_matrix'}, {}), '(self.transform_matrix)', True, 'import numpy as np\n'), ((37, 40, 37, 56), 'numpy.zeros', 'np.zeros', ({(37, 49, 37, 55): '(3, 3)'}, {}), '((3, 3))', True, 'import numpy as np\n'), ((53, 35, 53, 73), 'numpy.array', 'np.array', ({(53, 44, 53, 72): '[mean_f * self.area / 3] * 3'}, {}), '([mean_f * self.area / 3] * 3)', True, 'import numpy as np\n'), ((65, 22, 65, 38), 'numpy.array', 'np.array', ({(65, 31, 65, 37): 'points'}, {}), '(points)', True, 'import numpy as np\n'), ((66, 36, 66, 56), 'numpy.array', 'np.array', ({(66, 45, 66, 55): 'boundaries'}, {}), '(boundaries)', True, 'import numpy as np\n'), ((85, 36, 85, 56), 'numpy.array', 'np.array', ({(85, 45, 85, 55): 'boundaries'}, {}), '(boundaries)', True, 'import numpy as np\n'), ((127, 34, 127, 58), 'numpy.zeros', 'np.zeros', ({(127, 43, 127, 57): 'self.point_num'}, {}), '(self.point_num)', True, 'import numpy as np\n'), ((28, 42, 28, 72), 'numpy.linalg.inv', 'np.linalg.inv', ({(28, 56, 28, 71): 'reference_coord'}, {}), '(reference_coord)', True, 'import numpy as np\n'), ((97, 29, 97, 50), 'triangulation.delaunay.delaunay', 'delaunay', ({(97, 38, 97, 49): 'self.points'}, {}), '(self.points)', False, 'from triangulation.delaunay import delaunay\n'), ((99, 28, 99, 49), 'scipy.spatial.Delaunay', 'Delaunay', ({(99, 37, 99, 48): 'self.points'}, {}), '(self.points)', False, 'from scipy.spatial import Delaunay\n'), ((147, 28, 148, 50), 'pyamg.solve', 'pyamg.solve', (), '', False, 'import pyamg\n'), ((153, 28, 154, 80), 'linsolver.sparse_solver.sparse_gauss_seidel', 'sparse_solver.sparse_gauss_seidel', (), '', False, 'from linsolver import sparse_solver\n'), ((32, 24, 32, 60), 'numpy.linalg.det', 'np.linalg.det', ({(32, 38, 32, 59): 'self.transform_matrix'}, {}), '(self.transform_matrix)', True, 'import numpy as np\n'), ((42, 36, 42, 92), 'numpy.dot', 'np.dot', ({(42, 43, 42, 65): 'transform_matrix_inv.T', (42, 67, 42, 91): 'self.reference_grad[col]'}, {}), '(transform_matrix_inv.T, self.reference_grad[col])', True, 'import numpy as np\n'), ((150, 39, 150, 81), 'numpy.array', 'np.array', ({(150, 48, 150, 80): 'self.global_stiffness_matrix_row'}, {}), '(self.global_stiffness_matrix_row)', True, 'import numpy as np\n'), ((151, 39, 151, 81), 'numpy.array', 'np.array', ({(151, 48, 151, 80): 'self.global_stiffness_matrix_col'}, {}), '(self.global_stiffness_matrix_col)', True, 'import numpy as np\n'), ((152, 39, 152, 82), 'numpy.array', 'np.array', ({(152, 48, 152, 81): 'self.global_stiffness_matrix_data'}, {}), '(self.global_stiffness_matrix_data)', True, 'import numpy as np\n'), ((41, 42, 41, 84), 'numpy.dot', 'np.dot', ({(41, 49, 41, 59): 'self.fem.A', (41, 61, 41, 83): 'transform_matrix_inv.T'}, {}), '(self.fem.A, transform_matrix_inv.T)', True, 'import numpy as np\n'), ((44, 42, 44, 85), 'numpy.dot', 'np.dot', ({(44, 49, 44, 65): 'part_u_left_grad', (44, 67, 44, 84): 'part_u_right_grad'}, {}), '(part_u_left_grad, part_u_right_grad)', True, 'import numpy as np\n'), ((145, 47, 146, 85), 'scipy.sparse.coo_matrix', 'sparse.coo_matrix', ({(145, 65, 146, 84): '(self.global_stiffness_matrix_data, (self.global_stiffness_matrix_row, self\n    .global_stiffness_matrix_col))'}, {}), '((self.global_stiffness_matrix_data, (self.\n    global_stiffness_matrix_row, self.global_stiffness_matrix_col)))', False, 'from scipy import sparse\n')]