#!/usr/bin/env python
# coding: utf-8

# ## Online Low-Rank Tensor Learning (Online-LRTL)
# 
# This notebook shows how to implement the following paper:
# 
# <div class="alert alert-block alert-info">
# <font color="black">
# Rose Yu, Dehua Cheng, Yan Liu (2015). <b>Accelerated Online Low-Rank Tensor Learning for Multivariate Spatio-Temporal Streams</b>. ICML 2015. <a href="http://proceedings.mlr.press/v37/yua15.pdf" title="PDF">[PDF]</a> 
# </font>
# </div>
# 
# The authors provide Matlab code, if you have any interest, please download at [http://roseyu.com/Materials/accelerate_online_low_rank_tensor.zip](http://roseyu.com/Materials/accelerate_online_low_rank_tensor.zip).

# ### Necessarity

# In[ ]:


import numpy as np

def ten2mat(tensor, mode):
    return np.reshape(np.moveaxis(tensor, mode, 0), (tensor.shape[mode], -1), order = 'F')

def mat2ten(mat, dim, mode):
    index = list()
    index.append(mode)
    for i in range(dim.shape[0]):
        if i != mode:
            index.append(i)
    return np.moveaxis(np.reshape(mat, list(dim[index]), order = 'F'), 0, mode)


# ### Kernel

# **1) How to initialize paramters?**

# In[ ]:


def mapTensorKInit(W, nRank):
    dim = W.shape
    U = []
    for i in range(len(dim)):
        u, _, _ = np.linalg.svds(ten2mat(W, i), min(nRank, dim[i]))
        U.append(u)
    if nRank ~= dim[0]:
        W = np.einsum('pqm, pr, qs, mt -> rst', W, U[0], U[1], U[2])
        W = np.einsum('rst, pr, qs, mt -> pqm', W, U[0], U[1], U[2])
    return W, U

def solve_init(X, Y, nRank):
    """Get the initial estimation from the first batch."""
    P, _ = Y[0].shape
    Q, M = X[0].shape
    nSample = len(Y)
        
    W_mat = []
    W_mat.append(np.zeros((P, Q * M)))
    W_mat.append(np.zeros((Q, P * M)))
    W_mat.append(np.zeros((M, P * Q)))
    
    Sigma_X = np.zeros((Q, Q, M))
    Sigma_Y = np.zeros((P, Q, M))
    
    for m in range(M):
        X_m = np.zeros((Q, nSample))
        Y_m = np.zeros((P, nSample))
        for n in range(nSample):
            X_m[:, n] = X[n][:, m]
            Y_m[:, n] = Y[n][:, m]
        S_X = np.linalg.inv(X_m @ X_m.T + 1e-5 * np.eye(Q)) # Q x Q
        S_Y = Y_m @ X_m.T # P x Q
        W_est = S_Y @ S_X # P x Q
        
        W_mat[0][:, (m - 1) * Q : m * Q] = W_est   # P x MQ
        W_mat[1][:, (m - 1) * P : m * P] = W_est.T # Q x MP
        W_mat[2][m, :] = W_est.T.reshape(P * Q)    # M x PQ
        Sigma_X[:, :, m] = S_X
        Sigma_Y[:, :, m] = S_Y
    W = mat2ten(W_mat[0], np.array([P, Q, M]), 0)  # P x Q x M
    
    W, U = mapTensorKInit(W, nRank, [1, 2, 3])
    W = W.data
    D_L = U
    D_R = 0
    
    return W, Sigma_X, Sigma_Y, D_L, D_R


# **2) How to update parameters?**

# In[ ]:


def mapTensorKGreedy(W, nRank):
    dim = W.shape
    U = []
    W0 = np.zeros(dim)
    for i in range(len(dim)):
        u, s, v = np.linalg.svds(ten2mat(T, i), min(nRank, dim[i]))
        U.append(u)
        W0 += mat2ten(u @ np.diag(s) @ v.T, np.array(dim), i)
    W = W0 / 3
    return W, U

def ALTO(W, U):
    D_L = []
    dim = W.shape
    nRank = U[0].shape[1]
    for i in range(len(dim)):
        tmp = np.random.randn(U[i].shape)
        tmp = tmp - U[i] @ U[i].T @ tmp
        tmp = tmp / np.linalg.norm(tmp)
        D_L.append(np.append(U[i], tmp.reshape([dim[i], 1]), axis = 1))
    T = np.einsum('pqm, pr, qs, mt -> rst', W, D_L[0], D_L[1], D_L[2]) # R+1 x R+1 x R+1
    T, U = mapTensorKGreedy(T, nRank)
    W_out = np.einsum('rst, pr, qs, mt -> pqm', T, D_L[0], D_L[1], D_L[2])
    for i in range(len(dim)):
        D_L[i] = D_L[i] @ U[i]
    return W_out, D_L

def solve_update(W, X, Y, Sigma_X, Sigma_Y, D_L, D_R, nRank, nIter):
    nSample = len(Y)
    Q, _ = X[0].shape
    P, M = Y[0].shape
    
    Sigma_X_out = np.zeros((Q, Q, M))
    Sigma_Y_out = np.zeros((P, Q, M))
    Delta = np.zeros((P, Q, M))
    
    for m in range(M): # solve for each variable seperately
        X_m = np.zeros((Q, nSample))
        Y_m = np.zeros((P, nSample))
        for n in range(nSample):
            X_m[:, n] = X[n][:, m]
            Y_m[:, n] = Y[n][:, m]
        # update Sigma_X and Sigma_Y
        B = Sigma_X[:, :, m] @ X_m # Q x nSample
        E = B @ np.linalg.lstsq(np.eye(nSample) + X_m.T @ B, B.T)[0]
        Delta[:, :, m] = 1 / nIter @ (Y_m @ B.T - Simga_Y[:, :, m] @ E
                                      - Y_m @ X_m.T @ E + Sigma_Y[:, :, m] @ Sigma_X[:, :, m])
        Sigma_X_out[:, :, m] = Sigma_X[:, :, m] - E
        Sigma_Y_out[:, :, m] = Sigma_Y[:, :, m] + Y_m @ X_m.T
    
    W_out, D_L = ALTO(nIter * Delta, D_L)
    
    return W_out, Sigma_X_out, Sigma_Y_out, D_L, D_R


# **3) How to build the kernel online low-rank tensor learning?**

# In[ ]:


def online_tensor_learning(X, Y, At, nInit, batch_size, nRank):
    """
    Main routine for Online Low-Rank Tensor Learning (Online_LRTL)
    Refer to "Accelerated Online Low-Rank Tensor Learning for Multivariate Spatio-Temporal Streams"
    Rose Yu, Dehua Cheng, Yan Liu (ICML 2015)
    Input:
    - X: predictor, list of length 1 x T, each of size Q x M (array)
    - Y: response, list of length 1 x T, each of size P x M (array)
    - At:
      - if matrix, Laplacian matrix
      - if tensor: ground truth tensor, for synthetic evaluation
    - nInit: size of the initial batch
    - batch_size: size of mini-batch
    - nRank: upper bound of the rank
    Output:
    - W: model tensor: P x Q x M
    - obj: objective function value
    """
    
    nSample = len(X)
    nBatch = int((nSample - nInit) / batch_size)
    objs = np.zeros(nBatch)
    
    W, Sigma_X, Sigma_Y, D_L, D_R = solve_init(X[:nInit], Y[:nInit], nRank)
    
    for nIter in range(nBatch):
        if nIter == nBatch:
            X_new = X[nInit + (nBatch - 1) * batch_size - 1 :]
            Y_new = Y[nInit + (nBatch - 1) * batch_size - 1 :]
        else:
            X_new = X[nInit + (nIter - 1) * batch_size - 1 : nIter * batch_size - 1]
            Y_new = Y[nInit + (nIter - 1) * batch_size - 1 : nIter * batch_size - 1]
        W, Sigma_X, Sigma_Y, D_L, D_R = solve_update(W, X_new, Y_new, Sigma_X, Sigma_Y, D_L, D_R, nRank, nIter)
        if len(At.shape) == 3:
            objs[nIter] = np.linalg.norm(W - At) / np.linalg.norm(At)
        else:
            X_test = X[nInit + nIter * batch_size - 1 :]
            Y_test = Y[nInit + nIter * batch_size - 1 :]
            tmp = 0
            for s in range(nSample):
                for m in range(M):
                    tmp += np.linalg.norm(Y_test[s][:, m] - W[:, :, m] @ X_test[s][:, m]) ** 2
            objs[nIter] = np.sqrt(tmp / (nSample * P * M))
        print('Batch: {}'.format(nIter))
        print('Objective function: {:.6}'.format(objs[nIter]))
        print()
        
    return W, objs


# ### Example

# I have no idea for providing simple example currently, but I will try my best to give some!