import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
import numpy as np
import h5py
import scipy.io
from timeit import default_timer
import sys
import math

import operator
from functools import reduce
from timeit import default_timer


################################################################
# lowrank layer
################################################################
class LowRank1d(nn.Module):
    def __init__(self, in_channels, out_channels, s, width, rank=1):
        super(LowRank1d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.s = s
        self.n = s
        self.rank = rank

        self.phi = DenseNet([2, 64, 128, 256, width*width*rank], torch.nn.ReLU)
        self.psi = DenseNet([2, 64, 128, 256, width*width*rank], torch.nn.ReLU)


    def forward(self, v, a):
        # a (batch, n, 2)
        # v (batch, n, f)
        batch_size = v.shape[0]

        phi_eval = self.phi(a).reshape(batch_size, self.n, self.out_channels, self.in_channels, self.rank)
        psi_eval = self.psi(a).reshape(batch_size, self.n, self.out_channels, self.in_channels, self.rank)

        # print(psi_eval.shape, v.shape, phi_eval.shape)
        v = torch.einsum('bnoir,bni,bmoir->bmo',psi_eval, v, phi_eval) / self.n

        return v


