text: # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
  
  import argparse
  import itertools
  import logging
  import os
  import numpy as np
  
  import pytorch_lightning
  from torch.utils.data.sampler import SequentialSampler
  import pytorchvideo.data
  import pytorchvideo.models.resnet
  import pytorchvideo.models.s3dg
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from pytorch_lightning.callbacks import LearningRateMonitor
  from pytorchvideo.transforms import (
      ApplyTransformToKey,
      Normalize,
      RandomShortSideScale,
      RemoveKey,
      ShortSideScale,
      UniformTemporalSubsample,
  )
  # from slurm import copy_and_run_with_config
  from torch.utils.data import DistributedSampler, RandomSampler
  from torchaudio.transforms import MelSpectrogram, Resample
  from torchvision.transforms import (
      CenterCrop,
      Compose,
      Lambda,
      RandomCrop,
      RandomHorizontalFlip,
  )
  
  from gensim.models.keyedvectors import KeyedVectors
  
  pytorch_lightning.utilities.seed.seed_everything(0)
  
  
  """
  This video classification example demonstrates how PyTorchVideo models, datasets and
  transforms can be used with PyTorch Lightning module. Specifically it shows how a
  simple pipeline to train a Resnet on the Kinetics video dataset can be built.
  
  Don't worry if you don't have PyTorch Lightning experience. We'll provide an explanation
  of how the PyTorch Lightning module works to accompany the example.
  
  The code can be separated into three main components:
  1. VideoClassificationLightningModule (pytorch_lightning.LightningModule), this defines:
      - how the model is constructed,
      - the inner train or validation loop (i.e. computing loss/metrics from a minibatch)
      - optimizer configuration
  
  2. KineticsDataModule (pytorch_lightning.LightningDataModule), this defines:
      - how to fetch/prepare the dataset
      - the train and val dataloaders for the associated dataset
  
  3. pytorch_lightning.Trainer, this is a concrete PyTorch Lightning class that provides
    the training pipeline configuration and a fit(<lightning_module>, <data_module>)
    function to start the training/validation loop.
  
  All three components are combined in the train() function. We'll explain the rest of the
  details inline.
  """
  
  
  
  class S3DGVideoTextRepresentator(pytorch_lightning.LightningModule):
      def __init__(self, args):
          """
          This LightningModule implementation constructs a PyTorchVideo ResNet,
          defines the train and val loss to be trained with (cross_entropy), and
          configures the optimizer.
          """
          self.args = args
          super().__init__()
  
          #############
          # PTV Model #
          #############
  
          # Instantiate the model
          self.net = pytorchvideo.models.s3dg.S3D(args.s3d_dict, 512)
  
          # Load the model weights
          self.net.load_state_dict(torch.load(args.s3d_model))
          self.net.eval()
      
      def forward(self, x):
          # Video inference
          video_input = x["video"] / 255
          video_output = self.net(video_input)["video_embedding"]
          # video_emb = F.normalize(video_output)
          video_emb = video_output
  
          # Text inference
          text_output = self.net.text_module(x["sentence"])["text_embedding"]
          # text_emb = F.normalize(text_output)
          text_emb = text_output
          return [video_emb, text_emb]
  
      def on_train_epoch_start(self):
          """
          For distributed training we need to set the datasets video sampler epoch so
          that shuffling is done correctly
          """
          epoch = self.trainer.current_epoch
  
      def sim_matrix(self, x):
          """
          Forward defines the prediction/inference actions.
          """
          video_emb = self.video_model(x["video"])
          video_emb = self.video_head(video_emb)
          video_emb = video_emb.view(video_emb.shape[0], -1)
          video_emb = self.GU_video(video_emb)
          # text_emb = self.GU_text(self.text_pooling(x["sentence"]))
          text_emb = self.GU_text(self.text_pooling(x["text_emb"]))
          return torch.matmul(text_emb, video_emb.t())
  
      def training_step(self, batch, batch_idx):
          """
          This function is called in the inner loop of the training epoch. It must
          return a loss that is used for loss.backwards() internally. The self.log(...)
          function can be used to log any training metrics.
  
          PyTorchVideo batches are dictionaries containing each modality or metadata of
          the batch collated video clips. Kinetics contains the following notable keys:
             {
                 'video': <video_tensor>,
                 'audio': <audio_tensor>,
                 'label': <action_label>,
             }
  
          - "video" is a Tensor of shape (batch, channels, time, height, Width)
          - "audio" is a Tensor of shape (batch, channels, time, 1, frequency)
          - "label" is a Tensor of shape (batch, 1)
  
          The PyTorchVideo models and transforms expect the same input shapes and
          dictionary structure making this function just a matter of unwrapping the dict and
          feeding it through the model/loss.
          """
          sim_matrix = self.sim_matrix(batch)
          d = torch.diag(sim_matrix)
          margin = 0.1
          max_margin = F.relu(margin + sim_matrix - d.view(-1, 1)) + \
                       F.relu(margin + sim_matrix - d.view(1, -1))
          self.log("train_loss", max_margin)
          return max_margin.mean()
  
      def validation_step(self, batch, batch_idx):
          """
          This function is called in the inner loop of the evaluation cycle. For this
          simple example it's mostly the same as the training loop but with a different
          metric name.
          """
          sim_matrix = self.sim_matrix(batch)
          d = torch.diag(sim_matrix)
          margin = 0.1
          max_margin = F.relu(margin + sim_matrix - d.view(-1, 1)) + \
                       F.relu(margin + sim_matrix - d.view(1, -1))
          self.log("val_loss", max_margin)
          return max_margin.mean()
  
      def configure_optimizers(self):
          """
          We use the SGD optimizer with per step cosine annealing scheduler.
          """
          optimizer = torch.optim.SGD(
              self.parameters(),
              lr=self.args.lr,
              momentum=self.args.momentum,
              weight_decay=self.args.weight_decay,
          )
          scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
              optimizer, self.args.max_epochs, last_epoch=-1
          )
          return [optimizer], [scheduler]
  
  class Gated_Embedding_Unit(nn.Module):
      def __init__(self, input_dimension, output_dimension, gating=True):
          super(Gated_Embedding_Unit, self).__init__()
          self.fc = nn.Linear(input_dimension, output_dimension)
          self.cg = Context_Gating(output_dimension)
          self.gating = gating
  
      def forward(self, x):
          x = self.fc(x)
          if self.gating:
              x = self.cg(x)
          x = F.normalize(x)
          return x
  
  class Sentence_Maxpool(nn.Module):
      def __init__(self, word_dimension, output_dim, relu=True):
          super(Sentence_Maxpool, self).__init__()
          self.fc = nn.Linear(word_dimension, output_dim)
          self.out_dim = output_dim
          self.relu = relu
  
      def forward(self, x):
          x = self.fc(x)
          if self.relu:
              x = F.relu(x)
          return torch.max(x, dim=1)[0]
  
  
  class Context_Gating(nn.Module):
      def __init__(self, dimension, add_batch_norm=False):
          super(Context_Gating, self).__init__()
          self.fc = nn.Linear(dimension, dimension)
          self.add_batch_norm = add_batch_norm
          self.batch_norm = nn.BatchNorm1d(dimension)
  
      def forward(self, x):
          x1 = self.fc(x)
          if self.add_batch_norm:
              x1 = self.batch_norm(x1)
          x = torch.cat((x, x1), 1)
          return F.glu(x, 1)
  
  
  class YouCook2DataModule(pytorch_lightning.LightningDataModule):
      """
      This LightningDataModule implementation constructs a PyTorchVideo YouCook2 dataset for both
      the train and val partitions. It defines each partition's augmentation and
      preprocessing transforms and configures the PyTorch DataLoaders.
      """
  
      def __init__(self, args):
          self.args = args
          if args.word2vec_path is not None:
              print('Loading word vectors: {}'.format(args.word2vec_path))
              self.we = KeyedVectors.load_word2vec_format(args.word2vec_path, binary=True)
          else:
              self.we = None
          super().__init__()
  
      def _make_transforms(self, mode: str):
          """
          ##################
          # PTV Transforms #
          ##################
  
          # Each PyTorchVideo dataset has a "transform" arg. This arg takes a
          # Callable[[Dict], Any], and is used on the output Dict of the dataset to
          # define any application specific processing or augmentation. Transforms can
          # either be implemented by the user application or reused from any library
          # that's domain specific to the modality. E.g. for video we recommend using
          # TorchVision, for audio we recommend TorchAudio.
          #
          # To improve interoperation between domain transform libraries, PyTorchVideo
          # provides a dictionary transform API that provides:
          #   - ApplyTransformToKey(key, transform) - applies a transform to specific modality
          #   - RemoveKey(key) - remove a specific modality from the clip
          #
          # In the case that the recommended libraries don't provide transforms that
          # are common enough for PyTorchVideo use cases, PyTorchVideo will provide them in
          # the same structure as the recommended library. E.g. TorchVision didn't
          # have a RandomShortSideScale video transform so it's been added to PyTorchVideo.
          """
          if self.args.data_type == "video":
              transform = [
                  self._video_transform(mode),
                  RemoveKey("audio"),
              ]
          elif self.args.data_type == "audio":
              transform = [
                  self._audio_transform(),
                  RemoveKey("video"),
              ]
          else:
              raise Exception(f"{self.args.data_type} not supported")
  
          return Compose(transform)
  
      def _video_transform(self, mode: str):
          """
          This function contains example transforms using both PyTorchVideo and TorchVision
          in the same Callable. For 'train' mode, we use augmentations (prepended with
          'Random'), for 'val' mode we use the respective determinstic function.
          """
          args = self.args
          return ApplyTransformToKey(
              key="video",
              transform=Compose(
                  [
                      UniformTemporalSubsample(args.video_num_subsampled),
                      # Normalize(args.video_means, args.video_stds),
                  ]
                  + (
                      [
                          RandomShortSideScale(
                              min_size=args.video_min_short_side_scale,
                              max_size=args.video_max_short_side_scale,
                          ),
                          RandomCrop(args.video_crop_size),
                          RandomHorizontalFlip(p=args.video_horizontal_flip_p),
                      ]
                      if mode == "train"
                      else [
                          ShortSideScale(args.video_min_short_side_scale),
                          CenterCrop(args.video_crop_size),
                      ]
                  )
              ),
          )
  
      def _audio_transform(self):
          """
          This function contains example transforms using both PyTorchVideo and TorchAudio
          in the same Callable.
          """
          args = self.args
          n_fft = int(
              float(args.audio_resampled_rate) / 1000 * args.audio_mel_window_size
          )
          hop_length = int(
              float(args.audio_resampled_rate) / 1000 * args.audio_mel_step_size
          )
          eps = 1e-10
          return ApplyTransformToKey(
              key="audio",
              transform=Compose(
                  [
                      Resample(
                          orig_freq=args.audio_raw_sample_rate,
                          new_freq=args.audio_resampled_rate,
                      ),
                      MelSpectrogram(
                          sample_rate=args.audio_resampled_rate,
                          n_fft=n_fft,
                          hop_length=hop_length,
                          n_mels=args.audio_num_mels,
                          center=False,
                      ),
                      Lambda(lambda x: x.clamp(min=eps)),
                      Lambda(torch.log),
                      UniformTemporalSubsample(args.audio_mel_num_subsample),
                      Lambda(lambda x: x.transpose(1, 0)),  # (F, T) -> (T, F)
                      Lambda(
                          lambda x: x.view(1, x.size(0), 1, x.size(1))
                      ),  # (T, F) -> (1, T, 1, F)
                      Normalize((args.audio_logmel_mean,), (args.audio_logmel_std,)),
                  ]
              ),
          )
  
      def train_dataloader(self):
          """
          Defines the train DataLoader that the PyTorch Lightning Trainer trains/tests with.
          """
          sampler = RandomSampler
          train_transform = self._make_transforms(mode="train")
          self.train_dataset = pytorchvideo.data.YouCook2(
                  json_file=self.args.data_path,
                  word2vec=self.we,
                  video_path_prefix=self.args.video_path_prefix,
                  transform=train_transform,
                  video_sampler=sampler,
              )
          return torch.utils.data.DataLoader(
              self.train_dataset,
              batch_size=self.args.batch_size,
              num_workers=self.args.workers,
          )
  
      def val_dataloader(self):
          """
          Defines the train DataLoader that the PyTorch Lightning Trainer trains/tests with.
          """
          sampler = RandomSampler
          val_transform = self._make_transforms(mode="val")
          self.val_dataset = pytorchvideo.data.YouCook2(
              json_file=self.args.data_path,
              word2vec=self.we,
              video_path_prefix=self.args.video_path_prefix,
              transform=val_transform,
              video_sampler=sampler,
          )
          return torch.utils.data.DataLoader(
              self.val_dataset,
              batch_size=self.args.batch_size,
              num_workers=self.args.workers,
          )
  
  
  class LimitDataset(torch.utils.data.Dataset):
      """
      To ensure a constant number of samples are retrieved from the dataset we use this
      LimitDataset wrapper. This is necessary because several of the underlying videos
      may be corrupted while fetching or decoding, however, we always want the same
      number of steps per epoch.
      """
  
      def __init__(self, dataset):
          super().__init__()
          self.dataset = dataset
          self.dataset_iter = itertools.chain.from_iterable(
              itertools.repeat(iter(dataset), 2)
          )
  
      def __getitem__(self, index):
          return next(self.dataset_iter)
  
      def __len__(self):
          return self.dataset.num_videos
  
  
  def main():
      """
      To train the ResNet with the Kinetics dataset we construct the two modules above,
      and pass them to the fit function of a pytorch_lightning.Trainer.
  
      This example can be run either locally (with default parameters) or on a Slurm
      cluster. To run on a Slurm cluster provide the --on_cluster argument.
      """
      setup_logger()
  
      pytorch_lightning.trainer.seed_everything()
      parser = argparse.ArgumentParser()
  
      #  Cluster parameters.
      parser.add_argument("--on_cluster", action="store_true")
      parser.add_argument("--job_name", default="ptv_video_classification", type=str)
      parser.add_argument("--working_directory", default=".", type=str)
      parser.add_argument("--partition", default="dev", type=str)
  
      # Model parameters.
      parser.add_argument("--lr", "--learning-rate", default=0.1, type=float)
      parser.add_argument("--momentum", default=0.9, type=float)
      parser.add_argument("--weight_decay", default=1e-4, type=float)
      parser.add_argument(
          "--arch",
          default="video_resnet",
          choices=["video_resnet", "audio_resnet"],
          type=str,
      )
  
      # Data parameters.
      parser.add_argument("--data_path", default=None, type=str, required=True)
      parser.add_argument("--video_path_prefix", default="", type=str)
      parser.add_argument("--workers", default=8, type=int)
      parser.add_argument("--batch_size", default=32, type=int)
      parser.add_argument("--clip_duration", default=2, type=float)
      parser.add_argument(
          "--data_type", default="video", choices=["video", "audio"], type=str
      )
      parser.add_argument("--video_num_subsampled", default=8, type=int)
      parser.add_argument("--video_means", default=(0.45, 0.45, 0.45), type=tuple)
      parser.add_argument("--video_stds", default=(0.225, 0.225, 0.225), type=tuple)
      parser.add_argument("--video_crop_size", default=224, type=int)
      parser.add_argument("--video_min_short_side_scale", default=256, type=int)
      parser.add_argument("--video_max_short_side_scale", default=320, type=int)
      parser.add_argument("--video_horizontal_flip_p", default=0.5, type=float)
      parser.add_argument("--audio_raw_sample_rate", default=44100, type=int)
      parser.add_argument("--audio_resampled_rate", default=16000, type=int)
      parser.add_argument("--audio_mel_window_size", default=32, type=int)
      parser.add_argument("--audio_mel_step_size", default=16, type=int)
      parser.add_argument("--audio_num_mels", default=80, type=int)
      parser.add_argument("--audio_mel_num_subsample", default=128, type=int)
      parser.add_argument("--audio_logmel_mean", default=-7.03, type=float)
      parser.add_argument("--audio_logmel_std", default=4.66, type=float)
      parser.add_argument("--word2vec_path", default=None, type=str)
      parser.add_argument("--s3d_dict", default="data/s3d_dict.npy", type=str)
      parser.add_argument("--s3d_model", default="data/s3d_howto100m.pth", type=str)
  
      # Trainer parameters.
      parser = pytorch_lightning.Trainer.add_argparse_args(parser)
      parser.set_defaults(
          max_epochs=200,
          callbacks=[LearningRateMonitor()],
          replace_sampler_ddp=False,
          reload_dataloaders_every_epoch=False,
      )
  
      # Build trainer, ResNet lightning-module and Kinetics data-module.
      args = parser.parse_args()
  
      represent(args)
  
  
  def train(args):
      trainer = pytorch_lightning.Trainer.from_argparse_args(args)
      classification_module = S3DGVideoTextRepresentator(args)
      data_module = YouCook2DataModule(args)
      trainer.fit(classification_module, data_module)
  
  def compute_metrics(x):
      sx = np.sort(-x, axis=1)
      d = np.diag(-x)
      d = d[:, np.newaxis]
      ind = sx - d
      ind = np.where(ind == 0)
      ind = ind[1]
      metrics = {}
      metrics['R1'] = float(np.sum(ind == 0)) / len(ind)
      metrics['R5'] = float(np.sum(ind < 5)) / len(ind)
      metrics['R10'] = float(np.sum(ind < 10)) / len(ind)
      metrics['MR'] = np.median(ind) + 1
      return metrics
  
  def get_topk_result(x, topk, text_list, video_list, segment_list):
      topk_ind = np.argsort(-x, axis=1)
      for i in range(len(text_list)):
          print(text_list[i])
          for ind in topk_ind[i][:topk]:
              print(video_list[ind], segment_list[ind], text_list[ind], x[i][ind])
  
  def print_computed_metrics(metrics):
      r1 = metrics['R1']
      r5 = metrics['R5']
      r10 = metrics['R10']
      mr = metrics['MR']
      print('R@1: {:.4f} - R@5: {:.4f} - R@10: {:.4f} - Median R: {}'.format(r1, r5, r10, mr))
  
  def represent(args): 
      representator_module = S3DGVideoTextRepresentator(args)
      data_module = YouCook2DataModule(args)
      video_emb_list = []
      text_emb_list = []
      video_list = []
      segment_list = []
      text_list = []
      for data in data_module.val_dataloader():
          video_emb, text_emb = representator_module(data)
          video_emb_list.append(video_emb)
          text_emb_list.append(text_emb)
          video_list.extend([os.path.join(args.video_path_prefix, video_name)
                              for video_name in data["video_name"]])
          segment_list.extend(torch.stack(data["segment"], 1).numpy().tolist())
          text_list.extend(data["sentence"])
      all_video_emb = torch.cat(video_emb_list)
      all_text_emb = torch.cat(text_emb_list)
      sim_matrix = torch.matmul(all_text_emb, all_video_emb.t())
      m  = sim_matrix.cpu().detach().numpy()
      metrics = compute_metrics(m)
      print_computed_metrics(metrics)
      get_topk_result(m, 3, text_list, video_list, segment_list)
  
  def setup_logger():
      ch = logging.StreamHandler()
      formatter = logging.Formatter("\n%(asctime)s [%(levelname)s] %(name)s: %(message)s")
      ch.setFormatter(formatter)
      logger = logging.getLogger("pytorchvideo")
      logger.setLevel(logging.DEBUG)
      logger.addHandler(ch)
  
  
  if __name__ == "__main__":
      main()
  