#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue May 19 16:26:11 2020

@author: junbingao

To do list:
    
1. How to extend this for tensorial LSTMs, see https://mlexplained.com/2019/02/15/building-an-lstm-from-scratch-in-pytorch-lstms-in-depth-part-1/
   also https://github.com/pytorch/benchmark/blob/master/rnns/fastrnns/custom_lstms.py
   
2. How to give a name for each layer weights

3. How to use jit to speed up https://pytorch.org/docs/stable/jit.html
"""

import torch
import numpy as np



class MyTuckerLayer(torch.nn.Module):
    def __init__(self, in_dims, out_dims, bias = True):
        super(MyTuckerLayer, self).__init__()
        self.nDims = len(in_dims)
        # The first dimension of inputs is batch size
        assert len(in_dims) == len(out_dims)  
        self.layers = torch.nn.ModuleList()
        self.bias = bias
        # We add 1 to the batach size dimension
        self.permuteIdx = np.arange(self.nDims + 1).tolist()
        # [0, 1, 2, ..., nDims]  where 0 refers to batch size
        # 1, 2, ..., nDims refer to Tensor data mode
        # We want to apply mode multiplication one mode by one mode
        # As pytorch Linear layer apply to the last mode, so we will
        # permute the first mode to the last mode
        self.permuteIdx.append(self.permuteIdx.pop(self.permuteIdx.index(1)))
        
        for i in range(self.nDims):
            layer = torch.nn.Linear(in_dims[i], out_dims[i], bias=False)
            torch.nn.init.xavier_normal_(layer.weight)
            self.layers.append(layer)
        if self.bias:
            #self.BiasValue = torch.autograd.Variable(torch.zeros(out_dims), requires_grad=True)
            self.BiasValue = torch.nn.Parameter(torch.zeros(out_dims)) 
        else:
            self.BiasValue = torch.tensor(0.0)
            
    def forward(self, x):
        # tensor x must have consistent input dimension with in_dims
        # The following does the Tucker 
        #  Y = X \times_1 W_1 \times_2 W_2 \cdots \times_d W_d + B 
        for i in range(self.nDims):
            # Each time we permute a mode to the last mode for mode \times_i W_i
            x = x.permute(self.permuteIdx)
            x = self.layers[i](x)
        if self.bias:
            x = x + self.BiasValue
        return x

#Let us try a two layer 3D tensor network for simple regression
#from X (size [3,4,5]) to Y (size [2,3,4]) with LS error objective
class MyTensorNetWork(torch.nn.Module):
    def __init__(self, in_dims, out_dims):
        super().__init__()
        self.layer1 = MyTuckerLayer(in_dims, [10, 12, 15], bias = False)
        self.layer2 = MyTuckerLayer([10, 12, 15], out_dims)

    def forward(self, x):
        x = self.layer1(x)
        x = torch.relu(x)
        x = self.layer2(x)
        return torch.sigmoid(x)
 
# Data ...    for Six Data 
X = torch.rand((6, 3, 4, 5))
targetY = torch.rand((6, 2, 3, 4))

myNet = MyTensorNetWork([3, 4, 5], [2,3,4])

loss_fn = torch.nn.MSELoss(reduction='sum')

learning_rate = 1e-3
optimizer = torch.optim.Adam(myNet.parameters(), lr=learning_rate)
for t in range(5000):
    # Forward pass: compute predicted y by passing x to the model.
    y_pred = myNet(X)

    # Compute and print loss.
    loss = loss_fn(y_pred, targetY)
    if t % 100 == 99:
        print(t, loss.item())

    optimizer.zero_grad()

    # Backward pass: compute gradient of the loss with respect to model
    # parameters
    loss.backward()

    # Calling the step function on an Optimizer makes an update to its
    # parameters
    optimizer.step()

# Check the trained paratemers Ws
Ws = list(myNet.parameters())