import scipy
from scipy import linalg
import numpy

# la,v = linalg.eigh( A)
# v0 = v[:,0] # first eigenvector

def partn( rls, eps):
    dc={rls[0]:[0]}
    keys=[rls[0]]
    for i in [1..len(rls)-1]:
        x = rls[i]
        key_exists = False
        for j in [0..len(keys)-1]:
            if abs( x-keys[j] ) < eps:
                key_exists = True
                dc[keys[j]] += [i]
                break
        if not key_exists:
            dc[x] = [i]
            keys += [x]
    return dc 
   
def spcdec( A):
    eps = 1e-6
    n = A.nrows()
    la, v = linalg.eigh(A)
    dc = partn( la, eps)
    evals = dc.keys()
    Es = []
    vmat = Matrix(v)
    for ev in dc.keys():
        E  = zero_matrix(RDF,n,n)
        for ind in dc[ev]:
            E += vmat[:,ind]*vmat[:,ind].T
        Es += [(ev, E)]
    return Es 
    
# use M.apply_map(lambda x: round(x,6)) to zero out entries of M that are nearly zero

# compute entropy of a probability density given as a non-negative vector whose
# entries sum to 1  

def entropy( pv):
    return -sum([ vv* log(vv) for vv in pv if vv != 0])

# M had better be stochastic!
def matrix_entropy( M):
    return sum([entropy(vv) for vv in M])

# this returns a function whose value at t is exp(I*t*A)
def transmat( sp):
    return lambda t: sum([numpy.exp(I*t*it[0])*it[1] for it in sp])

# this returns a function whose value at t is the a-row/col of exp(I*t*A) - don't need it though!
def transmat_col( sp, a):
    return lambda t: sum([numpy.exp(I*t*it[0])*it[1][a] for it in sp])

# this returns a function whose value at t is the ab-entry of exp(I*t*A)
def transmat_entry( sp, a, b):
    return lambda t: sum([numpy.exp(I*t*it[0])*it[1][a][b] for it in sp])

# get function whose value at time t is U(t) \circ U(-t)
def stateprob( sp):
    return lambda t: sum([numpy.exp(I*t*it[0])*it[1] for it in sp]).apply_map(norm)

# this returns a function whose value at t is the i-th row/column? of the
# Schur product of exp(I*t*A) and exp(-I*t*A)
def stateprob_row( sp, a):
    return lambda t: map( norm, sum([numpy.exp(I*t*it[0])*it[1][a] for it in sp]))  

def stateprob_entry( sp, a,b):
    return lambda t: norm( sum([numpy.exp(I*t*it[0])*it[1][a][b] for it in sp]))

# entropy of a row, normalized to a max of 1
def entropy_row( sp, a):
    nn = sp[0][1].nrows()     
    return lambda t: entropy( map( norm, sum([numpy.exp(I*t*it[0])*it[1][a] for it in sp])))/log(float(nn))

# returns the sum of the Schur squares of the idempotents  
def avg( sp):
    return sum([it[1].elementwise_product(it[1]) for it in sp])
    
def abs_sum( sp):
    return sum([it[1].apply_map(lambda x: abs(x)) for it in sp])
  
# compute \sum_r E_rDE_r for density matrix D  
def av_state(density, sp):
    return sum([it[1]*density*it[1] for it in sp])
    
# average state of a vertex - D = e_ae_a^T
def av_vstate(vx, sp):
    n = len(sp)
    vc = vector(n*[0])
    vc[ vx] = 1
    return sum(it[1]*vc.outer_product(it[1]*vc) for it in sp)
    
def av_vxstate(vx, sp):
    return sum( it[1][vx].outer_product(it[1][vx]) for it in sp)
    
def mixmat( grf):
    return avg(spcdec(grf.am()))

def parplot(sp,a,b,t0,t1, pltpts=None):
    if pltpts==None: pltpts = max(1000,10*t1)
    t=var('t')
    return parametric_plot( (transmat_entry(sp,a,b)(t).real(), transmat_entry(sp,a,b)(t).imag()), (t,t0,t1),\
        plot_points=pltpts)

# Von Neumann entropy - doesn't work well on the empty graph        
def vne( grf):
    e = grf.num_edges()
    L = (1/(2*e))*grf.laplacian_matrix()
    evs = L.eigenvalues()
    return sum([-it*log_b(it,2) for it in evs if it != 0])
	
def ev_supp(grf, vx):
    gp = grf.charpoly()
    h = grf.copy()
    h.delete_vertex(vx)
    return (gp/gp.gcd(h.charpoly())).factor()