"""Defines a Python class for hypersurfaces"""  

from multiprocessing import Pool
import mpmath
import numpy as np
import sympy as sp
import tensorflow as tf

__all__ = ['Hypersurface', 'diff', 'diff_conjugate']

class Hypersurface():
    """A hypersurface or patch defined both symbolically and numerically.   

    The Hypersuface class contains the symbolic polynomial representation of a
    hypersurface in sympy. It is also numerically defined as a collection of
    points on the hypersurface. The points are sperated into patches, which are
    also collectons of points. Therefore, recursively, the patches can also be
    defined as instances of the Hypersurface class. 

                           Hypersurface
                           /     |     \
                      patch1   patch2  patch3   (Also Hypersurface class)
                      /  |  \  / | \   / | \
               subpatch ..  ..  .. ..   .. .. 

    Attribute:
    -------------
    coordinates: 
        The homogeneous coordinates as a list of sympy symbols, e.g. 
        z0, z1, z2, z3, z4 = sp.symbols('z0, z1, z2, z3, z4')    
        Z = [z0, z2, z3, z3, z4]
    function: 
        A function of the homogeneous coordiantes, e.g.
        f = z0**5 + z1**5 + z2**5 + z3**5 + z4**5 + 0.5*z0*z1*z2*z3*z4 
        The hypersurface is defined by f = 0
    n_dim: The number of homogeneous coordiantes.
    norm_coordinate:
        Applicable if the instance is a patch. An integer representing the
        index of the coordinate set to 1 on the affine patch. The first level 
        of patches are defined with this coordinate. The corresponding symbolic 
        coordinate is self.coordiante[self.norm_coordiante].
    affine_coordinates:
        The coordiantes on the affine patches (withouth the norm_coordinate). 
    max_grad_coordinate:
        The index of the affine coordinate with the largest |∂f/∂z_i|. The 
        second level of patches (subpatches) are defined using this coordinate,
        so that when one computes the holomorphic n-form
        Omega = 1/(∂f/∂z_i) * (dz_1 ^ ... dz_{i-1} ^ dz_{i+1} ^ ... dz_N),
        it is less likely to have a small number on the denominator.
        The corresponding symbolic coordinate is 
        self.affine_coordinate[self.max_grad_coordinate].
    patches: A list of instances of the subpatches of the hypersurface/patch
    points: The numerical points generated by Monte Carlo.
    n_points: The total number of points on the hypersurface or a patch. 
    grad: A list of symbolic expressions of ∂f/∂z_i       

    Usage:
    ----------
    Firstly, one needs to define the coordinates and function with sympy,
    and the number of pairs of points used in Monte Carlo:

    z0, z1, z2, z3, z4 = sp.symbols('z0, z1, z2, z3, z4')    
    Z = [z0, z2, z3, z3, z4]
    f = z0**5 + z1**5 + z2**5 + z3**5 + z4**5 + 0.5*z0*z1*z2*z3*z4 
    n_pairs = 1000

    Then define the hypersurface with f = 0:

    HS = MLGeometry.hypersurface.Hypersurface(Z, f, npairs)  

    """
    def __init__(self, 
                 coordinates, 
                 function,
                 n_pairs=0, 
                 points=None, 
                 norm_coordinate=None,
                 max_grad_coordinate=None):
        """Initialize the hypersurface
       
        Given the sympy coordinates Z, function f and npairs, there are two
        main steps on the highest level:  
        1. Generate points using Monte Carlo methods with __solve_points()
        2. Define the patches automatically with __autopatch()
        On the patches, the points are calculated with __autopatch() 
        beforehand and passed as an argument. 

        """
        self.coordinates = np.array(coordinates)
        self.function = function
        self.n_dim = len(self.coordinates)
        self.norm_coordinate = norm_coordinate
        if norm_coordinate is not None:
            self.affine_coordinates = np.delete(self.coordinates, norm_coordinate)
        else:
            self.affine_coordinates = self.coordinates
        self.max_grad_coordinate = max_grad_coordinate
        self.patches = []
        if points is None:
            self.points = self.__solve_points(n_pairs)
            self.__autopatch()
        else:
            self.points = points
        #self.n_patches = len(self.patches)
        self.n_points = len(self.points)
        self.grad = self.__get_grad()

    def __solve_points(self, n_pairs):
        """Generates random points on the hypersurface with Monte Carlo

        #TODO explain the MC method or refer to the paper

        Args:
        --------
        n_pairs: The number of random pair used in Monte carlo.
  
        Returns:
        --------
        A list of random complex points on the hypersurface, where the points
        themselves are a list of complex coordiantes with dtype Complex128.

        """ 
        points = []
        zpairs = self.__generate_random_pair(n_pairs)
        coeff_a = [sp.symbols('a'+str(i)) for i in range(self.n_dim)]
        coeff_b = [sp.symbols('b'+str(i)) for i in range(self.n_dim)]
        c = sp.symbols('c')
        coeff_zip = zip(coeff_a, coeff_b)
        line = [c*a+b for (a, b) in coeff_zip]
        function_eval = self.function.subs([(self.coordinates[i], line[i])
                                            for i in range(self.n_dim)])
        poly = sp.Poly(function_eval, c)
        coeff_poly = poly.coeffs()
        get_coeff = sp.lambdify([coeff_a, coeff_b], coeff_poly)
        # Multiprocessing. Then append the points to the same list in the main process
        with Pool() as pool:
            for points_d in pool.starmap(Hypersurface.solve_poly,
                                         zip(zpairs, [get_coeff(zpair[0], zpair[1])
                                                      for zpair in zpairs])):
                points.extend(points_d)
        return points

    def __generate_random_pair(self, n_pairs):
        z_random_pair = []
        for i in range(n_pairs):
            zv = []
            for j in range(2):
                zv.append([complex(c[0],c[1]) for c in np.random.normal(0.0, 1.0, (self.n_dim, 2))])
            z_random_pair.append(zv)
        return z_random_pair

    def __generate_CPN(self, n_points):
        z_random = []
        for i in range(n_points):
            z_random.append([complex(c[0],c[1]) for c in np.random.normal(0.0, 1.0, (self.n_dim, 2))])
        return z_random

    @staticmethod
    def solve_poly(zpair, coeff):
        # For each zpair there are d solutions, where d is the n_dim
        points_d = []
        c_solved = mpmath.polyroots(coeff) 
        for pram_c in c_solved:
            points_d.append([complex(pram_c * a + b)
                             for (a, b) in zip(zpair[0], zpair[1])])
        return points_d
    
    def __autopatch(self):
        # projective patches
        points_on_patch = [[] for i in range(self.n_dim)]
        for point in self.points:
            norms = np.absolute(point)
            for i in range(self.n_dim):
                if norms[i] == max(norms):
                    point_normalized = self.normalize_point(point, i)
                    points_on_patch[i].append(point_normalized)
                    continue
        for i in range(self.n_dim):
            self.set_patch(points_on_patch[i], i)
        # Subpatches on each patch
        for patch in self.patches:
            points_on_patch = [[] for i in range(self.n_dim-1)]
            grad_eval = sp.lambdify(self.coordinates, patch.grad, 'numpy')
            for point in patch.points:
                grad = grad_eval(*point)
                grad_norm = np.absolute(grad)
                for i in range(self.n_dim-1):
                    if grad_norm[i] == max(grad_norm):
                        points_on_patch[i].append(point)
                        continue
            for i in range(self.n_dim-1):
                patch.set_patch(points_on_patch[i], patch.norm_coordinate,
                                max_grad_coord=i)

    def set_patch(self, points_on_patch, norm_coord=None, max_grad_coord=None):
        new_patch = Hypersurface(self.coordinates, 
                                 self.function, 
                                 points=points_on_patch, 
                                 norm_coordinate=norm_coord,
                                 max_grad_coordinate=max_grad_coord)
        self.patches.append(new_patch)

    def list_patches(self):
        print("Number of Patches:", len(self.patches))
        i = 1
        for patch in self.patches:
            print("Points on patch", i, ":", len(patch.points))
            i = i + 1

    def normalize_point(self, point, norm_coordinate):
        point_normalized = []
        for coordinate in point:
            norm_coefficient = point[norm_coordinate]
            coordinate_normalized = coordinate / norm_coefficient
            point_normalized.append(coordinate_normalized)
        return point_normalized

    def sum_on_patch(self, f, numerical, tensor):
        summation = 0
        if self.patches == []:
            if tensor is True:
                summation = tf.reduce_sum(f(self))
            else:
                points = np.array(self.points)
                if numerical is True:
                    # Currying if numerical is true
                    def f_patch(patch):
                        def f_point(point):
                            return f(patch, point) 
                        return  f_point
                    func = f_patch(self)

                else:
                    func = sp.lambdify([self.coordinates], f(self), "numpy") 

                for point in self.points:
                    value = func(point)
                    summation += value

                    #if np.absolute(value) < 5 and np.absolute(value) > -5:
                    #    summation += value
                    #else:
                    #    print("Possible division of a small number:", value)
        else:
            for patch in self.patches:
                summation += patch.sum_on_patch(f,numerical, tensor)

        return summation

    def integrate(self, f, holomorphic=True, numerical=False, tensor=False):
        # f should be a lambda expression given by the user
        # holomorphic=True means integrating over Omega_Omegabar
        if tensor is True:
            m = lambda patch: patch.Omega_Omegabar_tf / \
                              patch.num_FS_volume_form_tf('identity', k=1)
            weighted_f = lambda patch: f(patch) * m(patch)

        elif numerical is True:

            # m is the mass formula
            def m(patch, point):
                mass = patch.omega_omegabar(point) / \
                       patch.num_FS_volume_form('identity', point, k=1)
                return mass

            def weighted_f(patch, point):
                weighted_f = f(patch, point) * m(patch, point) 
                return weighted_f
        else:
            m = lambda patch: patch.get_omega_omegabar() / \
                              patch.get_FS_volume_form(k=1)
            weighted_f = lambda patch: f(patch) * m(patch)

        if holomorphic is True:
            # Define a new f with an extra argument user_f and immediatly pass f as
            # the default value, so that f can be updated as f(x) * m(x)
            #f = lambda patch, user_f=f: user_f(patch) * m(patch)
            #m = lambda patch: patch.get_omega_omegabar() / \
            #    patch.get_FS_volume_form(k=1)
            summation = self.sum_on_patch(weighted_f, numerical, tensor)
            norm_factor = 1 / self.sum_on_patch(m, numerical, tensor)
        else:
            summation = self.sum_on_patch(f, numerical, tensor)
            norm_factor = 1 / self.n_points

        integration = summation * norm_factor
        return integration

    def get_FS(self):
        FS_metric = self.kahler_metric(np.identity(self.n_dim), k=1)
        return FS_metric

    def __get_grad(self):
        grad = []
        for coord in self.affine_coordinates:
            grad_i = self.function.diff(coord)
            grad.append(grad_i)
        return grad

    def get_hol_n_form(self, coord):
        """

        Return:
        -------
        A or a list of symbolic expressions of the holomorphic n-form 1/(∂f/∂z_i)

        """
        hol_n_form = []
        if coord is not None:
            hol_n_form = 1 / self.grad[coord]
        else:
            for i in range(len(self.affine_coordinates)):
                hol_n_form.append(self.get_hol_n_form(i))
        return hol_n_form

    def get_omega_omegabar(self, lambdify=False):
        omega_omegabar = []
        if self.patches == [] and self.max_grad_coordinate is not None:
            hol_n_form = self.get_hol_n_form(self.max_grad_coordinate)
            omega_omegabar = hol_n_form * sp.conjugate(hol_n_form)
        else:
            for patch in self.patches:
                try:
                    omega_omegabar.append(patch.omega_omegabar)
                except AttributeError:
                    omega_omegabar.append(patch.get_omega_omegabar(lambdify=lambdify))

        if lambdify is True:
            OObar_func = sp.lambdify([self.coordinates], omega_omegabar,'numpy')
            omega_omegabar = lambda point: OObar_func(point).real
        return omega_omegabar

    def get_sections(self, k, lambdify=False):
        sections = []
        t = sp.symbols('t')
        GenSec = sp.prod(1/(1-(t*zz)) for zz in self.coordinates)
        poly = sp.series(GenSec, t, n=k+1).coeff(t**k)
        while poly!=0:
            sections.append(sp.LT(poly))
            poly = poly - sp.LT(poly)
        n_sections = len(sections)
        sections = np.array(sections)
        if lambdify is True:
            sections = sp.lambdify([self.coordinates], sections, 'numpy')
        return sections, n_sections

    def kahler_potential(self, h_matrix=None, k=1):
        sections, n_sec = self.get_sections(k)
        if h_matrix is None:
            h_matrix = np.identity(n_sec)
        # Check if h_matrix is a string
        elif isinstance(h_matrix, str):
            if h_matrix == "identity":
                h_matrix = np.identity(n_sec)
            elif h_matrix == "symbolic":
                h_matrix = sp.MatrixSymbol('H', n_sec, n_sec)
            elif h_matrix == "FS":
                h_matrix = np.diag(sp.Poly(sp.expand(sum(self.coordinates)**k)).coeffs())
        z_H_zbar = np.matmul(sections, np.matmul(h_matrix, sp.conjugate(sections)))
        if self.norm_coordinate is not None:
            z_H_zbar = z_H_zbar.subs(self.coordinates[self.norm_coordinate], 1)
        kahler_potential = sp.log(z_H_zbar)
        return kahler_potential

    def kahler_metric(self, h_matrix=None, k=1, point=None):
        if point is None:
            pot = self.kahler_potential(h_matrix, k)
            metric = []
            #i holomorphc, j anti-hol
            for coord_i in self.affine_coordinates:
                a_holo_der = []
                for coord_j in self.affine_coordinates:
                    a_holo_der.append(diff_conjugate(pot, coord_j))
                metric.append([diff(ah, coord_i) for ah in a_holo_der])
            metric = sp.Matrix(metric)

        return metric

    def get_restriction(self, ignored_coord=None, lambdify=False):
        if ignored_coord is None:
            ignored_coord = self.max_grad_coordinate
        ignored_coordinate = self.affine_coordinates[ignored_coord]
        local_coordinates = sp.Matrix(self.affine_coordinates).subs(ignored_coordinate, self.function)
        affine_coordinates = sp.Matrix(self.affine_coordinates)
        restriction = local_coordinates.jacobian(affine_coordinates).inv()
        restriction.col_del(ignored_coord)
        if lambdify is True:
            restriction = sp.lambdify([self.coordinates], restriction, 'numpy')
        return restriction

    def get_FS_volume_form(self, h_matrix=None, k=1, lambdify=False):
        kahler_metric = self.kahler_metric(h_matrix, k)
        restriction = self.get_restriction()
        FS_volume_form = restriction.T * kahler_metric * restriction.conjugate()
        FS_volume_form = FS_volume_form.det()
        if lambdify is True:
            FS_func = sp.lambdify([self.coordinates], FS_volume_form, 'numpy')
            FS_volume_form = lambda point: FS_func(point).real
        return FS_volume_form

    # Numerical Methods:

    def set_k(self, k):
        self.k = k
        sections, ns = self.get_sections(k, lambdify=False)
        sections_func, ns = self.get_sections(k, lambdify=True)
        self.n_sections = ns
        for patch in self.patches:
            # patch.k = k
            for subpatch in patch.patches:
                # subpatch.k = k
                subpatch.n_sections = ns
                subpatch.sections = sections_func
                jacobian = sp.Matrix(sections).jacobian(subpatch.affine_coordinates)
                subpatch.sections_jacobian = sp.lambdify([subpatch.coordinates],
                                                         jacobian,'numpy')
                subpatch.restriction = subpatch.get_restriction(lambdify=True)
                subpatch.omega_omegabar = subpatch.get_omega_omegabar(lambdify=True)
                subpatch.h_FS = np.diag(sp.Poly(sp.expand(sum(self.coordinates)**k)).coeffs())
                
                # Tensors
                subpatch.s_tf, subpatch.J_tf = subpatch.num_s_J_tf() 
                subpatch.s_tf_1, subpatch.J_tf_1 = subpatch.num_s_J_tf(k=1) 
                subpatch.Omega_Omegabar_tf = subpatch.num_Omega_Omegabar_tf()
                subpatch.r_tf = subpatch.num_restriction_tf()

    #@tf.function
    def num_s_J_tf(self, k=-1):

        s_vec = []
        J_vec = []

        for point in self.points:
            if k == 1:
                # k = 1 will be used in the mass formula during the integration
                s = [point]
                # Delete the correspoding row
                J = np.delete(np.identity(self.n_dim), self.norm_coordinate, 0)
            else:
                s = [self.sections(point)]
                J = self.sections_jacobian(point).T

            s_vec.append(s)
            J_vec.append(J)

        s_tf = tf.constant(np.array(s_vec, dtype=np.complex64))
        J_tf = tf.constant(np.array(J_vec, dtype=np.complex64))
        return s_tf, J_tf

    #@tf.function
    def num_Omega_Omegabar_tf(self):
        Omega_Omegabar = []
        for point in self.points:
            Omega_Omegabar.append(self.omega_omegabar(point))
        Omega_Omegabar = tf.constant(np.array(Omega_Omegabar, dtype=np.float32))
        return Omega_Omegabar

    #@tf.function
    def num_restriction_tf(self):
        r = []
        for point in self.points:
            r.append(self.restriction(point).T)
        r_tf = tf.constant(np.array(r, dtype=np.complex64))  
        return r_tf

    #@tf.function
    def num_kahler_metric_tf(self, h_matrix, k=-1):
        if isinstance(h_matrix, str):
            if h_matrix == 'identity':
                if k == 1:
                    h_matrix = np.identity(self.n_dim, dtype=np.complex64)
                else:
                    h_matrix = np.identity(self.n_sections, dtype=np.complex64)
            elif h_matrix == 'FS':
                h_matrix = self.h_FS.astype(np.complex64)
                #h_matrix = np.array(self.h_FS, dtype=np.complex64)

        #h_tf = tf.constant(h_matrix)
        #if isinstance(h_matrix, np.ndarray):
        h_tf = tf.convert_to_tensor(h_matrix, dtype=tf.complex64)

        if k == 1:
            s_tf = self.s_tf_1
            J_tf = self.J_tf_1
        else:
            s_tf = self.s_tf
            J_tf = self.J_tf    

        #h_tf = tf.cast(h_tf, tf.complex64)

        H_Jdag = tf.matmul(h_tf, J_tf, adjoint_b=True)
        A = tf.matmul(J_tf, H_Jdag)
        b = tf.matmul(s_tf, H_Jdag)
        B = tf.matmul(b, b, adjoint_a=True)
        alpha = tf.matmul(s_tf, tf.matmul(h_tf, s_tf, adjoint_b=True))
        G = A / alpha - B / alpha**2
        #if tf.reduce_min(tf.abs(alpha)) < 0.001:
        #    print('alpha: ', tf.reduce_min(alpha))
        return G

    #@tf.function
    def num_FS_volume_form_tf(self, h_matrix, k=-1):
        kahler_metric = self.num_kahler_metric_tf(h_matrix, k)
        r_tf = self.r_tf
        FS_volume_form = tf.matmul(r_tf, tf.matmul(kahler_metric, r_tf, adjoint_b=True))
        FS_volume_form = tf.linalg.det(FS_volume_form)
        FS_volume_form = tf.math.real(FS_volume_form)
        return FS_volume_form

    #@tf.function
    def num_eta_tf(self, h_matrix):
        FS_volume_form = self.num_FS_volume_form_tf(h_matrix)
        Omega_Omegabar = self.Omega_Omegabar_tf
        eta = FS_volume_form / Omega_Omegabar
        return eta

    def num_kahler_metric(self, h_matrix, point, k=-1):
        if k == 1:
            # k = 1 will be used in the mass formula during the integration
            s = point
            # Delete the correspoding row
            J = np.delete(np.identity(len(s)), self.norm_coordinate, 0)
        else:
            s = self.sections(point)
            J = self.sections_jacobian(point).T
        if isinstance(h_matrix, str):
            if h_matrix == 'identity':
                h_matrix = np.identity(len(s))
            elif h_matrix == 'FS':
                h_matrix = np.array(self.h_FS, dtype=int)

        H_Jdag = np.matmul(h_matrix, np.conj(J).T)
        A = np.matmul(J, H_Jdag)
        # Get the right half of B then reshape to transpose,
        # since b.T is still b if b is a 1d vector
        b = np.matmul(s, H_Jdag).reshape(-1, 1)
        B = np.matmul(np.conj(b), b.T)
        alpha = np.matmul(s, np.matmul(h_matrix, np.conj(s)))
        G = A / alpha - B / alpha**2
        return G

    def num_FS_volume_form(self, h_matrix, point, k=-1):
        kahler_metric = self.num_kahler_metric(h_matrix, point, k)
        r = self.restriction(point)
        FS_volume_form = np.matmul(r.T, np.matmul(kahler_metric, np.conj(r)))
        FS_volume_form = np.matrix(FS_volume_form, dtype=complex)
        FS_volume_form = np.linalg.det(FS_volume_form).real
        return FS_volume_form

    def num_eta(self, h_matrix, point):
        FS_volume_form = self.num_FS_volume_form(h_matrix, point)
        Omega_Omegabar = self.omega_omegabar(point)
        eta = FS_volume_form / Omega_Omegabar
        return eta

def diff_conjugate(expr, coordinate):
    coord_bar = sp.symbols('coord_bar')
    expr_diff = expr.subs(sp.conjugate(coordinate), coord_bar).diff(coord_bar)
    expr_diff = expr_diff.subs(coord_bar, sp.conjugate(coordinate))
    return expr_diff

def diff(expr, coordinate):
    coord_bar = sp.symbols('coord_bar')
    expr_diff = expr.subs(sp.conjugate(coordinate), coord_bar).diff(coordinate)
    expr_diff = expr_diff.subs(coord_bar, sp.conjugate(coordinate))
    return expr_diff

