'''
liguangyao
10/25/2023
guangyaoli@ruc.edu.cn
'''

import os
import torch
from torchvision import transforms, utils
from PIL import Image
import numpy as np
import glob

from imagebind import data
from imagebind.models import imagebind_model
from imagebind.models.imagebind_model import ModalityType

device = "cuda:1" if torch.cuda.is_available() else "cpu"

# Instantiate model
model = imagebind_model.imagebind_huge(pretrained=True)
model.eval()
model.to(device)

def VideoLevelPrompt(video_label_list, video_name):

    video_level_prompt = 'A photo of a dog.'

    return video_level_prompt

def ImageBind_feat_extract(args, dir_audio_path, dir_viusal_path, dir_text_path, dst_audio_path, dst_visual_path, dst_text_path):

    # 此处为文本
    video_label_list = []
    with open(dir_text_path, 'r') as dpp:
        for line in dpp:
            video_label_list.append(line.replace("\n", ""))
    # print(video_label_list)

    video_list = os.listdir(dir_viusal_path)
    video_idx = 0
    total_nums = len(video_list)

    for video_name in video_list:

        video_idx = video_idx + 1
        print("\n--> ", video_idx, video_name)

        audio_save_file = os.path.join(dst_audio_path, video_name + '.npy')
        frame_save_file = os.path.join(dst_visual_path, video_name + '.npy')
        text_save_file = os.path.join(dst_text_path, video_name + '.npy')

        if os.path.exists(audio_save_file):
            print(video_name + '.npy', "is already processed!")
            continue

        frame_list_load = sorted(glob.glob(os.path.join(dir_viusal_path, video_name, '*.jpg')))
        audio_list_load = sorted(glob.glob(os.path.join(dir_audio_path, video_name, '*.wav')))
        text_list = VideoLevelPrompt(video_label_list, video_name) # 例如：A photo of a dog. 保证文本是陈述语句即可，可自行设计

        # 为了保证模型训练可以批处理，故需要保证每个数据样本后的长度一致。
        # 然而由于不同的视频长度不一，采样出的帧数不一致，故此处对每个视频进行均匀采样。
        frame_nums = len(frame_list_load)
        if frame_nums < args.frame_nums:
            frame_samples = np.round(np.linspace(0, frame_nums-2, args.frame_nums))
        else:
            frame_samples = np.round(np.linspace(0, args.frame_nums-1, args.frame_nums))
        frame_list  = [frame_list_load[int(sample)] for sample in frame_samples]

        audio_nums = len(audio_list_load)
        if audio_nums < args.audio_nums:
            audio_samples = np.round(np.linspace(0, audio_nums-2, args.audio_nums))
        else:
            audio_samples = np.round(np.linspace(0, args.audio_nums-1, args.audio_nums))
        audio_list  = [audio_list_load[int(sample)] for sample in audio_samples]

        # Load data
        inputs = {
            ModalityType.TEXT: data.load_and_transform_text(text_list, device),
            ModalityType.VISION: data.load_and_transform_vision_data(frame_list, device),
            ModalityType.AUDIO: data.load_and_transform_audio_data(audio_list, device),
        }

        with torch.no_grad():
            embeddings = model(inputs)

        text_feat = embeddings['text']
        audio_feat = embeddings['audio']
        visual_feat = embeddings['vision']

        # print("\nimagebind text: ", text_feat.shape)
        # print("imagebind audio: ", audio_feat.shape)
        # print("imagebind visual: ", visual_feat.shape)

        text_feat = text_feat.float().cpu().numpy()
        np.save(text_save_file, text_feat)

        audio_feat = audio_feat.float().cpu().numpy()
        np.save(audio_save_file, audio_feat)

        visual_feat = visual_feat.float().cpu().numpy()
        np.save(frame_save_file, visual_feat)

        print("Process: ", video_idx, " / ", total_nums, " ----- video id: ", video_idx)
        print("T-A-V Feat shape: ", text_feat.shape, audio_feat.shape, visual_feat.shape)

def ImageBind_visaul_feat_extract(args, dir_viusal_path, dst_visual_path):

    video_list = os.listdir(dir_viusal_path)
    video_idx = 0
    total_nums = len(video_list)

    for video_name in video_list:

        video_idx = video_idx + 1
        print("\n--> ", video_idx, video_name)

        frame_save_file = os.path.join(dst_visual_path, video_name + '.npy')

        if os.path.exists(frame_save_file):
            print(video_name + '.npy', "is already processed!")
            continue

        frame_list_load = sorted(glob.glob(os.path.join(dir_viusal_path, video_name, '*.jpg')))

        frame_nums = len(frame_list_load)
        if frame_nums < args.frame_nums:
            frame_samples = np.round(np.linspace(0, frame_nums-2, args.frame_nums))
        else:
            frame_samples = np.round(np.linspace(0, args.frame_nums-1, args.frame_nums))
        frame_list  = [frame_list_load[int(sample)] for sample in frame_samples]

        # Load data
        inputs = {ModalityType.VISION: data.load_and_transform_vision_data(frame_list, device),}

        with torch.no_grad():
            embeddings = model(inputs)

        visual_feat = embeddings['vision']
        # print("imagebind visual: ", visual_feat.shape)

        visual_feat = visual_feat.float().cpu().numpy()
        np.save(frame_save_file, visual_feat)

        print("Process: ", video_idx, " / ", total_nums, " ----- video id: ", video_idx)
        print("V Feat shape: ", visual_feat.shape)

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument("--dir_audio_path", type=str, default='data/users/guangyao_li/MUSIC-AVQA/audio_16kHz_2sec',
                        help='audio file path')
    parser.add_argument("--dir_visual_path", type=str, default='/data/users/guangyao_li/MUSIC-AVQA/avqa-frames-1fps',
                        help='visual frames path')
    parser.add_argument("--dir_text_path", type=str, default='../../dataset/split_que_id/music_avqa.json',
                        help='text file path')

    parser.add_argument("--dst_audio_path", type=str, default='/data/users/guangyao_li/MUSIC-AVQA/imagebind_feat/imagebind_audio_16kHz',
                        help='audio feature path')
    parser.add_argument("--dst_visual_path", type=str, default='/data/users/guangyao_li/MUSIC-AVQA/imagebind_feat/imagebind_frame_1fps',
                        help='visual frames feature path')
    parser.add_argument("--dst_text_path", type=str, default='/data/users/guangyao_li/MUSIC-AVQA/imagebind_feat/imagebind_text',
                        help='text feature path')

    parser.add_argument("--frame_nums", type=int, default=60,
                        help='frame sample numbers')
    parser.add_argument("--audio_nums", type=int, default=60,
                        help='audio clip sample numbers')
    # parser.add_argument("--gpu", dest='gpu', type=str, default='0',
    #                     help='Set CUDA_VISIBLE_DEVICES environment variable, optional')

    args = parser.parse_args()
    # os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
    params = vars(args)

    # 同时提取audio, vsiual 和text的特征
    ImageBind_feat_extract(args, args.dir_audio_path, args.dir_visual_path, args.dir_text_path,
                           args.dst_audio_path, args.dst_visual_path, args.dst_text_path)

    # 只提取一个模态的特征，如visual
    ImageBind_visual_feat_extract(args, dir_viusal_path, dst_visual_path)


import dataclasses
import subprocess
from typing import Any, ClassVar, List, Optional

from fancy_dataclass.utils import DataclassMixin, issubclass_safe, obj_class_name

class SubprocessDataclass(DataclassMixin):
    """Mixin class providing a method for converting dataclass fields to command-line args that can be used to make a subprocess call.

    Other arguments can be passed into the `metadata` argument of a `dataclasses.field`, namely:

    - `exec` (boolean flag indicating that this field should be treated as the name of the executable, rather than an argument)
    - `args` (list of command-line arguments corresponding to the field—only the first will be used, and only if it starts with a hyphen)
    - `exclude` (boolean flag indicating that the field should not be included in the args)"""

    def __post_init__(self) -> None:
        exec_field = None
        for (name, field) in self.__dataclass_fields__.items():
            if field.metadata.get('exec', False):
                if (exec_field is None):
                    exec_field = name
                else:
                    raise TypeError("cannot have more than one field with 'exec' flag set to True")

    def get_arg(self, name: str, suppress_defaults: bool = False) -> List[str]:
        """Given the name of a dataclass field, gets the command-line args for that field.

        Args:
            name: Name of dataclass field
            suppress_defaults: If `True`, suppresses arguments that are equal to the default values

        Returns:
            List of command-line args corresponding to the field"""
        field = self.__dataclass_fields__[name]
        if field.metadata.get('exclude', False):  # exclude the argument
            return []
        if getattr(field.type, '__origin__', None) is ClassVar:
            # ignore fields associated with the class, rather than the instance
            return []
        val = getattr(self, name, None)
        if (val is None):  # optional value is None
            return []
        if issubclass_safe(field.type, SubprocessDataclass):  # get args via nested SubprocessDataclass
            return val.args(suppress_defaults = suppress_defaults)
        if field.metadata.get('exec', False):  # this field is the executable, so return no arguments
            return []
        if suppress_defaults:  # if value matches the default, suppress the argument
            default = None
            has_default = True
            if (field.default == dataclasses.MISSING):
                if (field.default_factory == dataclasses.MISSING):
                    has_default = False
                else:
                    default = field.default_factory()
            else:
                default = field.default
            if has_default and (val == default):
                return []
        if field.metadata.get('args'):  # use arg name provided by the metadata
            arg = field.metadata['args'][0]
            if (not arg.startswith('-')):
                arg = None
        else:  # use the field name (assume a single dash if it is a single letter)
            prefix = '-' if (len(name) == 1) else '--'
            arg = prefix + name.replace('_', '-')
        if isinstance(val, bool):
            # make it a boolean flag if True, otherwise omit it
            if (not val):
                arg = None
            val = []
        elif isinstance(val, (list, tuple)):
            if val:
                val = [str(x) for x in val]
            else:
                arg = None
        elif (val is not None):  # convert the field value to a string
            val = str(val)
        args = [arg] if arg else []
        args += val if isinstance(val, list) else [val]
        return args

    def get_executable(self) -> Optional[str]:
        """Gets the name of an executable to run with the appropriate arguments.

        By default, this returns the name of the first dataclass field whose `exec` metadata flag is set to `True`, if one exists, and `None` otherwise.

        Returns:
            Name of the executable to run"""
        name = None
        for (name, field) in self.__dataclass_fields__.items():
            if field.metadata.get('exec', False):
                return getattr(self, name, None)
        return None

    def args(self, suppress_defaults: bool = False) -> List[str]:
        """Converts dataclass fields to a list of command-line arguments for a subprocess call.

        Args:
            suppress_defaults: If `True`, suppresses arguments that are equal to the default values

        Returns:
            List of command-line args corresponding to the dataclass fields"""
        args = []
        for name in self.__dataclass_fields__:
            args += [arg for arg in self.get_arg(name, suppress_defaults = suppress_defaults) if arg]
        return args

    def run_subprocess(self, **kwargs: Any) -> subprocess.CompletedProcess:
        """Executes the full subprocess command corresponding to the dataclass parameters.

        Args:
            kwargs: Keyword arguments passed to `subprocess.run`

        Returns:
            `CompletedProcess` object produced by `subprocess.run`

        Raises:
            ValueError: If no executable was found from the `get_executable` method"""
        executable = self.get_executable()
        if (not executable):
            raise ValueError(f'No executable identified for use with {obj_class_name(self)!r} instance')
        args = [executable] + self.args()
        return subprocess.run(args, **kwargs)


import wx

class SlideshowFrame(wx.Frame):
    def __init__(self,**kwargs):
        wx.Frame.__init__(self, **kwargs)
        self.SetBackgroundColour(wx.BLACK)

        self.panel = wx.Panel(self, pos=self.Rect.GetPosition(), size=self.Rect.GetSize())
        self.empty_img = wx.EmptyImage(self.Rect.GetWidth(),
                                       self.Rect.GetHeight())
        self.imageCtrl = wx.StaticBitmap(self.panel, wx.ID_ANY,
                                         wx.BitmapFromImage(self.empty_img))
        #self.verSizer = wx.BoxSizer(wx.VERTICAL)
        #self.horSizer = wx.BoxSizer(wx.HORIZONTAL)
        #self.mainSizer.Add(self.imageCtrl, 0, wx.ALL|wx.ALIGN_CENTER, 0)
        #self.panel.SetSizer(self.mainSizer)
        #self.mainSizer.Fit(self)
        #self.panel.Layout()

    def load_img(self, img_path):
        if img_path is None:
            img = self.empty_img
        else:
            img = wx.Image(img_path, wx.BITMAP_TYPE_ANY)

        #
        # scale the image, preserving the aspect ratio
        #

        w = img.GetWidth()
        h = img.GetHeight()

        W = self.Rect.GetWidth()
        H = self.Rect.GetHeight()

        # scale w to match W, and see if height is over/under H. If so, scale
        # h to match H instead.
        w2, h2 = W, h*(float(W)/w)
        if h2 > H:
            w2, h2 = w*(float(H)/h), H

        img = img.Scale(w2,h2,quality=wx.IMAGE_QUALITY_HIGH)
        self.imageCtrl.SetBitmap(wx.BitmapFromImage(img))
        #self.panel.Layout()

        O = self.Rect.GetPosition() # frame origin
        X,Y = (O[0] + (W-w2)/2, O[1] + (H-h2)/2)
        self.panel.SetRect((X,Y,w2,h2))
        #self.mainSizer.Fit(self)
        #self.panel.Layout()
        self.panel.Refresh()


#!/usr/bin/env python3.8
# -*- coding: utf-8 -*-

"""
Created on Wed Feb 27 17:55:12 2023
@author: Carlos Gómez-Huélamo
"""

# General purpose imports

import sys
import os
import pdb
import git

if str(sys.version_info[0])+"."+str(sys.version_info[1]) >= "3.9": # Python >= 3.9
    from math import gcd
else:
    from fractions import gcd

# DL & Math imports

import math
import numpy as np
import torch
import pytorch_lightning as pl

from scipy import sparse

from torch import nn
from torch.nn import functional as F
from torch_geometric.nn import conv
from torch_geometric.utils import from_scipy_sparse_matrix

# Plot imports

# Custom imports

# Global variables

# https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
torch.backends.cudnn.benchmark = True
torch.set_float32_matmul_precision("medium") # highest, high, medium

#######################################

class TMFModel(pl.LightningModule):
    def __init__(self, args):
        super(TMFModel, self).__init__() # allows us to avoid using the base class name explicitly
        self.args = args

        # Save model in log_dir as backup

        self.save_hyperparameters() # It will enable Lightning to store all the provided arguments under the self.hparams attribute.
                                    # These hyperparameters will also be stored within the model checkpoint, which simplifies model re-instantiation after training.

        # Encoder

        ## Social

        self.linear_embedding = LinearEmbedding(3,self.args)
        self.pos_encoder= PositionalEncoding1D(self.args.social_latent_size)
        self.encoder_transformer = EncoderTransformer(self.args)
        self.agent_gnn = AgentGNN(self.args)

        ## Physical

        if self.args.use_map:

            self.map_sub_net = MapSubNet(self.args)

            assert self.args.social_latent_size == self.args.map_latent_size

            if self.args.final_latent_info == "concat":
                self.args.decoder_latent_size = self.args.social_latent_size + self.args.map_latent_size
            elif self.args.final_latent_info == "fuse":
                self.A2L_1 = TransformerDecoder(self.args.social_latent_size, head_num=self.args.num_attention_heads)
                self.L2A_1 = TransformerDecoder(self.args.social_latent_size, head_num=self.args.num_attention_heads)

                self.A2L_2 = TransformerDecoder(self.args.social_latent_size, head_num=self.args.num_attention_heads)
                self.L2A_2 = TransformerDecoder(self.args.social_latent_size, head_num=self.args.num_attention_heads)

                self.args.decoder_latent_size = self.args.social_latent_size
            else:
                raise AssertionError
        else:
            self.args.decoder_latent_size = self.args.social_latent_size

        if self.args.decoder == "decoder_residual": self.decoder = DecoderResidual(self.args)
        elif self.args.decoder == "decoder_temporal": self.decoder = Temporal_Multimodal_Decoder(self.args)

        # Metrics

        self.reg_loss = nn.SmoothL1Loss(reduction="none")
        if self.args.freeze_decoder:
            self.initial_lr_conf = self.args.initial_lr_conf
            self.min_lr_conf = self.args.min_lr_conf
        else:
            self.initial_lr_conf = 1e-3
            self.min_lr_conf = 1e-6

        self.is_frozen = False
        self.save_model_script = True

    @staticmethod
    def init_args(parent_parser, BASE_DIR, DATASET_DIR):
        parser_dataset = parent_parser.add_argument_group("dataset")
        parser_dataset.add_argument(
            "--BASE_DIR", type=str, default=BASE_DIR)
        parser_dataset.add_argument(
            "--DATASET_DIR", type=str, default=DATASET_DIR)
        parser_dataset.add_argument(
            "--LOG_DIR", type=str, default="non_specified")
        parser_dataset.add_argument(
            "--train_split", type=str, default=os.path.join(
                BASE_DIR, DATASET_DIR, "train"))
        parser_dataset.add_argument(
            "--val_split", type=str, default=os.path.join(
                BASE_DIR, DATASET_DIR, "val"))
        parser_dataset.add_argument(
            "--test_split", type=str, default=os.path.join(
                BASE_DIR, DATASET_DIR, "test"))

        # Social preprocess

        parser_dataset.add_argument(
            "--train_split_pre_social", type=str, default=os.path.join(
                BASE_DIR, DATASET_DIR, "processed_social", "train_pre_clean.pkl"))
        parser_dataset.add_argument(
            "--val_split_pre_social", type=str, default=os.path.join(
                BASE_DIR, DATASET_DIR, "processed_social", "val_pre_clean.pkl"))
        parser_dataset.add_argument(
            "--test_split_pre_social", type=str, default=os.path.join(
                BASE_DIR, DATASET_DIR, "processed_social", "test_pre_clean.pkl"))

        # Map preprocess

        parser_dataset.add_argument(
            "--train_split_pre_map", type=str, default=os.path.join(
            BASE_DIR, DATASET_DIR, "processed_map", "train_map_data_rot_right_x_multi_agent.pkl"))
        parser_dataset.add_argument(
            "--val_split_pre_map", type=str, default=os.path.join(
                BASE_DIR, DATASET_DIR, "processed_map", "val_map_data_rot_right_x_multi_agent.pkl"))
        parser_dataset.add_argument(
            "--test_split_pre_map", type=str, default=os.path.join(
                BASE_DIR, DATASET_DIR, "processed_map", "test_map_data_rot_right_x_multi_agent.pkl"))

        # Whole preprocess

        parser_dataset.add_argument(
            "--train_split_pre", type=str, default=os.path.join(
            BASE_DIR, DATASET_DIR, "processed_full", "train_full_data_rot_right_x_multi_agent.pkl"))
        parser_dataset.add_argument(
            "--val_split_pre", type=str, default=os.path.join(
                BASE_DIR, DATASET_DIR, "processed_full", "val_full_data_rot_right_x_multi_agent.pkl"))
        parser_dataset.add_argument(
            "--test_split_pre", type=str, default=os.path.join(
                BASE_DIR, DATASET_DIR, "processed_full", "test_full_data_rot_right_x_multi_agent.pkl"))

        parser_dataset.add_argument("--reduce_dataset_size", type=int, default=0)
        parser_dataset.add_argument("--use_preprocessed", type=bool, default=False)
        parser_dataset.add_argument("--use_map", type=bool, default=False)
        parser_dataset.add_argument("--align_image_with_target_x", type=bool, default=True)

        parser_training = parent_parser.add_argument_group("training")
        parser_training.add_argument("--num_epochs", type=int, default=200)
        parser_training.add_argument("--check_val_every_n_epoch", type=int, default=10)
        parser_training.add_argument("--lr_values", type=list, default=[1e-3, 1e-4, 1e-3 , 1e-4])
        parser_training.add_argument("--lr_step_epochs", type=list, default=[10, 20, 45])
        parser_training.add_argument("--initial_lr_conf", type=float, default=5e-5)
        parser_training.add_argument("--min_lr_conf", type=float, default=1e-6)
        parser_training.add_argument("--wd", type=float, default=0.001)
        parser_training.add_argument("--batch_size", type=int, default=128)
        parser_training.add_argument("--val_batch_size", type=int, default=128)
        parser_training.add_argument("--workers", type=int, default=0) # TODO: Not working with >= 0
        parser_training.add_argument("--val_workers", type=int, default=0)
        parser_training.add_argument("--gpus", type=int, default=1)

        parser_model = parent_parser.add_argument_group("model")
        parser_dataset.add_argument("--MODEL_DIR", type=str, default="non_specified")
        parser_model.add_argument("--data_dim", type=int, default=2)
        parser_model.add_argument("--obs_len", type=int, default=50)
        parser_model.add_argument("--pred_len", type=int, default=60)
        parser_model.add_argument("--centerline_length", type=int, default=40)
        parser_model.add_argument("--num_centerlines", type=int, default=6)
        parser_model.add_argument("--num_attention_heads", type=int, default=8)
        parser_model.add_argument("--apply_dropout", type=float, default=0.2)
        parser_model.add_argument("--data_aug_gaussian_noise", type=float, default=0.01)
        parser_model.add_argument("--social_latent_size", type=int, default=64)
        parser_model.add_argument("--map_latent_size", type=int, default=64)
        parser_model.add_argument("--final_latent_info", type=str, default="non_specified")
        parser_model.add_argument("--decoder_latent_size", type=int, default=-1)
        parser_model.add_argument("--decoder_temporal_window_size", type=int, default=30) # 49
        parser_model.add_argument("--num_modes", type=int, default=6)
        parser_model.add_argument("--freeze_decoder", type=bool, default=False)
        parser_model.add_argument("--mod_steps", type=list, default=[1, 5]) # First unimodal -> Freeze -> Multimodal
        parser_model.add_argument("--mod_freeze_epoch", type=int, default=20)
        parser_model.add_argument("--mod_full_unfreeze_epoch", type=int, default=60)
        parser_model.add_argument("--reg_loss_weight", type=float, default=1) # xy predictions
        parser_model.add_argument("--cls_loss_weight", type=float, default=1) # classification = confidences
        parser_model.add_argument("--epsilon", type=float, default=0.0000001)

        return parent_parser

    def add_noise(self, input, factor=1):
        """_summary_
        Args:
            input (_type_): _description_
            factor (int, optional): _description_. Defaults to 1.
        Returns:
            _type_: _description_
        """

        noise = factor * torch.randn(input.shape).to(input)
        noisy_input = input + noise
        return noisy_input

    def forward(self, batch):
        # Set batch norm to eval mode in order to prevent updates on the running means,
        # if the weights are frozen
        if self.args.freeze_decoder:
            if self.is_frozen:
                for module in self.modules():
                    if isinstance(module, torch.nn.modules.BatchNorm1d):
                        module.eval()

        # Encoder

        ## Social

        ### Extract the social features in each sample of the current batch
        pdb.set_trace()
        displ, centers = batch["displ"], batch["centers"]
        rotation, origin = batch["rotation"], batch["origin"]
        agents_per_sample = [x.shape[0] for x in displ]
        batch_size = len(agents_per_sample)

        ### OBS: For each sequence, we always set the focal (target) agent as the first agent
        ###      of the scene, then our ego-vehicle (AV) and finally the remanining agents
        ###      (See extractor_proc.py preprocessing)

        focal_agent_id = np.cumsum(agents_per_sample)
        focal_agent_id = np.roll(focal_agent_id,1)
        focal_agent_id[0] = 0

        ### Convert the list of tensors to tensors

        displ_cat = torch.cat(displ, dim=0)
        centers_cat = torch.cat(centers, dim=0)

        ### Data augmentation (TODO: It should be in collate_fn_dict, in the DataLoader)

        if self.training:
            displ_cat[:,:,:2] = self.add_noise(displ_cat[:,:,:2], self.args.data_aug_gaussian_noise)
            centers_cat = self.add_noise(centers_cat, self.args.data_aug_gaussian_noise)

        linear_output = self.linear_embedding(displ_cat)
        pos_encoding = self.pos_encoder(linear_output)
        pos_encoding = pos_encoding + linear_output

        out_transformer = self.encoder_transformer(pos_encoding, agents_per_sample)
        out_agent_gnn = self.agent_gnn(out_transformer, centers_cat, agents_per_sample)

        social_info = torch.stack([x[0] for x in out_agent_gnn])

        if torch.any(torch.isnan(social_info)):
            pdb.set_trace()

        ## Physical

        if self.args.use_map:

            ### Get relevant centerlines (non-padded) per scenario

            rel_candidate_centerlines = batch["rel_candidate_centerlines"]
            rel_candidate_centerlines = torch.stack(rel_candidate_centerlines,dim=0)

            # Data augmentation (TODO: It should be in collate_fn_dict, in the DataLoader)

            # if self.training:
            #     rel_candidate_centerlines = self.add_noise(rel_candidate_centerlines, self.args.data_aug_gaussian_noise)

            ### Get the map latent vector associated

            _, num_centerlines, points_centerline, data_dim = rel_candidate_centerlines.shape
            rel_candidate_centerlines = rel_candidate_centerlines.contiguous().view(-1, points_centerline, data_dim)

            non_empty_mask = rel_candidate_centerlines.abs().sum(dim=1).sum(dim=1) # A padded-centerline must sum 0.0
            # in each dimension, and after that both dimensions together
            rows_mask = torch.where(non_empty_mask == 0.0)[0]
            non_masked_centerlines = rel_candidate_centerlines.shape[0] - len(rows_mask)

            rel_candidate_centerlines_mask = torch.zeros([rel_candidate_centerlines.shape[0]], device=rel_candidate_centerlines.device).type(torch.bool) # False
            rel_candidate_centerlines_mask[rows_mask] = True # Padded-centerlines
            rel_candidate_centerlines_mask_inverted = ~rel_candidate_centerlines_mask # Non-padded centerlines (so, relevant) to True

            centerlines_per_sample = [] # Relevant centerlines (non-padded) per sequence
            num_current_centerlines = 0

            for i in range(rel_candidate_centerlines_mask.shape[0]+1):
                if i % self.args.num_centerlines == 0 and i > 0: # Next traffic scenario
                    centerlines_per_sample.append(num_current_centerlines)
                    num_current_centerlines = 0

                    if i == rel_candidate_centerlines_mask.shape[0]:
                        break
                if rel_candidate_centerlines_mask_inverted[i]: # Non-masked
                    num_current_centerlines += 1

            assert non_masked_centerlines == sum(centerlines_per_sample), \
                "The number of relevant centerlines do not match"

            centerlines_per_sample = np.array(centerlines_per_sample)
            rel_candidate_centerlines_ = rel_candidate_centerlines[rel_candidate_centerlines_mask_inverted,:,:]
            rel_candidate_centerlines_mask_ = rel_candidate_centerlines_mask.reshape(-1,1).repeat_interleave(points_centerline,dim=1)

            physical_info = self.map_sub_net(rel_candidate_centerlines, rel_candidate_centerlines_mask_)

        # Decoder

        if self.args.use_map:
            if self.args.final_latent_info == "concat": # Concat info
                merged_info = torch.cat([social_info,
                                        physical_info],
                                        dim=1)

            if self.args.final_latent_info == "fuse": # Fuse info
                physical_info = physical_info + self.A2L_1(physical_info, social_info)
                social_info = social_info + self.L2A_1(social_info, physical_info)

                physical_info = physical_info + self.A2L_2(physical_info, social_info)
                social_info = social_info + self.L2A_2(social_info, physical_info)

                merged_info = social_info
        else:
            merged_info = social_info

        if torch.any(torch.isnan(merged_info)):
            pdb.set_trace()

        # If self.args.freeze_decoder is set to True, conf are useless

        if self.args.decoder == "decoder_residual":
            pred_traj, conf = self.decoder(merged_info, self.is_frozen, self.current_epoch)
        elif self.args.decoder == "decoder_temporal":
            traj_agent_abs_rel = displ_cat[focal_agent_id,:self.args.decoder_temporal_window_size,:self.args.data_dim]
            last_obs_agent = centers_cat[focal_agent_id,:]

            decoder_h = merged_info.unsqueeze(0)
            decoder_c = torch.zeros(tuple(decoder_h.shape)).to(decoder_h)
            state_tuple = (decoder_h, decoder_c)

            pred_traj_rel, conf = self.decoder(traj_agent_abs_rel, state_tuple)

            # Convert relative displacements to absolute coordinates (around origin)

            pred_traj = relative_to_abs_multimodal(pred_traj_rel, last_obs_agent)

        ### In this model we are only predicting
        ### the focal agent. We would actually
        ### have batch_size x num_agents x num_modes x pred_len x data_dim

        num_agents = 1
        out = pred_traj.contiguous().view(batch_size, num_agents, -1, self.args.pred_len, self.args.data_dim)
        if not self.args.freeze_decoder: conf = conf.view(batch_size, num_agents, -1)

        # Iterate over each batch and transform predictions into the global coordinate frame

        for i in range(len(out)):
            out[i] = torch.matmul(out[i], rotation[i]) + origin[i].view(
                1, 1, 1, -1
            )
        return out, conf

    # Aux class functions

    def freeze(self):
        for param in self.parameters():
            param.requires_grad = False

        self.decoder.unfreeze_layers()

        self.is_frozen = True

    def full_unfreeze(self):
        for param in self.parameters():
            param.requires_grad = True

        self.is_frozen = False

    def prediction_loss(self, preds, gts, conf=None):
        """_summary_

        Args:
            preds (torch.tensor): batch_size x num_agents x num_modes x pred_len x data_dim
                                  OBS: At this moment, num_agents = 1 since we are only predicting the focal agent
            gts (list): list of gt of each scenario (num_agents x pred_len x 2)
            conf (torch.tensor): batch_size x num_agents x 1

        Returns:
            _type_: _description_
        """

        if self.args.freeze_decoder:
            # # Stack all the predicted trajectories of the target agent
            # num_mods = preds.shape[2]
            # # [0] is required to remove the unneeded dimensions
            # preds = torch.cat([x[0] for x in preds], 0)

            # # Stack all the true trajectories of the target agent
            # # Keep in mind, that there are multiple trajectories in each sample,
            # # but only the first one ([0]) corresponds to the target agent

            # gt_target = torch.cat([torch.unsqueeze(x[0], 0) for x in gts], 0)
            # gt_target = torch.repeat_interleave(gt_target, num_mods, dim=0) # repeate the gt for all ks

            # loss_single = self.reg_loss(preds, gt_target)
            # loss_single = torch.sum(torch.sum(loss_single, dim=2), dim=1)
            # loss_single = torch.split(loss_single, num_mods)

            # # Tuple to tensor

            # loss_single = torch.stack(list(loss_single), dim=0)

            # min_loss_index = torch.argmin(loss_single, dim=1) # Get best mode
            # min_loss_combined = [x[min_loss_index[i]] for i, x in enumerate(loss_single)]

            # loss_out = torch.sum(torch.stack(min_loss_combined))
            # # loss_out = torch.mean(torch.stack(min_loss_combined))

            # return loss_out

            # Stack all the predicted trajectories of the target agent

            preds = preds.squeeze(1)

            batch_size, num_modes, pred_len, data_dim = preds.shape

            # Stack all the true trajectories of the target agent
            # Keep in mind, that there are multiple trajectories in each sample, but only the first one ([0]) corresponds
            # to the target agent

            gt_target = torch.cat([torch.unsqueeze(x[0], 0) for x in gts], 0) # batch_size x pred_len x data_dim
            gt_target_repeated = gt_target.unsqueeze(1).repeat(1,preds.shape[1],1,1) # repeate the gt for all ks
                                                                                # batch_size x num_modes x pred_len x data_dim

            fde_k = torch.sqrt((preds[:, :, -1, 0] - gt_target_repeated[:, :, -1, 0]) ** 2 + # x
                               (preds[:, :, -1, 1] - gt_target_repeated[:, :, -1, 1]) ** 2 + # y
                               self.args.epsilon) # to avoid division by zero
            k_hat = torch.argmin(fde_k, dim=1)

            index = torch.tensor(range(preds.shape[0]), dtype=torch.long)
            pred_fut_traj = preds[index, k_hat] # Best trajectory in terms of FDE per sequence

            batch_size, pred_len, _ = pred_fut_traj.shape
            num_modes = preds.shape[1]

            # Regression loss

            # reg_loss = torch.zeros(1, dtype=torch.float32).to(preds)

            mse_loss = F.mse_loss(pred_fut_traj, gt_target, reduction='none')
            mse_loss = mse_loss.sum(dim=2) + self.args.epsilon # sum epsilon to avoid division by zero
            mse_loss = torch.sqrt(mse_loss)
            mse_loss = mse_loss.mean(dim=1)

            fde_loss = fde_k[index, k_hat]

            reg_loss = mse_loss * 0.5 + fde_loss * 0.5
            reg_loss = reg_loss.mean()

            return reg_loss
        else:
            # Stack all the predicted trajectories of the target agent

            preds = preds.squeeze(1)
            conf = conf.squeeze(1)
            batch_size, num_modes, pred_len, data_dim = preds.shape

            # Stack all the true trajectories of the target agent
            # Keep in mind, that there are multiple trajectories in each sample, but only the first one ([0]) corresponds
            # to the target agent

            gt_target = torch.cat([torch.unsqueeze(x[0], 0) for x in gts], 0) # batch_size x pred_len x data_dim
            gt_target_repeated = gt_target.unsqueeze(1).repeat(1,preds.shape[1],1,1) # repeate the gt for all ks
                                                                                # batch_size x num_modes x pred_len x data_dim

            fde_k = torch.sqrt((preds[:, :, -1, 0] - gt_target_repeated[:, :, -1, 0]) ** 2 + # x
                               (preds[:, :, -1, 1] - gt_target_repeated[:, :, -1, 1]) ** 2 + # y
                               self.args.epsilon) # to avoid division by zero
            k_hat = torch.argmin(fde_k, dim=1)

            index = torch.tensor(range(preds.shape[0]), dtype=torch.long)
            pred_fut_traj = preds[index, k_hat] # Best trajectory in terms of FDE per sequence

            batch_size, pred_len, _ = pred_fut_traj.shape
            num_modes = preds.shape[1]

            # Regression loss

            # reg_loss = torch.zeros(1, dtype=torch.float32).to(preds)

            mse_loss = F.mse_loss(pred_fut_traj, gt_target, reduction='none')
            mse_loss = mse_loss.sum(dim=2) + self.args.epsilon # sum epsilon to avoid division by zero
            mse_loss = torch.sqrt(mse_loss)
            mse_loss = mse_loss.mean(dim=1)

            fde_loss = fde_k[index, k_hat]

            reg_loss = mse_loss * 0.5 + fde_loss * 0.5

            reg_loss = reg_loss.mean()

            # Classification loss (max-margin)

            score_hat = conf[index, k_hat].unsqueeze(-1)
            score_hat = score_hat.repeat(1, num_modes)
            cls_loss = conf + 0.2 - score_hat
            cls_loss[cls_loss < 0] = 0
            cls_loss = cls_loss.sum(dim=-1).sum(dim=-1)
            cls_loss = cls_loss /((num_modes-1) * batch_size)

            # Final loss

            loss = reg_loss * self.args.reg_loss_weight + \
                   cls_loss * self.args.cls_loss_weight

            return loss

    def get_lr(self, epoch):
        lr_index = 0
        for lr_epoch in self.args.lr_step_epochs:
            if epoch < lr_epoch:
                break
            lr_index += 1
        return self.args.lr_values[lr_index]

    def get_best_predictions(self, pred, best_pred_indeces):
        """
        pred: batch_size x num_modes x pred_len x data_dim
        best_pred_indeces: batch_size x 1
        Take the best prediction (best mode) according to the best confidence for each sequence
        """
        return pred[torch.arange(pred.shape[0]), best_pred_indeces, :, :].squeeze()

    def calc_prediction_metrics(self, preds, gts, conf=None):
        if self.args.freeze_decoder:
            # Calculate prediction error for each mode
            # Output has shape (batch_size, n_modes, n_timesteps)
            error_per_t = np.linalg.norm(preds - np.expand_dims(gts, axis=1), axis=-1)

            # Calculate the error for the first mode (at index 0)
            fde_1 = np.average(error_per_t[:, 0, -1])
            ade_1 = np.average(error_per_t[:, 0, :])

            # Calculate the error for all modes
            # Best mode is always the one with the lowest final displacement
            lowest_final_error_indices = np.argmin(error_per_t[:, :, -1], axis=1)
            error_per_t = error_per_t[np.arange(
                preds.shape[0]), lowest_final_error_indices]
            fde = np.average(error_per_t[:, -1])
            ade = np.average(error_per_t[:, :])
        else:
            # Calculate prediction error for each mode

            # K = 1
            # Calculate the error for the theoretically best mode (that with the highest confidence)

            best_pred_traj_indeces = conf.argmax(1)
            k1_predictions = self.get_best_predictions(preds,best_pred_traj_indeces)
            error_per_t_k1 = np.linalg.norm(k1_predictions - gts, axis=-1)

            fde_1 = np.average(error_per_t_k1[:, -1])
            ade_1 = np.average(error_per_t_k1[:, :])

            # K = 6
            # Calculate the error for all modes
            # Best mode is always the one with the lowest final displacement
            error_per_t = np.linalg.norm(preds - np.expand_dims(gts, axis=1), axis=-1)
            lowest_final_error_indices = np.argmin(error_per_t[:, :, -1], axis=1)
            error_per_t = error_per_t[np.arange(
                preds.shape[0]), lowest_final_error_indices]
            fde = np.average(error_per_t[:, -1])
            ade = np.average(error_per_t[:, :])

        return ade_1, fde_1, ade, fde

    # Overwrite Pytorch-Lightning functions

    def configure_optimizers(self):
        if self.args.freeze_decoder:
            if self.current_epoch == self.args.mod_freeze_epoch:
                optimizer = torch.optim.AdamW(
                    filter(lambda p: p.requires_grad, self.parameters()), weight_decay=self.args.wd) # Apply optimizer just to those parameters
                                                                                # that require to be trained
            else:
                optimizer = torch.optim.AdamW(
                    self.parameters(), weight_decay=self.args.wd)
            return optimizer
        else:
            optimizer = torch.optim.AdamW(self.parameters(),
                                          weight_decay=self.args.wd,
                                          lr=self.initial_lr_conf)

            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
                                                                   mode='min',
                                                                   factor=0.5,
                                                                   patience=5,
                                                                   min_lr=self.min_lr_conf,
                                                                   verbose=True)

            return {"optimizer": optimizer, "lr_scheduler": scheduler, "monitor": "ade_val"}

    def on_train_epoch_start(self):
        if self.args.freeze_decoder:
            # Trigger weight freeze and optimizer reinit on mod_freeze_epoch
            if self.current_epoch == self.args.mod_freeze_epoch:
                self.freeze()
                self.trainer.strategy.setup_optimizers(self.trainer)

            if self.current_epoch == self.args.mod_full_unfreeze_epoch:
                self.args.freeze_decoder = False
                self.full_unfreeze()
                self.trainer.strategy.setup_optimizers(self.trainer)

            # Set learning rate according to current epoch
            for single_param in self.optimizers().param_groups:
                single_param["lr"] = self.get_lr(self.current_epoch)
                self.log("lr", single_param["lr"], prog_bar=True, sync_dist=True)
        else:
            # Get learning rate according to current epoch
            for single_param in self.optimizers().param_groups:
                self.log("lr", single_param["lr"], prog_bar=True, sync_dist=True)

    def training_step(self, train_batch, batch_idx):
        out, conf = self.forward(train_batch)
        loss = self.prediction_loss(out, train_batch["gt"], conf)
        self.log("loss_train", loss, sync_dist=True)

        return loss

    def validation_step(self, val_batch, batch_idx):
        out, conf = self.forward(val_batch)
        loss = self.prediction_loss(out, val_batch["gt"], conf)
        self.log("loss_val", loss, sync_dist=True)

        # Extract target agent only

        pred = [x[0].detach().cpu().numpy() for x in out]
        gt = [x[0].detach().cpu().numpy() for x in val_batch["gt"]]
        if not self.args.freeze_decoder: conf = [x[0].detach().cpu().numpy() for x in conf]

        # if self.save_model_script:
        #     model_filename = os.path.join(self.args.BASE_DIR,
        #                                   self.args.MODEL_DIR,
        #                                   "TFMF_TGR.py")
        #     os.system(f"cp {model_filename} {self.args.LOG_DIR}")
        #     self.save_model_script = False

        return {"predictions": pred,
                "groundtruth": gt,
                "confidences": conf} # = validation_outputs

    def validation_epoch_end(self, validation_outputs):
        # Extract predictions
        pred = [out["predictions"] for out in validation_outputs]
        pred = np.concatenate(pred, 0) # get predictions along all validation steps
        gt = [out["groundtruth"] for out in validation_outputs]
        gt = np.concatenate(gt, 0) # get ground-truth along all validation steps

        if self.args.freeze_decoder:
            conf = None
        else:
            conf = [out["confidences"] for out in validation_outputs]
            conf = np.concatenate(conf, 0) # get confidences along all validation steps

        ade1, fde1, ade, fde = self.calc_prediction_metrics(pred, gt, conf)
        self.log("ade1_val", ade1, prog_bar=True, sync_dist=True)
        self.log("fde1_val", fde1, prog_bar=True, sync_dist=True)
        self.log("ade_val", ade, prog_bar=True, sync_dist=True)
        self.log("fde_val", fde, prog_bar=True, sync_dist=True)

# Layers

class LinearEmbedding(nn.Module):
    def __init__(self,input_size,args):
        super(LinearEmbedding, self).__init__()
        self.args = args
        self.input_size = input_size
        self.output_size = args.social_latent_size

        self.encoder_input_layer = nn.Linear(
                in_features=self.input_size,
                out_features=self.output_size
                    )
    def forward(self,linear_input):

        linear_out = F.relu(self.encoder_input_layer(linear_input))

        return linear_out

class PositionalEncoding1D(nn.Module):
    def __init__(self, channels):
        """
        :param channels: The last dimension of the tensor you want to apply pos emb to.
        """
        super(PositionalEncoding1D, self).__init__()
        self.org_channels = channels
        channels = int(np.ceil(channels / 2) * 2)
        self.channels = channels
        inv_freq = 1.0 / (10000 ** (torch.arange(0, channels, 2).float() / channels))
        self.register_buffer("inv_freq", inv_freq)
        self.cached_penc = None

    def forward(self, tensor):
        """
        :param tensor: A 3d tensor of size (batch_size, x, ch)
        :return: Positional Encoding Matrix of size (batch_size, x, ch)
        """
        if len(tensor.shape) != 3:
            raise RuntimeError("The input tensor has to be 3d!")

        if self.cached_penc is not None and self.cached_penc.shape == tensor.shape:
            return self.cached_penc

        self.cached_penc = None
        batch_size, x, orig_ch = tensor.shape
        pos_x = torch.arange(x, device=tensor.device).type(self.inv_freq.type())
        sin_inp_x = torch.einsum("i,j->ij", pos_x, self.inv_freq)
        emb_x = torch.cat((sin_inp_x.sin(), sin_inp_x.cos()), dim=-1)
        emb = torch.zeros((x, self.channels), device=tensor.device).type(tensor.type())
        emb[:, : self.channels] = emb_x

        self.cached_penc = emb[None, :, :orig_ch].repeat(batch_size, 1, 1)

        return self.cached_penc

class EncoderTransformer(nn.Module):
    def __init__(self, args):
        super(EncoderTransformer, self).__init__()
        self.args = args

        self.d_model = self.args.social_latent_size # embedding dimension
        # self.nhead = self.args.num_attention_heads # TODO: Is this correct?
        self.nhead = self.args.social_latent_size
        self.d_hid = 1 ## dimension of the feedforward network model in nn.TransformerEncoder
        self.num_layers = 1
        self.dropout = self.args.apply_dropout

        self.encoder_layer = nn.TransformerEncoderLayer(self.d_model, self.nhead, self.d_hid , self.dropout, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=self.num_layers)

    def forward(self, transformer_in, agents_per_sample):

        transformer_out = F.relu(self.transformer_encoder(transformer_in))
        return transformer_out[:,-1,:]

class AgentGNN(nn.Module):
    def __init__(self, args):
        super(AgentGNN, self).__init__()
        self.args = args
        self.latent_size = args.social_latent_size

        self.gcn1 = conv.CGConv(self.latent_size, dim=2, batch_norm=True)
        self.gcn2 = conv.CGConv(self.latent_size, dim=2, batch_norm=True)

    def forward(self, gnn_in, centers, agents_per_sample):
        # gnn_in is a batch and has the shape (batch_size, number_of_agents, latent_size)

        x, edge_index = gnn_in, self.build_fully_connected_edge_idx(
            agents_per_sample).to(gnn_in.device)
        edge_attr = self.build_edge_attr(edge_index, centers).to(gnn_in.device)

        x = F.relu(self.gcn1(x, edge_index, edge_attr))
        gnn_out = F.relu(self.gcn2(x, edge_index, edge_attr))

        edge_index_out1 = []
        for i in agents_per_sample:
            edge_index_out1.append(gnn_out[0:i,:])
            gnn_out = gnn_out[i:,:]

        return edge_index_out1

    def build_fully_connected_edge_idx(self, agents_per_sample):
        edge_index = []

        # In the for loop one subgraph is built (no self edges!)
        # The subgraph gets offsetted and the full graph over all samples in the batch
        # gets appended with the offsetted subgrah
        offset = 0
        for i in range(len(agents_per_sample)):

            num_nodes = agents_per_sample[i]

            adj_matrix = torch.ones((num_nodes, num_nodes))
            adj_matrix = adj_matrix.fill_diagonal_(0)

            sparse_matrix = sparse.csr_matrix(adj_matrix.numpy())
            edge_index_subgraph, _ = from_scipy_sparse_matrix(sparse_matrix)

            # Offset the list
            edge_index_subgraph = torch.Tensor(
                np.asarray(edge_index_subgraph) + offset)
            offset += agents_per_sample[i]

            edge_index.append(edge_index_subgraph)

        # Concat the single subgraphs into one
        edge_index = torch.LongTensor(np.column_stack(edge_index))

        return edge_index
    def build_edge_attr(self, edge_index, data):
        edge_attr = torch.zeros((edge_index.shape[-1], 2), dtype=torch.float)

        rows, cols = edge_index
        # goal - origin
        edge_attr = data[cols] - data[rows]

        return edge_attr

class DecoderResidual(nn.Module):
    def __init__(self, args):
        super(DecoderResidual, self).__init__()

        self.args = args
        self.latent_size = self.args.decoder_latent_size
        self.num_modes = self.args.num_modes

        output = []
        for i in range(sum(args.mod_steps)):
            output.append(PredictionNet(args))

        self.output = nn.ModuleList(output) # is just like a Python list. It was designed to store any desired number of nn.Module’s

        if not self.args.freeze_decoder or self.args.mod_full_unfreeze_epoch != -1:
            # Classification

            norm = "BN"
            ng = 1

            self.latent_predictions = nn.Linear(self.args.num_modes * self.args.pred_len * self.args.data_dim,
                                                self.latent_size)
            self.confidences = nn.Sequential(LinearRes(self.latent_size*2, self.latent_size*2, norm=norm, ng=ng),
                                             nn.Linear(self.latent_size*2, self.num_modes))

    def forward(self, decoder_in, is_frozen, current_epoch):
        batch_size = decoder_in.shape[0]

        if self.args.freeze_decoder:
            sample_wise_out = []

            if self.training is False: # If you are validating or test, use all decoders
                for out_subnet in self.output:
                    sample_wise_out.append(out_subnet(decoder_in))
            elif is_frozen: # If the first decoder has been frozen, decode and train the remaining ones
                for i in range(self.args.mod_steps[0], sum(self.args.mod_steps)):
                    sample_wise_out.append(self.output[i](decoder_in))
            else: # If you are training and is_frozen = False, use only the first decoder
                sample_wise_out.append(self.output[0](decoder_in))

            decoder_out = torch.stack(sample_wise_out)
            decoder_out = torch.swapaxes(decoder_out, 0, 1)

            return decoder_out, []
        else:
            sample_wise_out = []

            for out_subnet in self.output:
                sample_wise_out.append(out_subnet(decoder_in))

            decoder_out = torch.stack(sample_wise_out)
            decoder_out = torch.swapaxes(decoder_out, 0, 1)

            latent_predictions = self.latent_predictions(decoder_out.contiguous().view(batch_size,-1))
            conf_latent = torch.cat([decoder_in,
                                     latent_predictions],
                                     dim=1)

            conf = self.confidences(conf_latent)
            conf = torch.softmax(conf.view(batch_size,-1), dim=1) # batch_size, num_modes
            if not torch.allclose(torch.sum(conf, dim=1), conf.new_ones((batch_size,))):
                pdb.set_trace()

            return decoder_out, conf

    def unfreeze_layers(self):
        for layer in range(self.args.mod_steps[0], sum(self.args.mod_steps)): # Unfreeze all decoders except the first one
            for param in self.output[layer].parameters():
                param.requires_grad = True

class LinearRes(nn.Module):
    def __init__(self, n_in, n_out, norm='GN', ng=32):
        super(LinearRes, self).__init__()
        assert(norm in ['GN', 'BN', 'SyncBN'])

        self.linear1 = nn.Linear(n_in, n_out)
        self.linear2 = nn.Linear(n_out, n_out)
        self.linear3 = nn.Linear(n_out, n_out)
        self.relu = nn.ReLU(inplace=True)

        if norm == 'GN':
            self.norm1 = nn.GroupNorm(gcd(ng, n_out), n_out)
            self.norm2 = nn.GroupNorm(gcd(ng, n_out), n_out)
        elif norm == 'BN':
            self.norm1 = nn.BatchNorm1d(n_out)
            self.norm2 = nn.BatchNorm1d(n_out)
            self.norm3 = nn.BatchNorm1d(n_out)
        else:
            exit('SyncBN has not been added!')

        if n_in != n_out:
            if norm == 'GN':
                self.transform = nn.Sequential(
                    nn.Linear(n_in, n_out, bias=False),
                    nn.GroupNorm(gcd(ng, n_out), n_out))
            elif norm == 'BN':
                self.transform = nn.Sequential(
                    nn.Linear(n_in, n_out, bias=False),
                    nn.BatchNorm1d(n_out))
            else:
                exit('SyncBN has not been added!')
        else:
            self.transform = None

    def forward(self, x):
        out = self.linear1(x)
        out = self.norm1(out)
        out = self.relu(out)
        out = self.linear2(out)
        out = self.norm2(out)
        out = self.relu(out)
        out = self.linear3(out)
        out = self.norm3(out)

        if self.transform is not None:
            out += self.transform(x)
        else:
            out += x

        out = self.relu(out)
        return out

class PredictionNet(nn.Module):
    def __init__(self, args):
        super(PredictionNet, self).__init__()

        self.args = args

        self.latent_size = args.decoder_latent_size

        self.weight1 = nn.Linear(self.latent_size, self.latent_size)
        self.norm1 = nn.GroupNorm(1, self.latent_size)

        self.weight2 = nn.Linear(self.latent_size, self.latent_size)
        self.norm2 = nn.GroupNorm(1, self.latent_size) # Batch normalization solves a major problem called internal covariate shift.

        self.output_fc = nn.Linear(self.latent_size, args.pred_len * 2)

    def forward(self, prednet_in):
        # Residual layer
        x = self.weight1(prednet_in)
        x = self.norm1(x)
        x = F.relu(x)
        x = self.weight2(x)
        x = self.norm2(x)

        x += prednet_in

        x = F.relu(x)

        # Last layer has no activation function
        prednet_out = self.output_fc(x)

        return prednet_out

class map_smooth_decoder(nn.Module):
    def __init__(self, args):
        super(map_smooth_decoder, self).__init__()

        self.args = args
        self.latent_size = self.args.map_latent_size

        self.norm0 = nn.BatchNorm1d(self.latent_size)

        self.conv1 = nn.Conv1d(self.latent_size, self.latent_size // 4, kernel_size=3, padding=1)
        self.norm1 = nn.BatchNorm1d(self.latent_size // 4)

        self.conv2 = nn.Conv1d(self.latent_size // 4, self.latent_size // 8, kernel_size=3, padding=1)
        self.norm2 = nn.BatchNorm1d(self.latent_size // 8)

        self.linear3 = nn.Linear(self.args.centerline_length * (self.latent_size // 8), self.latent_size // 8)
        self.norm3 = nn.BatchNorm1d(self.latent_size // 8)

        self.linear4 = nn.Linear(self.args.num_centerlines * (self.latent_size // 8), self.latent_size)

    def forward(self, x):
        total_centerlines = x.shape[0]
        batch_size = x.shape[0] // self.args.num_centerlines

        x = x.permute(0, 2, 1)
        x = self.norm0(x)

        x = self.norm1(F.relu(self.conv1(x)))
        x = self.norm2(F.relu(self.conv2(x)))

        x = self.norm3(F.relu(self.linear3(x.contiguous().view(total_centerlines,-1))))
        x = self.linear4(x.contiguous().view(batch_size,-1))

        return x

class MLP(nn.Module):
    def __init__(self, input_size, output_size) -> None:
        super(MLP, self).__init__()
        self.linear1 = nn.Linear(input_size, output_size // 2)
        self.norm = nn.LayerNorm(output_size // 2)
        self.GELU = nn.GELU()
        self.linear2 = nn.Linear(output_size // 2, output_size)

        # self.linear1 = nn.Linear(input_size, output_size)

    def forward(self, x):
        x = self.linear1(x)
        x = self.norm(x)
        x = self.GELU(x)
        x = self.linear2(x)
        return x

class MapSubNet(nn.Module):

    def __init__(self, args, depth=None):
        super(MapSubNet, self).__init__()

        self.args = args

        if depth is None:
            depth = 2

        self.hidden_size = self.args.map_latent_size
        self.input_dim = self.args.data_dim
        self.dropout = self.args.apply_dropout

        self.MLPs = nn.ModuleList([MLP(self.input_dim, self.hidden_size // 8), MLP(self.hidden_size // 4, self.hidden_size // 2)])
        self.Attn = nn.ModuleList([nn.MultiheadAttention(self.hidden_size // 8, self.args.num_attention_heads, dropout=self.dropout),
                                   nn.MultiheadAttention(self.hidden_size // 2, self.args.num_attention_heads, dropout=self.dropout)])
        self.Norms = nn.ModuleList([nn.LayerNorm(self.hidden_size // 4), nn.LayerNorm(self.hidden_size)])

        self.final_layer = map_smooth_decoder(self.args)

    def forward(self, inputs, inputs_mask):

        hidden_states_batch = inputs
        hidden_states_mask = inputs_mask

        for layer_index, layer in enumerate(self.Attn):
            hidden_states_batch = self.MLPs[layer_index](hidden_states_batch)
            if torch.any(torch.isnan(hidden_states_batch)):
                pdb.set_trace()
            temp = hidden_states_batch
            query = key = value = hidden_states_batch.permute(1,0,2)

            # hidden_states_batch = layer(query, key, value=value, attn_mask=None, key_padding_mask=hidden_states_mask)[0].permute(1,0,2)
            hidden_states_batch = layer(query, key, value=value)[0].permute(1,0,2)
            if torch.any(torch.isnan(hidden_states_batch)):
                pdb.set_trace()
            hidden_states_batch = torch.cat([hidden_states_batch, temp], dim=2)
            hidden_states_batch = self.Norms[layer_index](hidden_states_batch)
            hidden_states_batch = F.relu(hidden_states_batch)
            if torch.any(torch.isnan(hidden_states_batch)):
                pdb.set_trace()
        if torch.any(torch.isnan(hidden_states_batch)):
            pdb.set_trace()

        hidden_states_batch = self.final_layer(hidden_states_batch)
        return hidden_states_batch

class TransformerDecoder(nn.Module):
    def __init__(self, hidden_size, head_num=8, dropout=0.1) -> None:
        super(TransformerDecoder, self).__init__()
        self.self_attn = nn.MultiheadAttention(hidden_size, head_num, dropout)
        self.cross_attn = nn.MultiheadAttention(hidden_size, head_num, dropout)

        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        self.dropout4 = nn.Dropout(dropout)

        self.norm1 = nn.LayerNorm(hidden_size)
        self.norm2 = nn.LayerNorm(hidden_size)
        self.norm3 = nn.LayerNorm(hidden_size)

        self.linear1 = nn.Linear(hidden_size, 256)
        self.linear2 = nn.Linear(256, hidden_size)

    def forward(self, x_padding, y_padding):

        self_attn_output = self.self_attn(query=x_padding,
                                     key=x_padding,
                                     value=x_padding)[0]
        x_padding = x_padding + self.dropout1(self_attn_output)
        x_padding = self.norm1(x_padding)

        cross_attn_output = self.cross_attn(query=x_padding,
                                            key=y_padding,
                                            value=y_padding)[0]

        x_padding = x_padding + self.dropout2(cross_attn_output)
        x_padding = self.norm2(x_padding)

        output = self.linear1(x_padding)
        output = F.relu(output)
        output = self.dropout3(output)
        output = self.linear2(output)

        x_padding = x_padding + self.dropout4(output)
        x_padding = self.norm3(x_padding)

        return x_padding

class Temporal_Multimodal_Decoder(nn.Module):
    def __init__(self, args):
        super(Temporal_Multimodal_Decoder, self).__init__()

        self.args = args
        self.data_dim = self.args.data_dim
        self.obs_len = self.args.obs_len
        self.pred_len = self.args.pred_len
        self.window_size = self.args.decoder_temporal_window_size

        self.decoder_h_dim = self.args.decoder_latent_size
        self.num_modes = self.args.num_modes

        self.spatial_embedding = nn.Linear(self.window_size*2, self.window_size*4)

        self.decoder = nn.LSTM(self.window_size*4,
                               self.decoder_h_dim,
                               num_layers=1)

        pred = []
        for _ in range(self.num_modes):
            pred.append(nn.Linear(self.decoder_h_dim,self.data_dim))
        self.hidden2pos = nn.ModuleList(pred)

        norm = "BN"
        ng = 1

        # Confidences

        self.latent_predictions = nn.Linear(self.args.num_modes*self.args.pred_len*self.args.data_dim,
                                            self.decoder_h_dim)
        self.confidences = nn.Sequential(LinearRes(self.decoder_h_dim*2, self.decoder_h_dim*2, norm=norm, ng=ng),
                                         nn.Linear(self.decoder_h_dim*2, self.num_modes))

    def forward(self, traj_rel, state_tuple, num_mode=None, current_centerlines=None):
        """_summary_

        Args:
            traj_rel (_type_): _description_
            state_tuple (_type_): _description_
            num_mode (_type_, optional): _description_. Defaults to None.
            current_centerlines (_type_, optional): _description_. Defaults to None.

        Returns:
            _type_: _description_
        """

        traj_rel = traj_rel.permute(1,0,2)
        num_displacements, batch_size, data_dim = traj_rel.shape
        state_tuple_h, state_tuple_c = state_tuple

        pred_traj_fake_rel = []

        for num_mode in range(self.num_modes):
            traj_rel_ = torch.clone(traj_rel)
            decoder_input = F.leaky_relu(self.spatial_embedding(traj_rel_.permute(1,0,2).contiguous().view(batch_size,-1))) # bs x window_size·2

            decoder_input = decoder_input.unsqueeze(0)
            decoder_input = F.dropout(decoder_input, p=self.args.apply_dropout, training=self.training)

            state_tuple_h_ = torch.clone(state_tuple_h)
            state_tuple_c_ = torch.zeros(tuple(state_tuple_h_.shape)).to(state_tuple_h_)

            curr_pred_traj_fake_rel = []
            for _ in range(self.pred_len):
                output, (state_tuple_h_, state_tuple_c_) = self.decoder(decoder_input, (state_tuple_h_, state_tuple_c_))
                rel_pos = self.hidden2pos[num_mode](output.contiguous().view(-1, self.decoder_h_dim))

                traj_rel_ = torch.roll(traj_rel_, -1, dims=(0))
                traj_rel_[-1] = rel_pos

                curr_pred_traj_fake_rel.append(rel_pos)

                decoder_input = F.leaky_relu(self.spatial_embedding(traj_rel_.permute(1,0,2).contiguous().view(batch_size,-1))) # bs x window_size·2
                decoder_input = decoder_input.unsqueeze(0)
                decoder_input = F.dropout(decoder_input, p=self.args.apply_dropout, training=self.training)

            curr_pred_traj_fake_rel = torch.stack(curr_pred_traj_fake_rel,dim=0)
            curr_pred_traj_fake_rel = curr_pred_traj_fake_rel.permute(1,0,2)
            pred_traj_fake_rel.append(curr_pred_traj_fake_rel)

        pred_traj_fake_rel = torch.stack(pred_traj_fake_rel, dim=0) # num_modes, batch_size, pred_len, data_dim
        pred_traj_fake_rel = pred_traj_fake_rel.permute(1,0,2,3) # batch_size, num_modes, pred_len, data_dim

        # Obtain confidences based on the initial latent state and the predictions

        predictions_latent = self.latent_predictions(pred_traj_fake_rel.contiguous().view(batch_size, -1))
        state_tuple_h = state_tuple_h.squeeze(0)
        conf_latent = torch.cat([state_tuple_h,
                                 predictions_latent],
                                 dim=1)

        conf = self.confidences(conf_latent)
        conf = torch.softmax(conf.view(batch_size,-1), dim=1) # batch_size, num_modes
        if not torch.allclose(torch.sum(conf, dim=1), conf.new_ones((batch_size,))):
            pdb.set_trace()

        return pred_traj_fake_rel, conf

# Aux functions

def relative_to_abs_multimodal(rel_traj, start_pos):
    """
    Inputs:
    - rel_traj: pytorch tensor of shape (batch_size, num_modes, seq_len, 2)
    - start_pos: pytorch tensor of shape (batch_size, 2)
      N.B. If you only have the predictions, this must be the last observation.
           If you have the whole trajectory (obs+pred), this must be the first observation,
           since you must reconstruct the relative displacements from this position
    Outputs:
    - abs_traj: pytorch tensor of shape (seq_len, batch, 2) (around 0,0, not map coordinates)
    """

    displacement = torch.cumsum(rel_traj, dim=2) # Sum along the seq_len dimension!
    start_pos = torch.unsqueeze(torch.unsqueeze(start_pos, dim=1), dim=1) # batch, 1 (only one position) x 1 (same for all modes) x 2
    abs_traj = displacement + start_pos

    return abs_traj


import pandas as pd
import plotly.graph_objects as go
import prepare_data

population = {
    'NSW':8089526,
    'QLD':5095100,
    'VIC':6594804,
    'SA':1751693,
    'WA':2621680,
    'TAS':534281,
    'ACT':426709,
    'NT':245869,
    'Total':25359662,
    'DeathsNationally':25359662,
}

df_aus = prepare_data.australia()
df_aus_change = prepare_data.australia_change(df_aus)

# Let's plot this mofo
fig = go.Figure()

# Plot all the states!
for state in list(df_aus):
    fig.add_trace(go.Scatter(
        x=df_aus.index,
        y=pd.to_numeric(df_aus[state]).divide(population[state])*100000,
        name=state,
    ))

# Make the plot look fancy.
fig.update_layout(title='Per Capita COVID-19 Cases by State/Territory in Austalia',
                   xaxis_title='Date',
                   yaxis_title='Cases per 100,000 people')
fig.show()

# Let's plot this mofo
fig_change = go.Figure()

# Plot all the states!
for state in list(df_aus_change):
    fig_change.add_trace(go.Scatter(
        x=df_aus_change.index,
        y=pd.to_numeric(df_aus_change[state]).divide(population[state])*100000,
        name=state,
    ))

# Make the plot look fancy.
fig_change.update_layout(title='Per Capita Change in COVID-19 Cases by State/Territory in Austalia',
                         xaxis_title='Date',
                         yaxis_title='Change in cases per 100,000 people')
fig_change.show()

# Roll those numbers over a week
df_aus_change = df_aus_change.rolling(7).mean()

# Let's plot this mofo
fig_rolling_change = go.Figure()

# Plot all the states!
for state in list(df_aus):
    fig_rolling_change.add_trace(go.Scatter(
        x=df_aus_change.index,
        y=pd.to_numeric(df_aus_change[state]).divide(population[state])*100000,
        name=state,
    ))

# Make the plot look fancy.
fig_rolling_change.update_layout(
    title='7-day Rolling Per Capita Change in COVID-19 Cases by State/Territory in Austalia',
    xaxis_title='Date',
    yaxis_title='Change in cases per 100,000 people'
)

fig_rolling_change.show()


import collections
import difflib
import logging
import os
import re

from pytype.platform_utils import path_utils
from pytype.tools.merge_pyi import merge_pyi
import unittest

__all__ = ('TestBuilder', 'load_tests')

PY, PYI, EXPECTED = 'py', 'pyi', 'pep484.py'
OVERWRITE_EXPECTED = 0  # flip to regenerate expected files

def load_tests(unused_loader, standard_tests, unused_pattern):
  root = path_utils.join(path_utils.dirname(__file__), 'test_data')
  standard_tests.addTests(TestBuilder().build(root))
  return standard_tests

class TestBuilder:

  def build(self, data_dir):
    """Return a unittest.TestSuite with tests for the files in data_dir."""

    suite = unittest.TestSuite()

    files_by_base = self._get_files_by_base(data_dir)

    for base, files_by_ext in sorted(files_by_base.items()):
      if not (PY in files_by_ext and PYI in files_by_ext):
        continue

      if not OVERWRITE_EXPECTED and EXPECTED not in files_by_ext:
        continue

      py, pyi = (files_by_ext[x] for x in (PY, PYI))
      outfile = path_utils.join(data_dir, base + '.' + EXPECTED)

      test = build_regression_test(py, pyi, outfile)
      suite.addTest(test)

    return suite

  def _get_files_by_base(self, data_dir):
    files = os.listdir(data_dir)

    file_pat = re.compile(r'(?P<filename>(?P<base>.+?)\.(?P<ext>.*))$')
    matches = [m for m in map(file_pat.match, files) if m]
    ret = collections.defaultdict(dict)
    for m in matches:
      base, ext, filename = m.group('base'), m.group('ext'), m.group('filename')
      ret[base][ext] = path_utils.join(data_dir, filename)

    return ret

def build_regression_test(py, pyi, outfile):

  def regression_test(test_case):
    py_input, pyi_src = (_read_file(f) for f in (py, pyi))
    try:
      output = merge_pyi.merge_sources(py=py_input, pyi=pyi_src)
    except merge_pyi.MergeError:
      pass

    if OVERWRITE_EXPECTED:
      with open(outfile, 'w') as f:
        f.write(output)
    else:
      expected = _read_file(outfile)
      test_case.assertEqual(expected, output, _get_diff(expected, output))

  name = path_utils.splitext(path_utils.basename(outfile))[0].replace('.', '_')
  test = f'test_{name}'
  case = type('RegressionTest', (unittest.TestCase,), {test: regression_test})
  return case(test)

def _read_file(filename):
  with open(filename) as f:
    return f.read()

def _get_diff(a, b):
  a, b = a.split('\n'), b.split('\n')

  diff = difflib.Differ().compare(a, b)
  return '\n'.join(diff)

if __name__ == '__main__':
  logging.basicConfig(level=logging.CRITICAL)
  unittest.main()


import datetime
import time

from gpiozero import LED, Device
from gpiozero.pins.pigpio import PiGPIOFactory

Device.pin_factory = PiGPIOFactory()

# NOTE: Change this to match the GPIO pin you're connecting the LED to
led = LED(18)
# NOTE: Change these values to set the time you want the light to turn on and off at
weekday_on_time = datetime.time(hour=7, minute=0, second=0)
weekday_off_time = datetime.time(hour=17, minute=0, second=0)
weekend_on_time = datetime.time(hour=7, minute=30, second=0)
weekend_off_time = datetime.time(hour=17, minute=0, second=0)

while True:
    dayOfWeek = datetime.datetime.now().weekday()
    currentTime = datetime.datetime.now().time()
    on_time = weekday_on_time if dayOfWeek < 5 else weekend_on_time
    off_time = weekday_off_time if dayOfWeek < 5 else weekend_off_time
    if currentTime > on_time and currentTime < off_time:
        led.on()
    else:
        led.off()

    time.sleep(60)


"""Module to index columns of the paper-summarized CSV file."""
import pandas as pd
from loguru import logger
from omegaconf import OmegaConf

from utils import create_embeddings

# Load the configuration
cfg = OmegaConf.load("conf/config.yaml")

FILE_PATH = cfg.data.path
INDEXED_FILE_PATH = cfg.data.indexed_path

df = pd.read_csv(cfg.data.path, compression="gzip", header=0)
logger.info(f"Loaded {len(df)} rows from {cfg.data.path}")
logger.info(f"columns: {df.columns}")

logger.info("Creating embeddings for experiment time")
df = create_embeddings(df, ["experiment time"])

logger.info("Creating embeddings for device")
df = create_embeddings(df, ["device"])

print(df.head(5))

logger.info(f"Saving indexed file to {INDEXED_FILE_PATH}")
df.to_csv(INDEXED_FILE_PATH, compression="gzip", index=False)

logger.success("Done!")


import re
from io import StringIO

from flask import Flask, request, Response, redirect
import pandas as pd

app = Flask(__name__)

def is_valid_query(q):
    '''
    A query is valid if it is strictly consisted of the following three entries:
    [(A-Z, a-z, 0-9)+ or *] [==, !=, $=, &=] ["..."]
    Queries can be concat with 'and' & 'or'. The 'and' 'or' operators are executed in sequential order.
    Entries and operators must be separated by at least one single space. i.e. {C1=="a"} is not acceptable.
    Additional white spaces are allowed between entries.
    For the "" inside the query, the query term is the content wrapped by the first and last occurance of "".
    In the processing of query term, any sequence of consecutive spaces is reduced to a single space for clarity.
    (since consecutive spaces usually do not convey any semantic meanings in phrases)

    Output:
    a message indicating the validity of query ("valid" or error message),
    a list of valid queries (each query is represented by a 3-element list),
    a list of and/or operators
    '''
    entries = q.split()
    valid_q = [] # a 3-element list consist of 3 valid entries
    queries = [] # list of valid_q
    operators = [] # operators between queries
    operand = ['==', '!=', '$=', '&='] # valid operand

    # check the valid status of three entries defined above
    valid_first = False
    valid_second = False
    valid_third = False

    i = 0
    while(i < len(entries)):
        if not valid_first:
            column = re.findall('[A-Za-z0-9]+|\*', entries[i])
            # if valid, must be exactly one match
            # i.e. "abc*123" will give three matches and is invalid
            if len(column) != 1 or column[0] != entries[i]:
                return "Invalid column name", queries, operators
            else:
                valid_q.append(entries[i])
                valid_first = True

        elif not valid_second:
            if entries[i] not in operand:
                return "Invalid operator, must be ==, !=, $=, &=", queries, operators
            else:
                # store as int if only numbers
                valid_q.append(entries[i])
                valid_second = True

        elif not valid_third:
            if entries[i][0] != '\"':
                return "Invalid query term, must begin with \"", queries, operators
            else: # traverse the list to find the last " before the next query
                term = ""
                # find the string before next query
                if entries[i:].count('and') > 0:
                    end = entries[i:].index('and') + i
                    term = " ".join(entries[i:end])
                elif entries[i:].count('or') > 0:
                    end = entries[i:].index('or') + i
                    term = " ".join(entries[i:end])
                else:
                    end = len(entries)
                    term = " ".join(entries[i:])
                # test the validity of term
                if term[-1] != '\"':
                    return "Invalid query term, must end with \"", queries, operators
                else:
                    i = end
                    valid_q.append(term[1:-1]) # remove the front and end "" when storing
                    valid_third = True
                    continue

        else:
            if i == len(queries) - 1:
                return "Extra term after queries", queries, operators
            if entries[i] == 'and' or entries[i] == 'or':
                queries.append(valid_q)
                operators.append(entries[i])
                valid_q = []
                valid_first = valid_second = valid_third = False
            else:
                return "Invalid and/or operand between queries", queries, operators
        i += 1

    # append the last valid query and check incomplete query
    if valid_first and valid_second and valid_third:
        queries.append(valid_q)
    else:
        return "Missing entries in queries", queries, operators

    return "valid", queries, operators

def match_query(queries, operators, df):
    '''
    This function matches the queries associated with the operators to df.
    Output:
    a message indicating the validity of query matching ('valid' or error message)
    matched rows in df
    '''
    columns = df.columns.tolist()
    res_df = pd.DataFrame(columns = columns) # empty df to append matching rows

    for i,q in enumerate(queries):
        # if this is the first query or the operator connecting pervious query is 'or', check the entire df
        if i - 1 < 0 or operators[i - 1] == 'or':
            cur_df = df.astype(str) # convert the content of df to string for comparison
        elif operators[i - 1] == 'and':
            cur_df = res_df

            # select rows from df
        if q[0] == "*":
            select_df = pd.DataFrame(columns = columns) # empty df to append matching rows
            for (col, _) in cur_df.iteritems():
                if q[1] == "==":
                    select_df = select_df.append(cur_df[cur_df[col] == q[2]])
                elif q[1] == "!=":
                    drop_df = cur_df[cur_df[col] == q[2]]
                    select_df = select_df.append(cur_df.drop(index=drop_df.index.tolist()))
                elif q[1] == "$=":
                    select_df = select_df.append(cur_df[cur_df[col].str.lower().isin([q[2].lower()])])
                elif q[1] == "&=":
                    select_df = select_df.append(cur_df[cur_df[col].str.contains(q[2], case=True)])
            cur_df = select_df.drop_duplicates(keep='first')
        else:
            if q[0] not in columns:
                return 'No corresponding column name in data', res_df
            elif q[0] not in cur_df.columns:
                cur_df = pd.DataFrame(columns = columns) # no matching column, set the cur_df to empty
            else:
                if q[1] == "==":
                    cur_df = cur_df[cur_df[q[0]] == q[2]]
                elif q[1] == "!=":
                    drop_df = cur_df[cur_df[q[0]] == q[2]]
                    cur_df = cur_df.drop(index=drop_df.index.tolist())
                elif q[1] == "$=":
                    cur_df = cur_df[cur_df[q[0]].str.lower().isin([q[2].lower()])]
                elif q[1] == "&=":
                    cur_df = cur_df[cur_df[q[0]].str.contains(q[2], case=True)]

        # update res_df according to 'and' 'or' operators
        if i - 1 < 0 or operators[i - 1] == 'or':
            res_df = res_df.append(cur_df)
            res_df.drop_duplicates(keep='first',inplace=True)
        elif operators[i - 1] == 'and':
            res_df = cur_df

    if res_df.empty:
        return 'No corresponding items for the query', res_df
    return 'valid', res_df

@app.route('/')
def get_info():
    args = request.args
    query = args['query']

    # '&' will separate the query to two items, append it back
    for key, val in args.items():
        if key == 'query':
            continue
        if key == '':
            query += '&=' + val
        else:
            query += '&' + key
    print(query)

    # Query error checking and parsing
    mes, queries, operators = is_valid_query(query)
    if mes != "valid":
        return mes
    print(queries)
    print(operators)

    # Query match
    df = pd.read_csv('data.csv')
    mes, res_df = match_query(queries, operators, df)
    if mes != "valid":
        return mes

    res_df.to_csv('res.csv')

    return '''

        The query has been successfully processed.
        To download the extracted results in a csv file, click me.

        '''

@app.route("/getCSV")
def getCSV():
    output = StringIO()
    df = pd.read_csv('res.csv')
    df.to_csv(output)

    return Response(
        output.getvalue(),
        mimetype="text/csv",
        headers={"Content-disposition":"attachment; filename=res.csv"})

if __name__ == '__main__':
   app.run(host='127.0.0.1', port=9527)


#coding:utf8

#login
import logging
logging.basicConfig(level=logging.DEBUG)
_logger = logging.getLogger(__name__)

#flask frame
from flask_restplus import Resource

#wechat frame
import flask_wechat_utils
from flask_wechat_utils.user.utils import auth
from flask_wechat_utils.config import api

#application config
import config as config_application

#application model
from models import MessageTemplate

#application
from utils import get_formid_and_delete

#-------------------------------------------
# blueprint/api/ns
#-------------------------------------------
ns = api.namespace(
	config_application.APPLICATION_NAME,
	description=config_application.APPLICATION_DESCRIPTION
)

# api = flask_wechat_utils.create_api()

# ns = api.namespace(
# 	config_application.APPLICATION_NAME,
# 	description=config_application.APPLICATION_DESCRIPTION
# )

#-------------------------------------------
# /parser/marshal
#-------------------------------------------
parser_messageTemplate_create = api.parser()
parser_messageTemplate_create.add_argument('form_id',type=str,required=True)

#-------------------------------------------
# route
#-------------------------------------------
@ns.route('/')
class MessageTemplateRoute(Resource):

	@api.doc(parser=parser_messageTemplate_create)
	@auth
	def post(self):

		args = parser_messageTemplate_create.parse_args()

		message_template = MessageTemplate(
			openid=self.wechat_user.openid,
			form_id=args.get('form_id'),
		)

		message_template.save()

		return {
			'code':0,
		}

	@auth
	def get(self):

		form_id_result = get_formid_and_delete(self.wechat_user.openid)

		return {
			'code':0,
			'openid':form_id_result.openid,
			'created_ts':str(form_id_result.created_ts),
			'_id':str(form_id_result.id),
		}


import rclpy
import rclpy.node
from airobot_interfaces.srv import StringCommand

from gtts import gTTS
import speech_recognition as sr
import subprocess

class SpeechService(rclpy.node.Node):
    def __init__(self):
        super().__init__('speech_service')

        self.get_logger().info('音声サーバーを起動しました')

        self.init_rec = sr.Recognizer()

        self.service = self.create_service(
            StringCommand, '/speech_service/wake_up', self.command_callback)

    def command_callback(self, request, response):

        self.synthesis('I\'m ready.')

        text = None
        while text is None:
            text = self.recognition()

        self.synthesis(text)

        response.answer = text
        return response

    def recognition(self):
        text = ''

        with sr.Microphone() as source:
            while text == '':
                audio_data = self.init_rec.record(source, duration=5)
                self.get_logger().info(f'音声認識を行います')

                try:
                    text = self.init_rec.recognize_google(audio_data)
                except sr.UnknownValueError:
                    pass

        self.get_logger().info(f'認識したテキストは "{text}" です')

        return text

    def synthesis(self, text):

        self.get_logger().info(f'音声合成を実行します')
        self.get_logger().info(f'発話内容は "{text}"')

        gTTS(text, lang='en').save('voice.mp3')

        subprocess.run(['mpg123 voice.mp3'], shell=True)

def main():
    rclpy.init()

    speech_service = SpeechService()

    try:
        rclpy.spin(speech_service)
    except KeyboardInterrupt:
        pass

    rclpy.shutdown()


import itertools
import numpy as np
import networkx as nx
from sklearn.neighbors import kneighbors_graph
from sklearn.metrics.pairwise import euclidean_distances
from scipy.sparse.csgraph import minimum_spanning_tree

from ggc.utils import *

def knn_graph(X, k):
    """Returns k-Nearest Neighbor (MkNN) graph from the feature matrix.

    Parameters
    ----------
    X : ndarray, shape (N, F)
        N samples and F-dimensional features.
    k : int, k >= 1
        Parameter for knn: the k-th nearest neighbour.

    Returns
    -------
    adj : ndarray, shape (N, N)
        The adjacency matrix of the constructed knn graph.
    """
    assert k < X.shape[0]

    adj_directed = kneighbors_graph(X=X,
                            n_neighbors=k,
                            p=2,
                            include_self=False,
                            ).toarray()
    adj = adj_directed + adj_directed.T
    adj[adj > 0] = 1
    np.fill_diagonal(adj,0)

    return adj

def mknn_graph(X, k):
    """Returns Mutual k-Nearest Neighbor (MkNN) graph from the feature matrix.

    Parameters
    ----------
    X : ndarray, shape (N, F)
        N samples and F-dimensional features.
    k : int, k >= 1
        Parameter for mknn: the k-th nearest neighbour.

    Returns
    -------
    adj : ndarray, shape (N, N)
        The adjacency matrix of the constructed mknn graph.
    """
    assert k < X.shape[0]

    adj_directed = kneighbors_graph(X=X,
                            n_neighbors=k,
                            p=2,
                            include_self=False,
                            ).toarray()
    adj = adj_directed + adj_directed.T
    adj[adj < 2] = 0
    adj[adj >= 2] = 1
    np.fill_diagonal(adj,0)

    return adj

def cknn_graph(X, delta, k):
    """Returns Continuous k-Nearest Neighbor (CkNN) graph from the feature matrix.

    Parameters
    ----------
    X : ndarray, shape (N, F)
        N samples and F-dimensional features.
    delta : float, delta > 0
        Parameter for cknn.
    k : int, k >= 1
        Parameter for cknn: the k-th nearest neighbour.

    Returns
    -------
    adj : ndarray, shape (N, N)
        The adjacency matrix of the constructed cknn graph.

    References
    ----------
    .. [1] Tyrus Berry, Timothy Sauer. Consistent manifold representation for topological data analysis.
           Foundations of Data Science, 2019, 1 (1) : 1-38. doi: 10.3934/fods.2019001
    """
    assert k < X.shape[0]

    D = euclidean_distances(X, X)
    N = D.shape[0]
    np.fill_diagonal(D,0)
    D_k = np.sort(D)

    adj = np.zeros([N, N])
    adj[np.square(D) < delta * delta * np.dot(D_k[:,k].reshape(-1,1),D_k[:,k].reshape(1,-1))] = 1
    np.fill_diagonal(adj,0)

    return adj

def mst_graph(X):
    """Returns Minimum Spanning Tree (MST) graph from the feature matrix.

    Parameters
    ----------
    X : ndarray, shape (N, F)
        N samples and F-dimensional features.

    Returns
    -------
    adj : ndarray, shape (N, N)
        The adjacency matrix of the constructed mst graph.
    """
    D = euclidean_distances(X, X)
    adj_directed = minimum_spanning_tree(D).toarray()
    adj = adj_directed + adj_directed.T
    adj[adj > 0] = 1
    np.fill_diagonal(adj,0)

    return adj

def rmst_graph(X, gamma, k):
    """Returns Relaxed Minimum Spanning Tree (RMST) graph from the feature matrix.

    Parameters
    ----------
    X : ndarray, shape (N, F)
        N samples and F-dimensional features.
    gamma : float, gamma > 0
        Parameter for rmst.
    k : int, k >= 1
        Parameter for rmst: the k-th nearest neighbour.

    Returns
    -------
    adj : ndarray, shape (N, N)
        The adjacency matrix of the constructed rmst graph.

    References
    ----------
    .. [1] Beguerisse-Díaz, Mariano, Borislav Vangelov, and Mauricio Barahona.
           "Finding role communities in directed networks using role-based similarity,
           markov stability and the relaxed minimum spanning tree."
           2013 IEEE Global Conference on Signal and Information Processing. IEEE, 2013.
    """
    D = euclidean_distances(X, X)
    N = D.shape[0]
    assert k < N
    np.fill_diagonal(D,0)

    adj = np.zeros([N, N])

    D_k = np.sort(D)
    D_k = np.tile(D_k[:,k],(N,1))
    D_k = gamma * (D_k + D_k.T)
    np.fill_diagonal(D_k,0)

    max_weight = np.zeros((N,N))
    G = nx.Graph(D)
    T = nx.minimum_spanning_tree(G)
    path = dict(nx.all_pairs_dijkstra_path(T))
    for i,j in itertools.combinations(range(N),2):
        p = path[i][j]
        path_weight = np.zeros(len(p)-1)
        for k in range(len(p)-1):
            path_weight[k] = T.edges[p[k],p[k+1]]['weight']
        max_weight[i][j] = np.amax(path_weight)
    max_weight = max_weight + max_weight.T
    np.fill_diagonal(max_weight,0)

    adj[D < (max_weight + D_k)] = 1
    np.fill_diagonal(adj,0)

    return adj


import streamlit as st

from src.plotgraphs import make_radar_graph
from src.sentanalysis import hf_analysis
from src.sentanalysis import spacy_sentiment

if __name__ == "__main__":
    st.write("Welcome")
    user_input = st.text_input("Enter a sentence", key="name")
    result = st.button("Submit")
    if result:
        new_sentiment = hf_analysis(user_input)
        output = new_sentiment.pop()
        doc = spacy_sentiment(user_input)
        st.plotly_chart(make_radar_graph(doc))
        st.write(doc)
        st.write(
            "Prediction by using SiEBERT - English-Language Sentiment Classification"
        )
        st.write(output)


import logging

logging.basicConfig(filename='test_logs.log', encoding='utf-8', level=logging.INFO)
logger = logging.getLogger('selenium')
logger.setLevel(logging.INFO)

disable_loggers = ['urllib3.connectionpool','faker.factory']

def pytest_configure():
    for logger_name in disable_loggers:
        logger_not = logging.getLogger(logger_name)
        logger_not.disabled = True


import cv2

cap = cv2.VideoCapture('vtest.avi')

hog = cv2.HOGDescriptor()  # 클래스 호출을 통해 객체 생성

# SVM: 머신러닝 기술 이름
hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())

while True:
    ret, frame = cap.read()  # 프레임을 읽어서 반환
    # ret: true / false, true: 동영상 frame을 정상적으로 읽었을 때, false: 비정상적으로 읽었을 때
    if not ret:
        break

    detected, _ = hog.detectMultiScale(frame)
    for (x, y, w, h) in detected:
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 200), 3)

    cv2.imshow('CCTV', frame)
    if cv2.waitKey(10) == 27:  # 10: ascii 코드로 esc키를 의미, esc키를 누르면 waitKey 함수는 27 반환
        break

cv2.destroyAllWindows()


# from threading import Thread

import speech_recognition as sr
import keyboard as k
import spotipy
import os
import pyttsx3
import random
import credentials
from spotipy.oauth2 import SpotifyOAuth
from spotipy.oauth2 import SpotifyClientCredentials

# from refresh import Refresh

# from googleText2Speech import synthesize_text

os.environ["SPOTIPY_CLIENT_ID"] = credentials.SPOTIPY_CLIENT_ID
os.environ["SPOTIPY_CLIENT_SECRET"] = credentials.SPOTIPY_CLIENT_SECRET
os.environ["SPOTIPY_REDIRECT_URI"] = credentials.SPOTIPY_REDIRECT_URI
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credentials.GOOGLE_APPLICATION_CREDENTIALS

deviceId = credentials.DEVICE_ID
scope = "user-modify-playback-state"

auth_manager = SpotifyClientCredentials()
sp = spotipy.Spotify(auth_manager=SpotifyOAuth(scope=scope))

# TTS engine
engine = pyttsx3.init('sapi5')
voices = engine.getProperty('voices')
engine.setProperty('voice', voices[1].id)

# Mic init
r = sr.Recognizer()
mic = sr.Microphone(device_index=2)

jarvisResponses = ["I'm on it.", "Consider it done.", "Right away, Sir.", "Yes sir."]

def speak(text):
    engine.say(text)
    engine.runAndWait()

def main():
    while 1:
        try:
            with mic as source:
                r.adjust_for_ambient_noise(source)
                audio = r.listen(source)
                response = (r.recognize_google(audio))
                print(response)

                if any(x in response for x in ["Jarvis", "Yaris", "Garvais", "Taurus"]):
                    speak("Sir?")
                    audio = r.listen(source)
                    response = (r.recognize_google(audio))
                    # print(response)
                    # Discord Functionality
                    if any(x in response for x in ["mute", "unmute", "mutiny"]):
                        k.press_and_release('F8')
                        speak("It's done.")
                    elif any(x in response for x in ["deafen", "undeafen", "quiet"]):
                        k.press_and_release('F9')
                        speak("It's done.")

                    # Spotify Functionality
                    if any(x in response for x in ["next", "skip"]):
                        speak(jarvisResponses[random.randint(0, 3)])
                        sp.next_track(deviceId)
                    if any(x in response for x in ["previous", "last", "replay"]):
                        speak(jarvisResponses[random.randint(0, 3)])
                        sp.previous_track(deviceId)
                    if any(x in response for x in ["pause", "stop"]):
                        try:
                            speak(jarvisResponses[random.randint(0, 3)])
                            sp.pause_playback(deviceId)
                        except spotipy.exceptions.SpotifyException:
                            pass
                    elif any(x in response for x in ["resume", "continue", "play"]):
                        try:
                            speak(jarvisResponses[random.randint(0, 3)])
                            sp.start_playback(deviceId)
                        except spotipy.exceptions.SpotifyException:
                            pass
                    if any(x in response for x in ["increase", "lower", "raise", "set", "volume"]) and any(
                            char.isdigit() for char in response):
                        speak(jarvisResponses[random.randint(0, 3)])
                        volume = [int(s) for s in response.split() if s.isdigit()]
                        sp.volume(volume[0], deviceId)

                    if any(x in response for x in ["fast-forward", "fast", "forward"]) and any(
                            char.isdigit() for char in response):
                        speak(jarvisResponses[random.randint(0, 3)])
                        time = [int(s) for s in response.split() if s.isdigit()]
                        sp.seek_track(time[0] * 1000, deviceId)

                    # Application Functionality
                    if "open" in response:
                        if "valorant" in response:
                            speak(jarvisResponses[random.randint(0, 3)])
                            os.startfile(r"C:\Users\Public\Desktop\VALORANT.lnk")
                        if any(x in response for x in ["Apex", "Legends", "legend"]):
                            speak(jarvisResponses[random.randint(0, 3)])
                            os.startfile(r"C:\Users\Nasir\Desktop\Apex Legends.url")
                        if any(x in response for x in ["aim", "labs", "lab"]):
                            speak(jarvisResponses[random.randint(0, 3)])
                            os.startfile(r"C:\Users\Nasir\Desktop\Aim Lab.url")
                        if "Spotify" in response:
                            speak(jarvisResponses[random.randint(0, 3)])
                            os.startfile(r"C:\Users\Nasir\AppData\Roaming\Spotify\Spotify.exe")

                    # PC Functionality
                    if "sleep" in response:
                        speak("Goodbye for now, sir.")
                        os.system("rundll32.exe powrprof.dll,SetSuspendState 0,1,0")
                    if "quit" in response:
                        speak("Goodbye for now, sir.")
                        break

        except sr.RequestError:
            # print("API unavailable")
            pass
        except sr.UnknownValueError:
            # print("Unable to recognize speech or nothing said")
            pass

if __name__ == '__main__':
    main()