# Imports

  import torch
  from torchvision import datasets, models, transforms # All torchvision modules
  import torch.nn as nn  # All neural network modules, nn.Linear, nn.Conv2d, Loss functions,..
  import torch.optim as optim  # For all Optimization algorithms, SGD, Adam,...
  import torch.nn.functional as F  # All functions that don't have any parameters
  from torch.utils.data import (DataLoader,Dataset)  # Gives easier dataset managment and creates mini batches
  import torchvision.datasets as datasets  # Has standard datasets we can import in a nice way
  import torchvision.transforms as transforms  # Transformations we can perform on our dataset
  import torchtext # Makes it easy to work with sequence data 
  from torchtext.data import get_tokenizer
  
  import re # regex library
  import os # Doing operating system operations
  import cv2 # Computer vision tasks with OpenCV
  import numpy as np # Powerful arrray computation library
  from PIL import Image # WOrking with image files
  import pandas # Extracting data from csv
  import math # Math package
  import pickle # Saving variables for later usage.

  from torchsummary import summary # Make understanding of models easier
  import torch # PyTorch library
  from time import time # Using timer in code


  # Set device
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Use Cuda if GPU available!
  
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 

        -------- Helper Classes
                           ++++ Utils class ( Which uses torchtext and other packages to treat images, videos and text)
                           ++++ TextProcessor class (which uses torchtext to treat text)
        --------  Custom Dataset class: PyTorch offers this possibility of creating and processing our custom datasets with much ease
        --------  Model definition classes: 
                                     +++++  VGG16 pretrained model will be used, although a slight modification will be made
                                     +++++  LSTM SEQ2SEQ again, we shall use PyTorch modules to easily create this class (from scratch, unlike the VGG PRETRAINED MODEL)
         It should be noted that one can rewrite a class for a pretrained model (but for our problem, that wasn't necessary)

        --------  Key parameter definition: Here parameters like the learning rate,number of epochs, ... are defined 
        --------  Optimizer definition
        --------  Training
        --------  Testing

         def output_text(self, train_corpus, video = None):

                '''
                Purpose: Generate all text present in video using the csv file which contains all the text in the videos
                Input(s):
                    train_corpus: The file path to the csv file
                    video: A video file whose text is to be generated
                Outputs(s):
                    Final description: the text representing a video caption

                '''

                df = pandas.read_csv(train_corpus)
                if (video):
                    video_id,start,end = self.get_video_id(video) # 

                    final_description=''
                    for i in range(len(df)):

                        if df['VideoID'][i]==str(video_id) and df['Start'][i]==int(start) and df['End'][i]==int(end) and df['Language'][i]=='English':

                            final_description=df['Description'][i]
                else:

                    final_description = []
                    for i in range(len(df)):
                        if (df['Language'][i]=='English'):
                            final_description.append(df['Description'][i])
                return final_description

  def vocab_creator(self,sentence_list):

         '''
         Purpose: From a give corpus, generate a WORD vocabulary which maps a givenn word to a given index
         Input(s): 
             sentence_list: A corpus of all sentences extracted from the videos in the dataset
         Outputs(s):
             word_to_index: the word to index of all words contained in the textual corpus

         '''
         frequencies = {}
         idx = 4
         stoi = {}

         tokenizer = self.get_tokenizer
         for sentence in sentence_list:
             try:
                 for word in tokenizer(sentence):
                     if word not in frequencies:
                         frequencies[word] = 1

                     else:
                         frequencies[word] += 1

                     if frequencies[word] == self.freq_threshold:
                         self.word_to_index[word] = idx
                         idx += 1
             except:
                 pass 
         return self.word_to_index

  def get_output(self, sentence_to_indices, NUMBER_OF_WORDS):

         '''
         Purpose: Generate one - hot representation of sentence, ready for model training
         Input(s): 
             sentence_to_indices: A dictionary which contains the words and indices as key, value pairs
             NUMBER_OF_WORDS: The maximum number of words a sentence can contain
         Outputs(s):
             One-hot vectors stacked into an array

         '''

         arr = np.zeros((NUMBER_OF_WORDS, self.VOCAB_SIZE))
         pad_number = 1 # The pad in sentence to index is seen as 1
         for i in range(len(arr)):
             if(i<len(sentence_to_indices)):
                 arr[i][sentence_to_indices[list(sentence_to_indices.keys())[i]]] = 1 # set a given key to 1, while leaving the others at zero
             else:
                 arr[i][pad_number] = 1 # pad to complete the remaining words to make up the NUMBER OF WORDS needed for the model

         return arr

   def __getitem__(self, idx):

          textprocessor = TextProcessor(VOCAB_SIZE = self.VOCAB_SIZE)
          utils = Utils()


          video_file = self.train_dir_list[idx] # get video file corresponding to the id, idx


          output_text = self.utils.output_text(self.train_corpus, video_file) # get the text contained in the video file


          #### generate input 2,  from the output_text
          sentence_to_index = textprocessor.sentence_to_indices(utils.tagger_input(utils.clean_text(output_text)), self.word_to_index)
          X_2 = textprocessor.get_output(sentence_to_index, NUMBER_OF_WORDS)

          #### generate output,  from the output_text
          sentence_to_index = textprocessor.sentence_to_indices(utils.tagger_output(utils.clean_text(output_text)), self.word_to_index) 
          y = textprocessor.get_output(sentence_to_index, NUMBER_OF_WORDS)

          video_path = self.train_dir + video_file

          # generate input 1
          X_1 = utils.video_to_frames(video_path, self.number_of_frames, self.device, self.INPUT_SIZE, self.model, self.transform)
          #X_1 = pre_data[idx] ### this may be used if we collect the pretrained video frames from a pickle file instead of doing inference during training
          return (X_1,torch.tensor(X_2)), torch.tensor(y)

    ### parametres

   LEARNING_RATE = 1e-3
   NUMBER_OF_FRAMES = 40
   BATCH_SIZE = 1
   EPOCH = 10
   TRAINING_DEVICE = 'cuda'
   VOCAB_SIZE = 200
   NUMBER_OF_WORDS = 10
   HIDDEN_SIZE = 300
   INPUT_SIZE = 4096
   NUMBER_OF_LAYERS = 1
   tsfm = transforms.Compose([
       transforms.Resize([224, 224]),
       transforms.ToTensor(),
       transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
   ])
   train_dir = 'D:/Machine_Learning/datasets/YouTubeClips_2/YouTubeClips/'
   train_corpus = 'D:/Machine_Learning/datasets/video_corpus/video_corpus.csv'

   ### optimizer and loss function
   optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
   criterion = nn.CrossEntropyLoss()

   ### Sequence to sequence model

   class Encoder_LSTM(nn.Module):
       def __init__(self, input_size, hidden_size, num_layers):
           super(Encoder_LSTM, self).__init__()
           self.hidden_size = hidden_size
           self.num_layers = num_layers
           self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)

       def forward(self, x):
           out, (h_n, c_n) = self.lstm(x)  # out: tensor of shape (batch_size, seq_length, hidden_size)

           return h_n, c_n

   class Decoder_LSTM(nn.Module):
       def __init__(self, input_size, hidden_size, num_layers, number_of_words):
           super(Decoder_LSTM, self).__init__()
           self.hidden_size = hidden_size
           self.num_layers = num_layers
           self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
           self.fc = nn.Linear(hidden_size, input_size)
       def forward(self, x, h_n, c_n):
           output, _ = self.lstm(x.float(),(h_n,c_n))  # out: tensor of shape (batch_size, seq_length, hidden_size)
           output = self.fc(output)                            

           return output

   class Seq2Seq(nn.Module):
       def __init__(self, encoder, decoder):
           super(Seq2Seq, self).__init__()
           self.encoder = encoder
           self.decoder = decoder

       def forward(self, X_1, X_2):
           h_n, c_n = self.encoder(X_1)
           output = self.decoder(X_2, h_n, c_n)
           return output

 #### Model Training

EPOCH = 10
import time
print_feq = 100
best_loss = np.inf
for epoch in range(1, EPOCH+1):
    model.train()
    epoch_loss = 0

    for step, (img,label) in enumerate(train_dl):


        time_1 = time.time() ## timing

        X_1, X_2 = img ### get inputs

        X_1 = X_1.to(device) # Set device 
        X_2 = X_2.to(device) # Set device


        label = label.to(device) # Set output device

        ### zero the parameter gradients
        optimizer.zero_grad()

        ### forward
        prediction = model(X_1, X_2)

        ### Optimize
        prediction = prediction.to(device)
        prediction = torch.squeeze(prediction,0)
        label = torch.squeeze(label,0)

        new_label = torch.zeros([label.shape[0]])
        for l in range(label.shape[0]):
            new_label[l] = np.argmax(label[l].cpu())
        new_label = new_label.to(device)
        loss = criterion(prediction, new_label.long())

        # Backward prop.
        loss.backward()
        optimizer.step()

        ### print out statistics
        epoch_loss += loss.item()
        if step % print_feq == 0:
            print('epoch:', epoch,
                  '\tstep:', step+1, '/', len(train_dl) + 1,
                  '\ttrain loss:', '{:.4f}'.format(loss.item()),
                  '\ttime:', '{:.4f}'.format((time.time()-time_1)*print_feq), 's')
        torch.save(model.state_dict(), 'model_lstm_2.pth')
    ### save best model
    if(epoch_loss < best_loss):
        best_loss = epoch_loss
        torch.save(model.state_dict(), 'model_lstm_best_loss.pth')
    print("The loss for this epoch is = :", epoch_loss/lent(train_dl))

 python train.py --NUMBER_OF_WORDS 32 --train_dir D:/Machine_Learning/datasets/YouTubeClips_2/YouTubeClips/ --train_corpus  D:/Machine_Learning/datasets/video_corpus/video_corpus.csv --load_weights model_lstm_2.pth

   python test.py --HIDDEN_SIZE 300 --video_file D:/Machine_Learning/datasets/YouTubeClips_2/YouTubeClips/_O9kWD8nuRU_45_49.avi --load_weights model_lstm_best_loss.pth --train_corpus D:/Machine_Learning/datasets/video_corpus/video_corpus.csv

     Definition
      POST https://api.wit.ai/speech

         Example request
      $ curl -XPOST 'https://api.wit.ai/speech?v=20200513' \
        -i -L \
        -H "Authorization: Bearer $TOKEN" \ #### which you get when you login to wit.ai and go to settings
        -H "Content-Type: audio/wav" \
        --data-binary "@sample.wav"

   {
    "msg_id" : ".....",
    "msg_body" : "People are fighting",
    "outcome" : {
      "intent" : "......",
      "entities" : {
        "object" : {
          "value" : "......",
          "body" : "....."
        }
      },
      "confidence" : ....
    }
  }