import numpy as np
import fastcluster


def dist(R):
    N = R.shape[0]
    d = R[np.triu_indices(N,1)]
    out = fastcluster.average(d).astype(int)

    #Genealogy Set
    dend = {i: (np.array([i]),) for i in range(N)}
    for i,(a,b,_,_) in enumerate(out):
        dend[i+N] = (np.concatenate(dend[a]),np.concatenate(dend[b]))

    
    [dend.pop(i,None) for i in range(N)]

    return dend.values()

def AvLinkC(Dend,R):

    N = R.shape[0]
    Rs = np.zeros((N,N))

    for (a,b) in Dend:
        Rs[np.ix_(a,b)] = R[a][:,b].mean()

    Rs = Rs+Rs.T

    np.fill_diagonal(Rs,1)
    return Rs	
    
def noise(N,T,epsilon=1e-10):
    return np.random.normal(0,epsilon ,size=(N,T))

def filterCorrelation(x,Nboot=100):

    N,T = x.shape
    rT = range(T)
    ns = noise( N,T ) 

    Cbav = np.zeros((N,N))
    for _ in range(Nboot):
        xboot = x[:,np.random.choice(rT,size=T,replace=True)] + ns
        Cboot  = np.corrcoef(xboot)
        Dend = dist(1-Cboot)
        Cbav += AvLinkC(Dend,Cboot)

    Cbav = Cbav/Nboot
    return Cbav

def filterCovariance(x,Nboot=100):

    N,T = x.shape
    rT = range(T)
    ns = noise( N,T ) 
	

    Sbav = np.zeros((N,N))
    for _ in range(Nboot):
        xboot = x[:,np.random.choice(rT,size=T,replace=True)] + ns
        Cboot  = np.corrcoef(xboot)
        Dend = dist(1-Cboot)
        std = xboot.std(axis=1)
        Sbav += AvLinkC(Dend,Cboot)*np.outer(std,std)

    Sbav = Sbav/Nboot
    return Sbav


def filterCovarianceArmonic(x,Nboot=100):

    N,T = x.shape
    rT = range(T)
	ns = noise( N,T ) 
	
    Sbav = np.zeros((N,N))
    for _ in range(Nboot):
        xboot = x[:,np.random.choice(rT,size=T,replace=True)] + ns
        Cboot  = np.corrcoef(xboot)
        Dend = dist(1-Cboot)
        std = xboot.std(axis=1)
        Sbav += np.linalg.pinv(AvLinkC(Dend,Cboot)*np.outer(std,std))

    Sbav = Sbav/Nboot
    return np.linalg.pinv(Sbav)

def HigherOrder(C,Norder):
    Cf = np.zeros(C.shape)
    for _ in range(Norder):
        res = C - Cf
        dend = bahc.main.dist(1-res)
        res = bahc.main.AvLinkC(dend,res)
        np.fill_diagonal(res,0)
        Cf += res

    np.fill_diagonal(Cf,1)
    return Cf

def filterCovarianceHO(x,Norder=3,Nboot=100):
    N,T = x.shape
    rT = range(T)

    #bootstrap
    Cs,std = [],[]
    for _ in range(Nboot):

        xboot = x[:,np.random.choice(rT,size=T,replace=True)]
        C = np.corrcoef(xboot)
        std.append(xboot.std(axis=1))

        #higher-order approximation
        Cs.append(HigherOrder(C,Norder))

    #Eigenvalue-Cleaning
    V = np.linalg.eigh(Cs)[1]

    Ls = []
    for v in V:
        L = np.array([np.dot(np.dot(v.T,Cf),v).diagonal() for Cf in Cs])
        L[L<0] = np.nan
        L[:,np.isnan(L).all(axis=0)] = 0

        Ls.append(np.nanmean(L,axis=0))

    #Average Matrix
    return sum(np.dot(np.dot(v,np.diag(l)),v.T)*np.outer(s,s) for v,l,s in zip(V,Ls,std)) 
