import numpy as np
import torch; torch.set_default_dtype(torch.float64)
import torch.nn as nn
import torch.optim as optim
import copy
import pickle as pkl
import math
import matplotlib.pyplot as plt
from tqdm import tqdm

device = 'cuda' if torch.cuda.is_available() else 'cpu'

# The following command loads the data. For this to work the file "movie_data.p"
# must be uploaded to this notebook. To do that, navigate to the Jupyter Notebook
# interface home page and use the “Upload” option.
contents = pkl.load(open("movie_data.p","rb"))
data = contents['data']

# The data file contains a graph of movie similarities. The graph has P nodes.
# This is the graph we will use to run the GNN.
S = data['S'].to(device)
P = S.shape[0]

# The data file contains a training set with entries (x_n, y_n, p_n).
# These entries are stored in the arrays xTrain, yTrain, and pTrain,
# respectively. All of these arrays have nTrain columns, with each
# column corresponding to a different entry in the dataset.
xTrain = data['train']['x'].to(device) # Available ratings for making predictions.
yTrain = data['train']['y'].to(device) # Rating to be predicted
pTrain = data['train']['p'].to(device) # Index of the movie whose rating is being predicted
nTrain = xTrain.shape[0]

# Same thing with test data
xTest = data['test']['x'].to(device)
yTest = data['test']['y'].to(device)
pTest = data['test']['p'].to(device)
nTest = xTest.shape[0]

def movie_mse_loss(yHat, y, idxMovie):
    """
    evaluates the square of comparing yHat(p) with y. The function
    accepts as an input tensors with multiple movie ratings.
    Inputs:
      yHat: torch.Tensor (NumUsers x 1 x NumMovies)
        A set of rating estimates. It includes estimates of all movies.
      y: torch.Tensor (NumUsers x 1)
        A specific rating of a specific movie
      idxMovie: torch.Tensor
        The index of the movie whose rating is given by y

    Outputs:
      mse: jnp.ndarray
        Computed mean squared error.

    """
    # the .squeeze() method is needed for dimension match between y and yHat
    yHat = yHat.squeeze()
    # Isolate the predictions in yHat that we will use for comparison
    prediction = yHat[:,idxMovie.cpu().numpy()]
    return torch.mean((prediction-y)**2)

def filter_function(x, h, S, b):
    '''
    The following function defines a Graph Filter.
    Inputs:
      x: Input to Graph Filter
      h: Filter
      S: Graph Shift Operator
      b: Bias term

      Outpus:
      y: Output of Graph Filter
    '''
    # X is  B x G x N
    # S is N x N
    B, G, N = x.shape
    K, _, F = h.shape
    y = torch.zeros((B, N, F), device=device)
    # The following for-loop is utilized to perform the graph diffusions
    for k in range(K):
        # sum S^k x * h_k
        # We permute x dimensions it to get dimensions B x N x G
        # The filter has dimensions G x F
        y += torch.matmul(x.permute((0,2,1)), h[k])
        # diffuse signal
        x = torch.matmul(x, S)
    # y has dims B x N x F,
    # add bias of shape F
    y = y + b
    # permute dimensions to get the desired dimensions B x F x N
    y = y.permute(0, 2, 1)
    return  y

class FilterModule(nn.Module):
    #First we initialize the graph filter class and define the attributes
    def __init__(self, k, f_in, f_out):
        """
        initialize graph filter
        :param k: order of the filter
        :param f_in: filter input dimension
        :param f_out:  filter output dimension
        """
        super().__init__()
        self.k     = k
        self.f_in  = f_in
        self.f_out = f_out
        # initialize weights and bias
        self.bias = nn.parameter.Parameter(torch.Tensor(self.f_out))
        self.h = nn.Parameter(torch.randn(self.k, self.f_in, self.f_out))
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.f_in * self.k)
        self.h.data.uniform_(-stdv, stdv)
        self.bias.data.uniform_(-stdv, stdv)

    def forward(self, x, S):
        """
        PyTorch module interface implementation; just apply the filter
        """
        return filter_function(x, self.h, S, self.bias)

class GNNModule(nn.Module):
    # setup method gets called on init
    def __init__(self, l, k,  f, sigma):
      """
      param:l: (int) number of layers
      param:k: (list of length n) number of taps for each layer
      param:f: (list of length n+1) input dimension to each layer
      sigma: (nn.Module) nonlinear activation function (same for all layers)
      """
      super().__init__()
      # non linearity
      self.sigma = sigma
      gml = []
      # append layers
      for layer in range(l-1):
          gml.append(FilterModule(k[layer], f[layer],f[layer+1]).to(device))
      self.gml = gml
      # add last linear layer
      self.readout = FilterModule(1,f[layer+1], f[layer+2])

    # Define the forward pass
    def forward(self, x, S):
        for layer in self.gml:
            # graph filter
            x = layer(x, S)
            # non-linearity
            x = self.sigma(x)
        # linear readout
        x = self.readout(x, S)
        return x

l = 3
k = [5,5,1]
f = [1, 64, 32, 1]
sigma = nn.ReLU()

model = GNNModule(l, k, f, sigma).to(device)

# learning parameters
nEpochs = 8
learningRate = 0.05
batchSize = 20
optimizer = optim.Adam(model.parameters(), lr=learningRate)

# Helper variables for data loading during training.
batchIndex = np.append(np.arange(0, nTrain, batchSize), nTrain)
nBatches = int(np.ceil(nTrain/batchSize))