# -*- coding: utf-8 -*-
"""
Copy of GAN with Generator: LSTM, Discriminator: Convolutional NN with ECG Data

 Introduction
 ------------
    The aim of this script is to use a convolutional neural network with 
    a max pooling layer in the discrimiantor. 
    This was found to work well with the Physionet ECG data in a paper. 
    They used two convolutional NN so we will compare the difference between the 
    images generated using a single layer of CNN in the discriminator and 2 CNN layers 
    to see if this improves the quality of series generated.

"""
"""
Bringing in required dependencies as defined in the GitHub repo: 
    https://github.com/josipd/torch-two-sample/blob/master/torch_two_sample/permutation_test.pyx"""
from __future__ import division

import torch
from tqdm import tqdm
import numpy as np
from matplotlib import pyplot as plt
import seaborn as sns

from torchvision import transforms
from torch.autograd.variable import Variable
sns.set(rc={'figure.figsize':(11, 4)})

import datetime 
from datetime import date
today = date.today()

import random
import json as js
import pickle
import os

from data import ECGData, PD_to_Tensor
from Model import Generator, Discriminator 

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

if device == 'cuda:0':
    print('Using GPU : ')
    print(torch.cuda.get_device_name(device))
else :
    print('Using CPU')


"""#MMD Evaluation Metric Definition
Using MMD to determine the similarity between distributions

PDIST code comes from torch-two-sample utils code: 
    https://github.com/josipd/torch-two-sample/blob/master/torch_two_sample/util.py
"""

def pdist(sample_1, sample_2, norm=2, eps=1e-5):
    r"""Compute the matrix of all squared pairwise distances.
    Arguments
    ---------
    sample_1 : torch.Tensor or Variable
        The first sample, should be of shape ``(n_1, d)``.
    sample_2 : torch.Tensor or Variable
        The second sample, should be of shape ``(n_2, d)``.
    norm : float
        The l_p norm to be used.
    Returns
    -------
    torch.Tensor or Variable
        Matrix of shape (n_1, n_2). The [i, j]-th entry is equal to
        ``|| sample_1[i, :] - sample_2[j, :] ||_p``."""
    n_1, n_2 = sample_1.size(0), sample_2.size(0)
    norm = float(norm)
    
    if norm == 2.:
        norms_1 = torch.sum(sample_1**2, dim=1, keepdim=True)
        norms_2 = torch.sum(sample_2**2, dim=1, keepdim=True)
        norms = (norms_1.expand(n_1, n_2) +
                 norms_2.transpose(0, 1).expand(n_1, n_2))
        distances_squared = norms - 2 * sample_1.mm(sample_2.t())
        return torch.sqrt(eps + torch.abs(distances_squared))
    else:
        dim = sample_1.size(1)
        expanded_1 = sample_1.unsqueeze(1).expand(n_1, n_2, dim)
        expanded_2 = sample_2.unsqueeze(0).expand(n_1, n_2, dim)
        differences = torch.abs(expanded_1 - expanded_2) ** norm
        inner = torch.sum(differences, dim=2, keepdim=False)
        return (eps + inner) ** (1. / norm)

def permutation_test_mat(matrix,
                         n_1,  n_2,  n_permutations,
                          a00=1,  a11=1,  a01=0):
    """Compute the p-value of the following statistic (rejects when high)
        \sum_{i,j} a_{\pi(i), \pi(j)} matrix[i, j].
    """
    n = n_1 + n_2
    pi = np.zeros(n, dtype=np.int8)
    pi[n_1:] = 1

    larger = 0.
    count = 0
    
    for sample_n in range(1 + n_permutations):
        count = 0.
        for i in range(n):
            for j in range(i, n):
                mij = matrix[i, j] + matrix[j, i]
                if pi[i] == pi[j] == 0:
                    count += a00 * mij
                elif pi[i] == pi[j] == 1:
                    count += a11 * mij
                else:
                    count += a01 * mij
        if sample_n == 0:
            statistic = count
        elif statistic <= count:
            larger += 1

        np.random.shuffle(pi)

    return larger / n_permutations

"""Code from Torch-Two-Samples at https://torch-two-sample.readthedocs.io/en/latest/#"""

class MMDStatistic:
    r"""The *unbiased* MMD test of :cite:`gretton2012kernel`.

    The kernel used is equal to:

    .. math ::
        k(x, x') = \sum_{j=1}^k e^{-\alpha_j\|x - x'\|^2},

    for the :math:`\alpha_j` proved in :py:meth:`~.MMDStatistic.__call__`.

    Arguments
    ---------
    n_1: int
        The number of points in the first sample.
    n_2: int
        The number of points in the second sample."""

    def __init__(self, n_1, n_2):
        self.n_1 = n_1
        self.n_2 = n_2

        # The three constants used in the test.
        self.a00 = 1. / (n_1 * (n_1 - 1))
        self.a11 = 1. / (n_2 * (n_2 - 1))
        self.a01 = - 1. / (n_1 * n_2)

    def __call__(self, sample_1, sample_2, alphas, ret_matrix=False):
        r"""Evaluate the statistic.

        The kernel used is

        .. math::

            k(x, x') = \sum_{j=1}^k e^{-\alpha_j \|x - x'\|^2},

        for the provided ``alphas``.

        Arguments
        ---------
        sample_1: :class:`torch:torch.autograd.Variable`
            The first sample, of size ``(n_1, d)``.
        sample_2: variable of shape (n_2, d)
            The second sample, of size ``(n_2, d)``.
        alphas : list of :class:`float`
            The kernel parameters.
        ret_matrix: bool
            If set, the call with also return a second variable.

            This variable can be then used to compute a p-value using
            :py:meth:`~.MMDStatistic.pval`.

        Returns
        -------
        :class:`float`
            The test statistic.
        :class:`torch:torch.autograd.Variable`
            Returned only if ``ret_matrix`` was set to true."""
        sample_12 = torch.cat((sample_1, sample_2), 0)
        distances = pdist(sample_12, sample_12, norm=2)

        kernels = None
        for alpha in alphas:
            kernels_a = torch.exp(- alpha * distances ** 2)
            if kernels is None:
                kernels = kernels_a
            else:
                kernels = kernels + kernels_a

        k_1 = kernels[:self.n_1, :self.n_1]
        k_2 = kernels[self.n_1:, self.n_1:]
        k_12 = kernels[:self.n_1, self.n_1:]

        mmd = (2 * self.a01 * k_12.sum() +
               self.a00 * (k_1.sum() - torch.trace(k_1)) +
               self.a11 * (k_2.sum() - torch.trace(k_2)))
        if ret_matrix:
            return mmd, kernels
        else:
            return mmd


    def pval(self, distances, n_permutations=1000):
        r"""Compute a p-value using a permutation test.

        Arguments
        ---------
        matrix: :class:`torch:torch.autograd.Variable`
            The matrix computed using :py:meth:`~.MMDStatistic.__call__`.
        n_permutations: int
            The number of random draws from the permutation null.

        Returns
        -------
        float
            The estimated p-value."""
        if isinstance(distances, Variable):
            distances = distances.data
        return permutation_test_mat(distances.cpu().numpy(),
                                    self.n_1, self.n_2,
                                    n_permutations,
                                    a00=self.a00, a11=self.a11, a01=self.a01)

"""

This paper 
https://arxiv.org/pdf/1611.04488.pdf says that the most common way to 
calculate sigma is to use the median pairwise distances between the joint data.

"""

def pairwisedistances(X,Y,norm=2):
    dist = pdist(X,Y,norm)
    return np.median(dist.numpy())


""" 

Function for loading ECG Data 

"""
def GetECGData(source_file,class_id):
  compose = transforms.Compose(
        [PD_to_Tensor()
        ])
  return ECGData(source_file ,class_id = class_id, transform = compose)

"""

Creating the training set of sine/ECG signals

"""

#Taking normal ECG data for now
source_filename = './mitbih_train.csv'
ecg_data = GetECGData(source_file = source_filename,class_id = 0)

sample_size = 119 #batch size needed for Data Loader and the noise creator function.

# Create loader with data, so that we can iterate over it

data_loader = torch.utils.data.DataLoader(ecg_data, batch_size=sample_size, shuffle=True)
# Num batches
num_batches = len(data_loader)
print(num_batches)

"""Creating the Test Set"""
test_filename =  './mitbih_test.csv'

ecg_data_test = GetECGData(source_file = test_filename,class_id = 0)

data_loader_test = torch.utils.data.DataLoader(ecg_data_test[:18088], batch_size=sample_size, shuffle=True)



"""##Defining the noise creation function"""

def noise(batch_size, features):
  noise_vec = torch.randn(batch_size, features).to(device)
  return noise_vec

"""#Initialising Parameters"""

seq_length = ecg_data[0].size()[0] #Number of features


#Params for the generator
hidden_nodes_g = 50
layers = 2
tanh_layer = False

#No. of training rounds per epoch
D_rounds = 3
G_rounds = 1
num_epoch = 35
learning_rate = 0.0002

#Params for the Discriminator
minibatch_layer = 0
minibatch_normal_init_ = True
num_cvs = 2
cv1_out= 10
cv1_k = 3
cv1_s = 1
p1_k = 3
p1_s = 2
cv2_out = 10
cv2_k = 3
cv2_s = 1
p2_k = 3
p2_s = 2

"""# Evaluation of GAN with 2 CNN Layer in Discriminator

##Generator and Discriminator training phase
"""

minibatch_out = [0,3,5,8,10]
for minibatch_layer in minibatch_out:
  path = ".../your_path/Run_"+str(today.strftime("%d_%m_%Y"))+"_"+ str(datetime.datetime.now().time()).split('.')[0]
  os.mkdir(path)
  
  dict = {'data' : source_filename, 
          'sample_size' : sample_size, 
          'seq_length' : seq_length,
          'num_layers': layers, 
          'tanh_layer': tanh_layer,
          'hidden_dims_generator': hidden_nodes_g, 
          'minibatch_layer': minibatch_layer,
          'minibatch_normal_init_' : minibatch_normal_init_,
          'num_cvs':num_cvs,
          'cv1_out':cv1_out,
          'cv1_k':cv1_k,
          'cv1_s':cv1_s,
          'p1_k':p1_k,
          'p1_s':p1_s,
          'cv2_out':cv2_out,
          'cv2_k':cv2_k,
          'cv2_s':cv2_s,
          'p2_k':p2_k,
          'p2_s':p2_s,
          'num_epoch':num_epoch,
          'D_rounds': D_rounds,
          'G_rounds': G_rounds,  
          'learning_rate' : learning_rate
         }

  json = js.dumps(dict)
  f = open(path+"/settings.json","w")
  f.write(json)
  f.close()

  generator_1 = Generator(seq_length,sample_size,hidden_dim =  hidden_nodes_g, tanh_output = tanh_layer).to(device)
  discriminator_1 = Discriminator(seq_length, sample_size ,minibatch_normal_init = minibatch_normal_init_, minibatch = minibatch_layer,num_cv = num_cvs, cv1_out = cv1_out,cv1_k = cv1_k, cv1_s = cv1_s, p1_k = p1_k, p1_s = p1_s, cv2_out= cv2_out, cv2_k = cv2_k, cv2_s = cv2_s, p2_k = p2_k, p2_s = p2_s).to(device)
  #Loss function 
  loss_1 = torch.nn.BCELoss()

  generator_1.train()
  discriminator_1.train()
  
  d_optimizer_1 = torch.optim.Adam(discriminator_1.parameters(),lr = learning_rate)
  g_optimizer_1 = torch.optim.Adam(generator_1.parameters(),lr = learning_rate)

  G_losses = []
  D_losses = []
  mmd_list = []
  series_list = np.zeros((1,seq_length))


  for n in tqdm(range(num_epoch)):
     # for k in range(1):

      for n_batch, sample_data in enumerate(data_loader):
      ### TRAIN DISCRIMINATOR ON FAKE DATA
        for d in range(D_rounds):
          discriminator_1.zero_grad()

          h_g = generator_1.init_hidden()

          #Generating the noise and label data
          noise_sample = Variable(noise(len(sample_data),seq_length))

          #Use this line if generator outputs hidden states: dis_fake_data, (h_g_n,c_g_n) = generator.forward(noise_sample,h_g)
          dis_fake_data = generator_1.forward(noise_sample,h_g).detach()

          y_pred_fake = discriminator_1(dis_fake_data)

          loss_fake = loss_1(y_pred_fake,torch.zeros([len(sample_data),1]).to(device))
          loss_fake.backward()    

          #Train discriminator on real data   
          real_data = Variable(sample_data.float()).to(device)    
          y_pred_real  = discriminator_1.forward(real_data)

          loss_real = loss_1(y_pred_real,torch.ones([len(sample_data),1]).to(device))
          loss_real.backward()

          d_optimizer_1.step() #Updating the weights based on the predictions for both real and fake calculations.



        #Train Generator  
        for g in range(G_rounds):
          generator_1.zero_grad()
          h_g = generator_1.init_hidden()

          noise_sample = Variable(noise(len(sample_data), seq_length))


          #Use this line if generator outputs hidden states: gen_fake_data, (h_g_n,c_g_n) = generator.forward(noise_sample,h_g)
          gen_fake_data = generator_1.forward(noise_sample,h_g)
          y_pred_gen = discriminator_1(gen_fake_data)

          error_gen = loss_1(y_pred_gen,torch.ones([len(sample_data),1]).to(device))
          error_gen.backward()
          g_optimizer_1.step()         
    
      if n_batch ==( num_batches - 1):
          G_losses.append(error_gen.item())
          D_losses.append((loss_real+loss_fake).item())
          
          torch.save(generator_1.state_dict(), path+'/generator_state_'+str(n)+'.pt')
          torch.save(discriminator_1.state_dict(),path+ '/discriminator_state_'+str(n)+'.pt')
          
        # Check how the generator is doing by saving G's output on fixed_noise
          with torch.no_grad():
              h_g = generator_1.init_hidden()
              fake = generator_1(noise(len(sample_data), seq_length),h_g).detach().cpu()
              generated_sample = torch.zeros(1,seq_length).to(device)
              
              for iter in range(0,int(len(ecg_data_test[:18088])/sample_size)):
                noise_sample_test = noise(sample_size, seq_length)
                h_g = generator_1.init_hidden()
                generated_data = generator_1.forward(noise_sample_test,h_g).detach().squeeze()
                generated_sample = torch.cat((generated_sample,generated_data),dim = 0)
             
              
              # Getting the MMD Statistic for each Training Epoch
              generated_sample = generated_sample[1:][:]
              sigma = [pairwisedistances(ecg_data_test[:18088].type(torch.DoubleTensor),generated_sample.type(torch.DoubleTensor).squeeze())] 
              mmd = MMDStatistic(len(ecg_data_test[:18088]),generated_sample.size(0))
              mmd_eval = mmd(ecg_data_test[:18088].type(torch.DoubleTensor),generated_sample.type(torch.DoubleTensor).squeeze(),sigma, ret_matrix=False)
              mmd_list.append(mmd_eval.item())
              
          
              series_list = np.append(series_list,fake[0].numpy().reshape((1,seq_length)),axis=0)
          
  #Dumping the errors and mmd evaluations for each training epoch.
  with open(path+'/generator_losses.txt', 'wb') as fp:
      pickle.dump(G_losses, fp)
  with open(path+'/discriminator_losses.txt', 'wb') as fp:
      pickle.dump(D_losses, fp)   
  with open(path+'/mmd_list.txt', 'wb') as fp:
      pickle.dump(mmd_list, fp)
  
  #Plotting the error graph
  plt.plot(G_losses,'-r',label='Generator Error')
  plt.plot(D_losses, '-b', label = 'Discriminator Error')
  plt.title('GAN Errors in Training')
  plt.legend()
  plt.savefig(path+'/GAN_errors.png')
  plt.close() 
  
  #Plot a figure for each training epoch with the MMD value in the title
  i = 0
  while i < num_epoch:
    if i%3==0:
      fig, ax = plt.subplots(3,1,constrained_layout=True)
      fig.suptitle("Generated fake data")
    for j in range(0,3):
      ax[j].plot(series_list[i][:])
      ax[j].set_title('Epoch '+str(i)+ ', MMD: %.4f' % (mmd_list[i]))
      i = i+1
     
    plt.savefig(path+'/Training_Epoch_Samples_MMD_'+str(i)+'.png')
    plt.close(fig) 
  #Checking the diversity of the samples:
  generator_1.eval()
  h_g = generator_1.init_hidden()
  test_noise_sample = noise(sample_size, seq_length)
  gen_data= generator_1.forward(test_noise_sample,h_g).detach()


  plt.title("Generated ECG Waves")
  plt.plot(gen_data[random.randint(0,sample_size-1)].tolist(),'-b')
  plt.plot(gen_data[random.randint(0,sample_size-1)].tolist(),'-r')
  plt.plot(gen_data[random.randint(0,sample_size-1)].tolist(),'-g')
  plt.plot(gen_data[random.randint(0,sample_size-1)].tolist(),'-', color = 'orange')
  plt.savefig(path+'/Generated_Data_Sample1.png')
  plt.close()