import numpy as np
from multi_power import MultiPower
from multi_cheb import MultiCheb
import itertools
from groebner_class import Groebner

'''
This module contains the tools necessary to find the points of the variety of the
ideal generated by a Groebner basis.
'''

permitted_round_error = 1.e-10

def roots(polys):
    '''
    Finds the roots of the given list of polynomials

    parameters
    ----------
    polys : list of polynomial objects
        polynomials to find the common roots of

    returns
    -------
    list of numpy arrays
        the common roots of the polynomials
    '''
    # Determine polynomial type
    poly_type = ''
    if (all(type(p) == MultiCheb for p in polys)):
        poly_type = 'MultiCheb'
    elif (all(type(p) == MultiPower for p in polys)):
        poly_type = 'MultiPower'
    else:
        raise ValueError('All polynomials must be the same type')

    # Calculate groebner basis
    G = Groebner(polys)
    GB = G.solve()
    dim = max(g.dim for g in GB)

    # Check for no solutions
    if len(GB) == 1 and GB[0].coeff == 1:
        print("No solutions")
        return -1

    # Make sure ideal is zero-dimensional
    var_list = _get_var_list(dim)
    if not _test_zero_dimensional(var_list, GB):
        print("Ideal is not zero-dimensional; cannot calculate roots.")
        return -1

    # Pick a random polynomial of the same type
    # and in the same number of variables as the Groebner basis
    f = _random_poly(poly_type, var_list)

    # Get multiplication matrix
    VB, var_dict = vectorSpaceBasis(GB)
    print("VB:", VB)
    print("var_dict:", var_dict)
    m_f = multMatrix(f, GB, VB)

    # Get list of indexes of single variables and store vars that were not
    # in the vector space basis.
    var_indexes = np.array([-1 for i in range(dim)])
    vars_not_in_basis = {}
    for i in range(len(var_list)):
        var = var_list[i]
        if var in var_dict:
            # note this assures x_i will be in the ith position in var_indexes
            var_indexes[i] = var_dict[var]
        else:
            # maps the position in the root to its variable
            vars_not_in_basis[i] = var

    vnib = False
    if len(vars_not_in_basis) != 0:
        vnib = True

    # Get left eigenvectors
    eig = np.linalg.eig(m_f.T)[1]
    num_vectors = eig.shape[1]
    eig_vectors = [eig[:,i].tolist() for i in range(num_vectors)] # columns of eig

    roots = []
    for v in eig_vectors:
        root = np.zeros(dim, dtype=complex)
        # This will always work because var_indexes and root have the
        # same length - dim - and var_indexes has the variables in the
        # order they should be in the root
        for i in range(dim):
            x = var_indexes[i]
            if x != -1:
                root[i] = v[x]/v[0]
        if vnib:
            print("cur_root:", root)
            for pos in list(vars_not_in_basis.keys())[::-1]:
                GB_poly = _get_poly_with_LT(vars_not_in_basis[pos], GB)
                print("GB_poly:\n", GB_poly.coeff)
                var_value = GB_poly.evaluate_at(root) * -1
                print("var_value:", var_value)
                root[pos] = var_value
        roots.append(root)

    return roots

def multMatrix(poly, GB, basis):
    '''
    Finds the matrix of the linear operator m_f on A = C[x_1,...,x_n]/I
    where f is the polynomial argument. The linear operator m_f is defined
    as m_f([g]) = [f]*[g] where [f] represents the coset of f in
    A. Since m_f is a linear operator on A, it can be represented by its
    matrix with respect to the vector space basis.

    parameters
    ----------
    poly : polynomial object
        The polynomial f for which to find the matrix m_f.
    GB: list of polynomial objects
        Polynomials that make up a Groebner basis for the ideal
    basis : list of tuples
        The monomials that make up a basis for the vector space A

    return
    ------
    multOperatorMatrix : square numpy array
        The matrix m_f
    '''

    # Reshape poly's coefficienet matrix if it is not in the same number
    # of variables as the polynomials in the Groebner basis.
    # (i.e. len(shape) is the number of variables the polynomial is in)

    numVars = max(g.dim for g in GB)

    polyVars = len(poly.coeff.shape)
    if polyVars != numVars:
        new_shape = [i for i in poly.coeff.shape]
        for j in range(numVars-polyVars): new_shape.append(1)
        if type(poly) is MultiPower:
            poly = MultiPower(poly.coeff.reshape(tuple(new_shape)))
        if type(poly) is MultiCheb:
            poly = MultiCheb(poly.coeff.reshape(tuple(new_shape)))

    dim = len(basis)
    operatorMatrix = np.zeros((dim, dim))

    for i in range(dim):
        monomial = basis[i]
        poly_ = poly.mon_mult(monomial)

        operatorMatrix[:,i] = coordinateVector(poly_, GB, basis)

    return operatorMatrix

def vectorSpaceBasis(GB):
    '''
    parameters
    ----------
    GB: list
        polynomial objects that make up a Groebner basis for the ideal

    return
    ------
    basis : list
        tuples representing the monomials in the vector space basis
    var_to_pos_dict : dictionary
        maps each variable to its position in the vector space basis
    '''
    LT_G = [f.lead_term for f in GB]
    possibleVarDegrees = [range(max(tup)) for tup in zip(*LT_G)]
    possibleMonomials = itertools.product(*possibleVarDegrees)
    basis = []
    var_to_pos_dict = {}

    for mon in possibleMonomials:
        divisible = False
        for LT in LT_G:
            if (divides(LT, mon)):
                 divisible = True
                 break
        if (not divisible):
            basis.append(mon)
            if (sum(mon) == 1):
                var_to_pos_dict[mon] = basis.index(mon)

    return basis, var_to_pos_dict

def coordinateVector(poly, GB, basis):
    '''
    parameters
    ----------
    reducedPoly : polynomial object
        The polynomial for which to find the coordinate vector of its coset.
    GB : list of polynomial objects
        Polynomials that make up a Groebner basis for the ideal
    basis : list of tuples
        The monomials that make up a basis for the vector space

    return
    ------
    coordinateVector : list
        The coordinate vector of the given polynomial's coset in
        A = C[x_1,...x_n]/I as a vector space over C
    '''
    print("orig_poly:\n", poly.coeff)
    vectorSpaceDim = len(basis)
    poly = reduce_poly(poly, GB)
    print("reduced_poly:\n", poly.coeff)

    # reverse the array since self.vectorBasis is in increasing order
    # and monomialList() gives a list in decreasing order
    poly_terms = poly.monomialList()[::-1]
    assert(len(poly_terms) <= vectorSpaceDim)

    coordinateVector = [0] * vectorSpaceDim
    for monomial in poly_terms:
        coordinateVector[basis.index(monomial)] = \
            poly.coeff[monomial]

    return coordinateVector

def divides(mon1, mon2):
    '''
    parameters
    ----------
    mon1 : tuple
        contains the exponents of the monomial divisor
    mon2 : tuple
        contains the exponents of the monomial dividend

    return
    ------
    boolean
        true if mon1 divides mon2, false otherwise
    '''
    return all(np.subtract(mon2, mon1) >= 0)

def reduce_poly(poly, divisors):
    '''
    Divides a polynomial by the Groebner basis using the standard
    multivariate division algorithm and returns the remainder

    parameters
    ----------
    polynomial : polynomial object
        the polynomial to be divided by the Groebner basis
    divisors : list of polynomial objects
        Polynomials to divid poly by

    return
    ------
    polynomial object
        the unique remainder of poly divided by self.GB
    '''

    remainder_coeff = np.zeros_like(poly.coeff, dtype=float)

    # while poly is not the zero polynomial
    while np.any(poly.coeff):
        divisible = False
        # Go through polynomials in set of divisors
        for divisor in divisors:
            # If the LT of the divisor divides the LT of poly
            if divides(divisor.lead_term, poly.lead_term):
                # Get the quotient LT(poly)/LT(divisor)
                LT_quotient = tuple(np.subtract(
                    poly.lead_term, divisor.lead_term))

                new = divisor.mon_mult(LT_quotient)

                # Get max value of shapes to know how much to pad
                max_shape = np.maximum(poly.coeff.shape, new.coeff.shape)

                poly_pad = np.subtract(max_shape, poly.coeff.shape)
                poly.__init__(_pad_matrix(poly_pad, poly.coeff), clean_zeros=False)

                new_pad = np.subtract(max_shape, new.coeff.shape)
                new.__init__(_pad_matrix(new_pad, new.coeff), clean_zeros=False)

                new_coeff = poly.coeff - \
                    (poly.lead_coeff/divisor.lead_coeff)*new.coeff
                new_coeff[np.where(abs(new_coeff) < permitted_round_error)]=0
                poly.__init__(new_coeff, clean_zeros=False)

                divisible = True
                break

        if not divisible:
            lcm = np.maximum(poly.coeff.shape, remainder_coeff.shape)
            remainder_pad = np.subtract(lcm, remainder_coeff.shape)
            remainder_coeff = \
                _pad_matrix(remainder_pad, remainder_coeff)

            # Add lead term to remainder
            polyLT = poly.lead_term
            remainder_coeff[polyLT] = poly.coeff[polyLT]

            # Subtract LT from poly
            new_coeff = poly.coeff
            new_coeff[poly.lead_term] = 0
            poly.__init__(new_coeff)

    if (type(poly) == MultiPower):
        return MultiPower(remainder_coeff)
    else:
        return MultiCheb(remainder_coeff)

def _pad_matrix(pad, matrix):
    _list = []
    for i in pad:
        _tuple = (0,i)
        _list.append(_tuple)
    return np.pad(matrix, _list, 'constant', constant_values=0)

def _get_var_list(dim):
    _vars = [] # list of the variables: [x_1, x_2, ..., x_n]
    for i in range(dim):
        var = np.zeros(dim, dtype=int)
        var[i] = 1
        _vars.append(tuple(var))
    return _vars

def _random_poly(_type, _vars):
    '''
    Generates a random polynomial that has the form
    c_1x_1 + c_2x_2 + ... + c_nx_n where n = len(_vars) and each c_i is a randomly
    chosen integer between 0 and 1000.

    _vars should be the result of the _get_var_list(dim) method where dim is
    the number of variables in which to create the polynomial.
    '''
    dim = len(_vars)

    random_poly_shape = [2 for i in range(dim)]

    random_poly_coeff = np.zeros(tuple(random_poly_shape), dtype=int)
    for var in _vars:
        random_poly_coeff[var] = np.random.randint(1000)

    if _type == 'MultiCheb':
        return MultiCheb(random_poly_coeff)
    else:
        return MultiPower(random_poly_coeff)

def _get_poly_with_LT(LT, GB):
    for poly in GB:
        if poly.lead_term == LT:
            return poly

def _test_zero_dimensional(_vars, GB):
    LT_list = [p.lead_term for p in GB]

    for var in _vars:
        exists_multiple = False
        for LT in LT_list:
            if np.linalg.matrix_rank(np.array([list(var), list(LT)])) == 1:
                exists_multiple = True
                break
        if not exists_multiple:
            return False

    return True
