# -*- coding: utf-8 -*-
"""
methods associated with the power iteration Tensor decomposition method
"""

import numpy as np

def readMarks(files):
    """Get the binary histone data, convert binary into decimal,
            delete unimportant columns"""    
    marks = []
    for filename in files:
        file = open(filename)
    
        file.readline()
        file.readline()    
    
        for line in file:
            line = line.replace('\n','')
            line = line.replace('\t','')
            line = line[1:4] + line[5]
            line = int(line,2)
            marks.append(line)
    
    
        file.close()
    return marks
    
#2,3,4,6



def getTriples(marks):
    """get all interlacing triples"""
    
    triples = []
    i = 0
    while i + 2 < len(marks):
        trip = [marks[i],marks[i+1],marks[i+2]]
        triples.append(trip)
        i += 1
    
    return triples
    


 
def pairs(triples,i,j,n):
    """get the frequency of all pairs
        
        arguments: triples = 3 combination documents
                   i = position of first element
                   j = position of second element
                   n = number of combinations """
    #initalize pairs matrix               
    pairsFreq = np.zeros([n,n])
    
    #get all pairs from triples
    for triple in triples:

        pairsFreq[triple[i-1],triple[j-1]] += 1
    #get frequency of pairs
    pairsFreq /= len(triples)
            
    return pairsFreq
    

def triples(triples,n):
    """get the frequency of all triples
    
        arguments: triples = 3 combination documents
                   n = number of combinations"""
                   
    triplesFreq = np.zeros([n,n,n])
    #get frequency of triples
    for triple in triples:
        triplesFreq[triple[0],triple[1],triple[2]] += 1
    triplesFreq /= len(triples)
                
    return triplesFreq
    
def x_i(triples,i,n):
    """returns x_i, the fequency of each combination occuring as 
            the first document"""
    x_i = np.zeros(n)
    #get singletons from triples
    for triple in triples:
        x_i[triple[i-1]] += 1

    #get frequency of singletons
    x_i /= len(triples)
    return x_i

def pairsTriples(triples,n):
    """returns the pairs matrices and triples tensor
    
        arguments: triples = 3 combination documents
                   n = number of combinations"""
    p12 = np.zeros([n,n])
    p13 = np.zeros([n,n])
    p23 = np.zeros([n,n])
    p123 = np.zeros([n,n,n])
    
    for triple in triples:
        p12[triple[0],triple[1]] += 1
        p13[triple[0],triple[2]] += 1
        p23[triple[1],triple[2]] += 1
        p123[triple[0],triple[1],triple[2]] += 1
    count = len(triples)
    p12 /= count   
    p13 /= count 
    p23 /= count 
    p123 /= count
    return (p12,p13,p23,p123)

def transform(p12,p13,p23,p123,k):
    """returns m_2 and m_3, the transformed moment tensors on which the
        tensor powr method can be used
        Hsu et al 2014 section 3.3
        arguments: x_i = the ith singleton freqeuency vector
                   pij = the i,j pairs frequency matrix
                   p123 = the triples freqeuency tensor
                   k = number of states"""
                
    d = len(p12[0])
    # choose a an b such that p13 is a * p13 * b transpose is invertible
    u,s,v = np.linalg.svd(p13)
    for i in range(len(u[0])-1,k-1,-1):
        u = np.delete(u,i,1)
        v = np.delete(v,i,0)
    a = np.transpose(u)
    b = v
    
    
       
    
    #creation of m2
    #m2 = np.dot(np.dot(p23,np.linalg.inv(p13)),p12)
    temp = np.linalg.inv(np.dot(np.dot(a,p13),np.transpose(b)))
    temp = np.dot(np.dot(p23,np.transpose(b)),temp)
    temp = np.dot(temp,np.transpose(np.dot(np.transpose(p12),np.transpose(a))))
    m2 = temp
    #m2 = np.dot(np.dot(p21,np.linalg.inv(p31)),p32)
    #create m2 as it says in AFH+12

    #m2 = np.dot(np.dot(p31,np.linalg.inv(p21)),np.transpose(p32))

    #m2prime = np.dot(p31,np.dot(np.linalg.inv(p21),np.transpose(p32)))    
    #create m3
    m3 = np.zeros([d,d,d])
    for i in range(0,d):
        temp = np.dot(np.dot(p23,np.linalg.inv(p13)),a)
        temp = np.dot(temp,np.transpose(b))
        temp = np.dot(np.dot(temp,np.linalg.inv(p13)),p12)
        m3[:,i,:] = temp
    
    return (m2, m3)
   
   

        
    
        
    
    
def whiten(m2,m3,k):
    """whiten m2 and m3 so that the eigenvectors are orthonormal
        Hsu et al 2014 section 4.3        
        
        arguments: m2,m3 
                   k = number of states        
        """
    d,v = np.linalg.eig(m2)
    d = d.astype(float)
    v = v.astype(float)        
    sqrtd = np.diag(d**(-.5))

    for i in range(len(d)-1,k-1,-1):
        sqrtd = np.delete(sqrtd,i,1)
        sqrtd = np.delete(sqrtd,i,0)
        v = np.delete(v,i,1)
    w = np.dot(v,sqrtd)

    m3tilda = np.zeros([k,k,k])
    for i1 in range(0,k):
        for i2 in range(0,k):
            for i3 in range(0,k):
                m3tilda[i1,i2,i3] = tensorMap(m3,w,i1,i2,i3)
    return (m3tilda,w)


def tensorMap(t,w,i1,i2,i3):
    """returns the i1,i2,i3 th element in a tensor map t(w,w,w)
        where w is a matrix,
        defined in Hsu et al section 2        
        """
    s = 0
    n = len(t[0,0])
    for i in range(0,n):
        for j in range(0,n):
            for k in range(0,n):
                s += t[i,j,k] * w[i,i1] * w[j,i2] * w[k,i3]
    return s

def tensorMapVector(t,w):
    """returns the result of the tensor map t(w,w,w) of a vector w,
    defined in Hsu et al section 2"""
    n = len(t[0,0])
    s=0
    for i in range(0,n):
        for j in range(0,n):
            for k in range(0,n):
                s += t[i,j,k] * w[i] * w[j] * w[k]
    return s


def randomUnitSphere(dim):
    """returns a vector sampled uniformly from the unit n-sphere Marsaglia (1972)"""
    x = np.zeros(dim)
    r = 0
    for i in range(0,dim):
        x[i] = np.random.normal(0,1)
        r += x[i]**2
    r = r**(1/2)
    return np.divide(x,r)
    
def operate(t,theta):
    """returns the application of the quadratic operator of a tensor on a vector"""
    n = len(theta)
    s = 0
    for i in range(0,n):
        for j in range(0,n):
            for k in range(0,n):
                e_i = np.zeros(n)            
                e_i[i] = 1
                s += t[i,j,k] * theta[j] * theta[k] * e_i
    return s

def powerIteration(t,l,n):
    """decomposes a tensor using the power method, returns an eigenvector 
        and eigenvalue
        Hsu et al 2014 section 5
        arguments: t = tensor
                   l = number of iterations
                   n = number of power iteration updates
    """
    dim = len(t[0,0])
    #get eigenvector apprximations
    thetas = []
    for tau in range(0,l):
        theta = randomUnitSphere(dim)
        for i in range(0,n):
            theta = operate(t,theta)
            theta = theta/np.linalg.norm(theta)
        thetas.append(theta)
    #get eigenvalue approximations 
    eigenvalues = []
    for theta in thetas:
        eigenvalues.append(tensorMapVector(t,theta))
    #get eigenvector/value pair with highest eigenvalue 
    index = np.argmax(eigenvalues)
    #return eigenvector/value pair and deflated tensor
    v = thetas[index]
    lam = eigenvalues[index]
    deflated = t - lam * tensorPower(v)
    return(v,lam,deflated)

def decompose(t,l,n):
    """uses power iteration to fully decompose a tensor, returns an array of 
        eigenvectors and an array of eigenvalues
        Hsu et al 2014 section 5
        arguments: t = tensor
                   l = number of iterations
                   n = number of power iteration updates
    """
    vectors = []
    values = []
    k = len(t[0,0])
    for i in range(0,k):
        v, lam, t = powerIteration(t,l,n)
        vectors.append(v)
        values.append(lam)
    return (vectors,values)
    

def getParams(m2,m3,l,n,k):
    """returns the eigenvectors and values of the third order tensor m3
        
        aguments: m2 = 2nd moment
                  m3 = 3rd moment
                  l = number of iterations
                  n = number of power iteration updates
                  k = number of states
        """
    m3_tilda, w = whiten(m2,m3,k)
    vectors, values = decompose(m3_tilda,l,n)
    vectors = convertVectors(vectors,values,w)    
    values = convertValues(values)
    return (vectors,values)

def convertValues(values):
    """converts the eigenvalues of the whitened tensor to the
        eigenvalues of the original tensor
        Hsu et al 2014 section 4.3
        """
    eigenvalues = []
    for value in values:
        eigenvalues.append(1/value**(2))
    return eigenvalues

def convertVectors(vectors,values,w):
    """converts the eigenvectors of the whitened tensor to the 
        eigenvectors of the original tensor
        Hsu et al 2014 section 4.3
        """
    
    b = np.linalg.pinv(np.transpose(w))
    eigenvectors = []
    for i in range(0,len(vectors)):
        eigenvector = values[i] * np.dot(b,vectors[i])
        eigenvectors.append(eigenvector)
        
    return eigenvectors

def getOTPi(m2,m3,l,n,p31,p21,p1,k):
    """Get the ommission and transition matrix from the parameters
        Spectacle section 2.2.1
        aguments: m2 = 2nd moment
                  m3 = 3rd moment
                  l = number of iterations
                  n = number of power iteration updates
                  k = numbe of states
                  """
    vectors, values = getParams(m2,m3,l,n,k)
    o = np.transpose(np.array(vectors))
    t = np.dot(np.dot(np.linalg.pinv(o),p31),
               np.linalg.pinv(np.dot(np.linalg.pinv(o),p21)))
    pi = np.dot(np.linalg.pinv(o),p1)
    
    return (o,t,pi)
    
def tensorPower(v):
    """returns the 3rd tensor power of a vector v"""
    n = len(v)
    tensor = np.zeros([n,n,n])
    for i in range (0,n):
        for j in range (0,n):
            for k in range (0,n):
                tensor[i,j,k] = v[i]*v[j]*v[k]
    return tensor