import numpy as np
import itertools
import warnings
from CHEBYSHEV.TVB_Method.cheb_class import MultiCheb, Term, TVBError, match_size, match_poly_dimensions, get_var_list
from CHEBYSHEV.TVB_Method.TVB import telen_van_barel

'''This module contains methods for finding the zero locus of a
   0-dimensional ideal defined by a list of Chebyshev polynomials.
   Roots are found using a variant of the Telen-vanBarel method.
   Speed and accuracy are good for real roots whose coordinates have
   norm 1 or less, but ill-conditioning makes them less accurate for
   roots with norm > 1.


Methods: 

    roots(polys) : Find common roots of 'polys'.

    sortVB(VB) : Sort a vector-space basis to improve speed of eigensolving.

    TVBMultMatrix(polys,f) : Find the matrix rep of the multiply-by-f
                             operator in the TvB basis
        
    _random_poly(dim) : Construct a random linear Chebyshev polynomial for TVBMultMatrix

    newton_polish(polys,roots) : Use Newton's method to polish a collection of 
                                 approximate roots.

    check_zeros(zeros, polys): Check whether 'zeros' are, indeed, all the zeros of 
                               'polys', and how many are outside the sup-norm unit 
                                ball |z|_\infty < 1 or are not real.

       
'''


def roots(polys):
    '''Find the common roots of the given list of polynomials.

        Uses a variant of the Telen-vanBarel method for reducing the
        Macaulay matrix to find a basis of the quotient ring
        C[x_1,...,x_n]/I, where I is the ideal generated by 'polys'.
        Given the basis, finds the roots by finding the left
        eigenvector of the multiply-by-f matrix, where f is a random
        (or pre-chosen) linear polynomial.

    Parameters
    ----------
    polys : list of polynomial objects
        Polynomials to find the common roots of.

    returns
    -------
    list of numpy arrays
        the common roots of the polynomials

    '''
    
    polys = match_poly_dimensions(polys) #polynomials might not all be the same dimension, initially.

    m_f, var_dict = TVBMultMatrix(polys) # m_f is the operator whose eigenvectors give the roots
                                         # var_dict gives the location of each variable in the
                                         #   basis (this is needed for identifying the root in the
                                         #   eigenvector) 

    # Get list of indices of single variables and store vars that were not
    # in the vector space basis (if x_i is not in the basis, this is indicated
    # by var_indices[i] = -1).

    dim = max(f.dim for f in polys) # Number of variables in C[x_1,...,x_n]
    var_list = get_var_list(dim)    # Variables x_1,...,x_n, given as tuples: x_1 is (1,0,0,...,0) etc
    var_indices = [-1]*dim          # Make a list of the indices
    for i in range(dim):
        var = var_list[i] # x_i
        var_indices[i] = var_dict[var] #location of x_i in the TvB basis

        
    # Get left eigenvectors of m_f

    eig = np.linalg.eig(m_f.T)[1]  # This is one of the two most expensive steps in this algorithm.
    num_vectors = eig.shape[1]

    eig_vectors = [eig[:,i] for i in range(num_vectors)] # columns of eig
    roots = []
    for v in eig_vectors:
        if v[var_dict[tuple(0 for i in range(dim))]] == 0:
            continue
        root = np.zeros(dim, dtype=complex)
        # This will always work because var_indices and root have the
        # same length - dim - and var_indices has the variables in the
        # order they should be in the root
        for i in range(dim):
            x_i_pos = var_indices[i]
            if x_i_pos != -1:
                root[i] = v[x_i_pos]/v[var_dict[tuple(0 for i in range(dim))]]
        roots.append(root)
        #roots.append(newton_polish(polys,root,niter=1000,tol=1e-10))
    return roots

def sortVB(VB):
    '''
    Sorts the Vector Basis into degrevlex order so the eigensolve is faster (in theory).

    Parameters
    ----------
    VB : numpy array
        Each row in VB is a term in the vector basis.

    Returns
    -------
    VB : numpy array
        The vector basis sorted so the lowest terms are at the top.
    '''

    VBList = list()
    for i in VB:
        VBList.append(Term(i))

    return VB[np.argsort(VBList)]

def TVBMultMatrix(polys,f=None):
    
    '''Find the multiplication-by-f matrix on C[x_1,...,x_n]/I, where I is
       the idea generated by 'polys', expressed in the TvB basis.

    Parameters
    ----------
    polys : array-like
        The polynomials to find the common zeros of
    
    f : MultiCheb polynomial (linear)
        The linear polynomial defining the multiplication-by-f operator.  
        If None, f is chosen randomly.

    Returns
    -------
    multiplicationMatrix : 2D numpy array
        The matrix representation of the multiplication-by-f operator on the TvB basis.

    var_dict : dictionary
        Maps each variable to its position in the TvB vector space basis.

    '''
    
    basisDict, VB, degree = telen_van_barel(polys)

    VB = sortVB(VB)

    dim = max(g.dim for g in polys)

    if not f:
        # Get random linear polynomial f
        f = _random_poly(dim)[0]

    slices = list()
    for i in range(len(VB[0])):
        slices.append(VB.T[i])

    VBset = {tuple(mon) for mon in VB}

    # Build multiplication-by-f matrix 'mMatrix'
    mMatrix = np.zeros((len(VB), len(VB)))
    remainder = np.zeros([degree]*dim)

    for i in range(VB.shape[0]):
        f_coeff = f.mon_mult(VB[i], return_type = 'Matrix')
        for term in zip(*np.where(f_coeff != 0)):
            if term in VBset:
                remainder[term] += f_coeff[term]
            else:
                remainder[slices] -= f_coeff[term]*basisDict[term][slices]
        mMatrix[:,i] = remainder[slices]
        remainder[slices] = 0

    # Construct var_dict
    var_dict = {}
    for i,mon in enumerate(VB):
        if np.sum(mon) == 1 or np.sum(mon) == 0:
            var_dict[tuple(mon)] = i

    return mMatrix, var_dict


def _random_poly(dim):
    '''Generate a random linear (Chebyshev) polynomial of the form 
    c_1 x_1 + c_2 x_2 + ... + c_n x_n, where n = dim and each c_i is a
    randomly chosen integer between 0 and 1000.

    Parameters
    ----------
    dim : int
        Number of variables

    Returns
    -------
    polynomial
        Randomly generated Chebyshev polynomial of degree 1 and dimension 'dim'.

    '''

    _vars = get_var_list(dim)  #list of tuples corresponding to the variables 

    random_poly_coeff = np.zeros([2]*dim, dtype=int)
    for var in _vars:
        random_poly_coeff[var] = np.random.randint(1000)
        
    return MultiCheb(random_poly_coeff), _vars



########## Newton ##########################################

def newton_polish(polys,root,niter=100,tol=1e-8):
    """
    Perform Newton's method on a system of N polynomials in M variables.

    Parameters
    ----------
    polys : list
        A list of polynomial objects of the same type (MultiPower or MultiCheb).
    root : ndarray
        An initial guess for Newton's method, intended to be a candidate root from root_finder.
    niter : int
        A maximum number of iterations of Newton's method.
    tol : float
        Tolerance for convergence of Newton's method.

    Returns
    -------
    x1 : ndarray
        The terminal point of Newton's method, an estimation for a root of the system
    """

    m = len(polys)
    dim = max(poly.dim for poly in polys)
    f_x = np.empty(m,dtype="complex_")
    jac = np.empty((m,dim),dtype="complex_")
    
    def f(x):
        #f_x = np.empty(m,dtype="complex_")
        for i, poly in enumerate(polys):
            f_x[i] = poly(x)
        return f_x

    def Df(x):
        #jac = np.empty((m,dim),dtype="complex_")
        for i, poly in enumerate(polys):
            jac[i] = poly.grad(x)
        return jac

    i = 0
    x0, x1 = root, root
    while True:
        if i == niter:
            break
        delta = np.linalg.solve(Df(x0),-f(x0))
        x1 = delta + x0
        if np.linalg.norm(delta) < tol:
            break
        x0 = x1
        i+=1
    return x1




#########  Testing Funtions   ##########

def check_zeros(zeros, polys, real=True, tol=1e-5):
    """
    Check whether 'zeros' are, indeed, all the zeros of 'polys', and how many are 
          outside the unit ball |z|_\infty < 1.

    Parameters
    ----------
    zeros : list
        Supposed roots (usually found using the root finder).
    polys : list
        Polynomials that 'zeros' should be roots of.
    real : bool
        Whether to check that the bad zeros are real (real=True) or 
         not to check (real=False)

    Prints
    ----------
    The number of correct zeroes found with the total number.
    The number of incorrect zeros that were out of range or nonreal.

    Returns
    ----------
    A list of bad zeros.


    """
    correct = 0
    outOfRange = 0
    bad = set()
    bad_inrange = []
    if zeros != -1:
        for zero in zeros:
            good = True
            for poly in polys:
                v = poly(zero)
                if np.abs(v) > tol:
                    good = False
                    bad.add(tuple(zero))
            if good:
                correct += 1


    real_status = ''
    bad_list = [np.array(zero) for zero in bad]
    for zero in bad_list:
        if (np.abs(zero) > 1).any():
            outOfRange += 1
        elif real and np.any(np.abs(np.imag(zero))>tol):
            outOfRange += 1
            real_status = 'not real or'
        else:
            bad_inrange.append(zero)
            
    print("{} zeros are correct to {}, out of {} total zeros.".format(correct, tol,len(zeros)))
    print("{} are bad, but {} of these were {} out of range (expected to be bad).".format(len(bad), outOfRange, real_status))
    print("{} might be lost".format(len(bad_inrange)))

    better = []
    diff = []
    newt_bad = set()
    for zero in bad_inrange:
        newt_zero = newton_polish(polys,zero,tol=1e-10)
        better.append(newt_zero)
        diff.append(zero - newt_zero)
        for poly in polys:
            v = poly(newt_zero)
            if np.abs(v) > tol:
                newt_bad.add(tuple(newt_zero))
        
    print('{} seem to be lost after newton polishing'.format(len(newt_bad)))
    print("Differences between the 'bad' inrange zeros and the polished ones are {}".format(diff)) 
                    
    #return bad_inrange, newt_bad

    

    
