from numpy import *
from numpy.linalg import *
import copy

def ggmd(A,N):
    """
    [U,T,V] = ggmd(A,N)
    Python implementation of the "K-user Geometric Mean Decomposition".
    version of Idan Livni, February 3, 2012.
    Copyright 2012, Tel Aviv University, Tel Aviv, Israel. 
    
    The function calculates the joint K-GMD of the given matrices.
    such that U[i].T * A[i] * V = T[i],
    where: 
    	U[i] and V have orthonormal.
    	T[i] are triangular with constant diagonals.
    @param A - list of K square matrices of the same dimension nxn to be jointly decomposed.
    @param N - number of duplications performed for each matrix, such that N > n^(K-1).
    @output V - right orthonormal-column matrix of the decomposition.
    @output U - list of size K with the left orthonormal-column matrices of the decomposition.
    @output T - list of size K with the triangular matrices of the decomposition.
    
    The function uses the scipy and numpy packages.
    """

    #find paramaters
    K=len(A)
    n=A[0].shape[0]

    if (N < n**(K-1)):
            raise Exception("N (num duplication)should be at least n^(k-1)")
    
    #full block gmd for the first matrix
    v=gmd(A[0])[2]
    nn=N*n
    
    U=[kron(eye(N),qrp(x*v)[0]) for x in A]
    T=[kron(eye(N),qrp(x*v)[1]) for x in A]
    V=kron(eye(N),v)
    
    #gmd from the second to the last matrix 
    for i in xrange (1,K):    
        delta=n**(K-i)-1

        #first step - permutatation
        place=[]
        j=0
        while (j*n+n-1+(n-1)*delta)<(nn):
            place += range (j*n+(n-1), j*n+(n-1)+(n-1)*delta+1,delta)
            j+=1
        v=matrix(eye(nn))[:,place]
        T=[v.T*x*v for x in T]

        V=V*v
        U=[U[m]*v for m in xrange(K)]

        #second step - GMD block wise

        nn=T[0].shape[0]
        v=gmd(T[i][0:n,0:n])[2]
        V=V*kron(eye(nn/n),v)

        t = [qrp(x[0:n,0:n]*v)[1] for x in T]
        u=[qrp(x[0:n,0:n]*v)[0] for x in T]

        T=[kron(eye(nn/n),x) for x in t]
        U=[U[m]*kron(eye(nn/n),u[m]) for m in xrange (K)]
    
    T = [ U[i].T * kron(eye(N), A[i]) * V for i in xrange(K) ]
    return U,T,V



def gmd(A):
    """
    A = Q*R*P'
    R is upper-triangular with constant diagonal.

    Based upon Matlab implementation of the "Geometric Mean Decomposition" of
    Hager, December 3, 2003.
    """

    n = len(A)
    sigma_bar = det(A)**(1.0/n)

    good = 0

    Q = matrix(eye(n))
    R = copy.deepcopy(A)
    P = matrix(eye(n))

    #return (Q,R,P)
    for good in xrange(0,n-1):    
        # fix column number "good".
        #
        # start with svd:
        
        tmp = R[good:,good:]
        (a,b,c) = svd(tmp)
        c = c.T
        bla = matrix(eye(n))
        for i in xrange(good,n):
            for j in xrange(good,n):
                bla[i,j] = c[i-good,j-good]
        P = P*bla
        (Q,R) = qrp(A*P)        

        # R should now be diagonal.

        
        for k in xrange(good+1,n):
            if R[good,good]<sigma_bar and R[k,k]>sigma_bar:
                break
            if R[good,good]>sigma_bar and R[k,k]<sigma_bar:
                break
        else:
            continue

        # replace good and k        

        a = R[good,good]**2
        b = R[k,k]**2        
        c2 = (sigma_bar**2 - b)/(a-b)
        if (c2<0) or (c2>1):
            raise Exception("not working...")

        c = c2**0.5
        s = (1-c2)**0.5
        

        tmp = matrix(eye(n))
        tmp[good,good] = c
        tmp[good,k] = s
        tmp[k,good] = s
        tmp[k,k] = -c

        P = P*tmp
        (Q,R) = qrp(A*P)
        
            

    return (Q,R,P)


def qrp(A):
    (q,r) = qr(A)
    n = len(r)
    for i in xrange(len(diag(r))):        
        T = matrix(eye(n))
        T[i,i] = conj(r[i,i])/abs(r[i,i])
        q = q*T
        r = T*r
    return (q,r)