# @author:                hehaoran
# @environment: python 3.7.4
# pylint: disable = no-member
from operator import index
import numpy as np
from numpy.core.fromnumeric import shape
from tensorly import tensor_to_vec

def kr_prod(a, b):
    return np.einsum('ir, jr -> ijr', a, b).reshape(a.shape[0] * b.shape[0], -1)

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

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

def mat2vec(mat):
    dim1, dim2 = mat.shape
    return mat.T.reshape([dim1 * dim2])

def vec2mat(vec, mat_size):
    return vec.reshape([mat_size[1], mat_size[0]]).T

def prod(arg):
    """ returns the product of elements in arg.
    arg can be list, tuple, set, and array with numerical values. """
    ret = 1
    for i in range(0, len(arg)):
        ret = ret * arg[i]
    return ret

def permute(x, order=None):
    """ Returns a Tensor permuted by the order specified.
    x:np.ndarray.
    """
    ndims = len(x.shape)
    if order is None:
        raise ValueError("Permute: Order must be specified.")

    if order.__class__ == list or order.__class__ == tuple:
        order = np.array(order)

    if ndims != len(order):
        raise ValueError("Permute: Invalid permutation order.")

    if not (sorted(order) == np.arange(ndims)).all():
        raise ValueError("Permute: Invalid permutation order.")

    x = x.transpose(order)

    return x

def tenmat(x=None, rdim=None, cdim=None):
    """
    Create a Tenmat object from a given Tensor X

    paramater:
    ----------
    x: 
        np.ndarray.
    rdim: 
        an one-dim array representing the arranged dimension index for the matrix column
    cdim: 
        an one-dim array representing the arranged dimension index for the matrix row
    
    see also
    ----------
    https://github.com/hanchengge/Helios-PyTen
    """
    if rdim is None:
        raise ValueError('Tenmat: second argument cannot be empty.')
    if rdim is None or cdim is None or rdim.__class__ != list or cdim.__class__ != list:
        raise ValueError("Tenmat: second and third arguments must be specified with list.")
    if rdim.__class__ == list or rdim.__class__ == int:
        rdim = np.array(rdim) - 1
    if cdim is None:
        cdim = np.array([y for y in range(0, x.ndims) if y not in np.zeros(x.ndims - 1) + rdim])
    elif cdim.__class__ == list or cdim.__class__ == int:
        cdim = np.array(cdim) - 1
    else:
        raise ValueError("Tenmat: incorrect specification of dimensions.")
    shape = x.shape
    x = permute(x, np.append(rdim, cdim))
    if type(rdim) != np.ndarray:
        row = prod([shape[y] for y in [rdim]])
    else:
        row = prod([shape[y] for y in rdim])

    if type(cdim) != np.ndarray:
        col = prod([shape[y] for y in [cdim]])
    else:
        col = prod([shape[y] for y in cdim])

    return x.reshape([row, col], order='F')

def ct(d, rank):
  E = np.zeros((rank,rank))
  E[d][d] = 1
  return E

def st(l1, l2, rank):
  T = np.zeros((rank,rank))
  T[l1][l2] =T[l2][l1] = 1
  return T

def inv(T):
    """
    the inverse of four-order tensor
    
    Parameters
    ----------
    T: array_like
        a four-order tensor, For example: (10,10,5,5)--->(50,50). 
        you should modify the shape of original tensor to the above form using the `np.einsum`.
    
    Returns
    -------
    the invere of a four tensor.  
    """
    K, L = T.shape[1:3]
    squ_matrix = np.einsum('ijkl->ikjl', T).reshape((K*L, K*L),order='F')
    t = np.linalg.inv(squ_matrix)
    return np.einsum('ijkl->ikjl', t.reshape((K,L,K,L), order='F'))

if __name__ == "__main__":
    X = np.arange(1, 49).reshape([3, 2, 2, 2, 2], order='F')
    print(X[:, :, 0, 0])
    print(tenmat(X, [1, 2], [3, 4, 5]))