from google.colab import drive

drive.mount("/content/drive")

import h5py
with h5py.File("gafs_dataset_very_small/0.hdf5", "r") as fread:
    keys = list(fread.keys())
    print(keys[:10])
    group = fread[keys[0]]
    print(group.keys())
    feature = group["feature"][:]
    token = group["token"][:]
    print(feature.shape)
    print(token)

import copy
import json
import os
import sys
import time
from functools import partial
from pathlib import Path

# Third party's modules
import numpy as np

import torch

from nltk.metrics.distance import edit_distance

from torch import nn
from torch.nn import functional as F
from torch.utils.data import (
    DataLoader)

from torchvision.transforms import Compose

# Local modules
sys.path.append("modules_gislr")
from modules_gislr.dataset import (
    HDF5Dataset,
    merge)
from modules_gislr.defines import (
    get_fullbody_landmarks
)
from modules_gislr.layers import (
    RNNEncoder,
    apply_norm,
    create_norm
)
from modules_gislr.train_functions import (
    LabelSmoothingCrossEntropyLoss
)
from modules_gislr.transforms import (
    PartsBasedNormalization,
    ReplaceNan,
    SelectLandmarksAndFeature,
    ToTensor
)
from modules_gislr.utils import (
    select_reluwise_activation
)

class BahdanauAttentionEnergy(nn.Module):
    def __init__(self,
                 key_dim,
                 query_dim,
                 att_dim,
                 add_bias=False):
        super().__init__()

        self.w_key = nn.Linear(key_dim, att_dim, bias=add_bias)
        self.w_query = nn.Linear(query_dim, att_dim, bias=add_bias)
        self.w_out = nn.Linear(att_dim, 1, bias=add_bias)

    def forward(self, key, query):
        # key: `[N, key_len, key_dim]`
        # query: `[N, 1, query_dim]`
        key = self.w_key(key)
        query = self.w_query(query)
        # Adding with broadcasting.
        # key: `[N, key_len, key_dim]`
        # query: `[N, 1, query_dim]`
        # query should be broadcasted to `[N, key_len, query_dim]`
        temp = key + query
        # `[N, key_len, att_dim] -> [N, key_len, 1] -> [N, 1, key_len]`
        energy = self.w_out(torch.tanh(temp))
        energy = torch.permute(energy, [0, 2, 1])
        return energy

class SingleHeadAttention(nn.Module):
    def __init__(self,
                 key_dim,
                 query_dim,
                 att_dim,
                 add_bias):
        super().__init__()

        self.att_energy = BahdanauAttentionEnergy(
            key_dim=key_dim,
            query_dim=query_dim,
            att_dim=att_dim,
            add_bias=add_bias)

        self.neg_inf = None

    def forward(self,
                key,
                value,
                query,
                mask=None):
        if self.neg_inf is None:
            self.neg_inf = float(np.finfo(
                torch.tensor(0, dtype=key.dtype).numpy().dtype).min)

        batch, klen, kdim = key.shape
        _, qlen, qdim = query.shape
        energy = self.att_energy(key=key, query=query)
        assert energy.shape == (batch, qlen, klen)

        # Apply mask.
        if mask is not None:
            if len(mask.shape) == 2:
                # `[N, klen] -> [N, qlen(=1), klen]`
                mask = mask.unsqueeze(1)
            # Negative infinity should be 0 in softmax.
            energy = energy.masked_fill_(mask==0, self.neg_inf)

        # Compute attention mask.
        attw = torch.softmax(energy, dim=-1)
        # attw: `[N, qlen, klen]`
        # value: `[N, klen, kdim]`
        # bmm: `[N, qlen, klen] x [N, klen, kdim] -> [N, qlen, kdim]`
        cvec = torch.bmm(attw, value)
        return cvec, attw

class BahdanauRNNDecoder(nn.Module):
    def __init__(self,
                 in_channels,
                 hidden_channels,
                 out_channels,
                 emb_channels,
                 att_dim,
                 att_add_bias,
                 rnn_type,
                 num_layers,
                 activation,
                 dropout,
                 padding_val,
                 proj_size=0):
        super().__init__()
        assert rnn_type in ["srnn", "lstm", "gru"]

        self.emb_layer = nn.Embedding(
            num_embeddings=out_channels,
            embedding_dim=emb_channels,
            padding_idx=padding_val)

        self.att_layer = SingleHeadAttention(
            key_dim=in_channels,
            query_dim=hidden_channels,
            att_dim=att_dim,
            add_bias=att_add_bias)

        if rnn_type == "srnn":
            self.rnn = nn.RNN(input_size=in_channels + emb_channels,
                              hidden_size=hidden_channels,
                              num_layers=num_layers,
                              nonlinearity=activation,
                              batch_first=True,
                              dropout=dropout,
                              bidirectional=False)
        elif rnn_type == "lstm":
            self.rnn = nn.LSTM(input_size=in_channels + emb_channels,
                               hidden_size=hidden_channels,
                               num_layers=num_layers,
                               batch_first=True,
                               dropout=dropout,
                               bidirectional=False,
                               proj_size=proj_size)
        elif rnn_type == "gru":
            self.rnn = nn.GRU(input_size=in_channels + emb_channels,
                              hidden_size=hidden_channels,
                              num_layers=num_layers,
                              batch_first=True,
                              dropout=dropout,
                              bidirectional=False)
        self.head = nn.Linear(hidden_channels, out_channels)

        self.num_layers = num_layers
        self.dec_hstate = None
        self.attw = None

    def init_dec_hstate(self, enc_hstate, init_as_zero=False):
        if init_as_zero:
            dec_hstate = torch.zeros_like(enc_hstate)
        else:
            dec_hstate = enc_hstate
        # To avoid error at RNN layer.
        self.dec_hstate = dec_hstate.contiguous()

    def forward(self,
                dec_inputs,
                enc_seqs,
                enc_mask):
        assert self.dec_hstate is not None, f"dec_hstate has not been initialized."
        dec_hstate = self.dec_hstate

        # Attention layer requires hidden state of 2nd rnn layer.
        # as `[N, 1, C]`
        query = dec_hstate[-1].unsqueeze(1)
        cvec, self.attw = self.att_layer(
            key=enc_seqs,
            value=enc_seqs,
            query=query,
            mask=enc_mask)

        emb_out = self.emb_layer(dec_inputs)
        # `[N, C] -> [N, 1, C]`
        emb_out = emb_out.reshape([-1, 1, emb_out.shape[-1]])
        feature = torch.cat([cvec, emb_out], dim=-1)
        if isinstance(self.rnn, nn.LSTM):
            hidden_seqs, (last_hstate, last_cstate) = self.rnn(feature,
                                                               dec_hstate)
        else:
            hidden_seqs, last_hstate = self.rnn(feature,
                                                dec_hstate)
            last_cstate = None

        output_dec = self.head(hidden_seqs)
        self.dec_hstate = last_hstate
        return output_dec

class RNNCSLR(nn.Module):
    def __init__(self,
                 enc_in_channels,
                 enc_hidden_channels,
                 enc_rnn_type,
                 enc_num_layers,
                 enc_activation,
                 enc_bidir,
                 enc_dropout,
                 enc_apply_mask,
                 enc_proj_size,
                 dec_in_channels,
                 dec_hidden_channels,
                 dec_out_channels,
                 dec_emb_channels,
                 dec_att_dim,
                 dec_att_add_bias,
                 dec_rnn_type,
                 dec_num_layers,
                 dec_activation,
                 dec_dropout,
                 dec_padding_val,
                 dec_proj_size):
        super().__init__()
        self.enc_bidir = enc_bidir

        self.linear = nn.Linear(enc_in_channels, enc_hidden_channels)
        self.enc_activation = nn.ReLU()

        self.encoder = RNNEncoder(
            in_channels=enc_hidden_channels,
            out_channels=enc_hidden_channels,
            rnn_type=enc_rnn_type,
            num_layers=enc_num_layers,
            activation=enc_activation,
            bidir=enc_bidir,
            dropout=enc_dropout,
            apply_mask=enc_apply_mask,
            proj_size=enc_proj_size)

        if enc_bidir:
            dec_in_channels *= 2
            dec_hidden_channels *= 2
            dec_att_dim *= 2

        self.decoder = BahdanauRNNDecoder(
            in_channels=dec_in_channels,
            hidden_channels=dec_hidden_channels,
            out_channels=dec_out_channels,
            emb_channels=dec_emb_channels,
            att_dim=dec_att_dim,
            att_add_bias=dec_att_add_bias,
            rnn_type=dec_rnn_type,
            num_layers=dec_num_layers,
            activation=dec_activation,
            dropout=dec_dropout,
            padding_val=dec_padding_val,
            proj_size=dec_proj_size)

    def _apply_encoder(self, feature, feature_pad_mask=None):
        # Feature extraction.
        # `[N, C, T, J] -> [N, T, C, J] -> [N, T, C*J] -> [N, T, C']`
        N, C, T, J = feature.shape
        feature = feature.permute([0, 2, 1, 3])
        feature = feature.reshape(N, T, -1)

        feature = self.linear(feature)
        feature = self.enc_activation(feature)

        # Apply encoder.
        enc_seqs, enc_hstate = self.encoder(feature, feature_pad_mask)[:2]

        # Basically, decoder should not be bidirectional.
        # So, we should concatenate backwarded feature.
        if self.enc_bidir:
            # `[2*layers, N, C] -> [layers, N, 2*C]`
            enc_hstate = torch.permute(enc_hstate, [1, 0, 2])
            enc_hstate = enc_hstate.reshape([enc_hstate.shape[0],
                                             enc_hstate.shape[1] // 2,
                                             -1])
            enc_hstate = torch.permute(enc_hstate, [1, 0, 2])
        return enc_seqs, enc_hstate

    def forward(self,
                feature, tokens,
                feature_pad_mask=None, tokens_pad_mask=None):
        """Forward computation for train.
        """
        enc_seqs, enc_hstate = self._apply_encoder(feature, feature_pad_mask)

        # Apply decoder.
        self.decoder.init_dec_hstate(enc_hstate)
        preds = None
        for t_index in range(0, tokens.shape[-1]):
            # Teacher forcing.
            dec_inputs = tokens[:, t_index].reshape([-1, 1])
            pred = self.decoder(
                dec_inputs=dec_inputs,
                enc_seqs=enc_seqs,
                enc_mask=feature_pad_mask)
            if preds is None:
                preds = pred
            else:
                # `[N, T, C]`
                preds = torch.cat([preds, pred], dim=1)
        return preds

    def inference(self,
                  feature,
                  start_id,
                  end_id,
                  feature_pad_mask=None,
                  max_seqlen=62):
        """Forward computation for test.
        """
        enc_seqs, enc_hstate = self._apply_encoder(feature, feature_pad_mask)

        # Apply decoder.
        self.decoder.init_dec_hstate(enc_hstate)
        dec_inputs = torch.tensor([start_id]).to(feature.device)
        # `[N, T]`
        dec_inputs = dec_inputs.reshape([1, 1])
        preds = None
        pred_ids = [start_id]
        for _ in range(max_seqlen):
            pred = self.decoder(
                dec_inputs=dec_inputs,
                enc_seqs=enc_seqs,
                enc_mask=feature_pad_mask)
            if preds is None:
                preds = pred
            else:
                # `[N, T, C]`
                preds = torch.cat([preds, pred], dim=1)

            pid = torch.argmax(pred, dim=-1)
            dec_inputs = pid

            pid = pid.reshape([1]).detach().cpu().numpy()[0]
            pred_ids.append(int(pid))
            if int(pid) == end_id:
                break

        # `[N, T]`
        pred_ids = np.array([pred_ids])
        return pred_ids, preds

def forward(model, feature, tokens, feature_pad_mask, tokens_pad_mask):
    if isinstance(model, RNNCSLR):
        preds = model(feature,
                      tokens,
                      feature_pad_mask=feature_pad_mask,
                      tokens_pad_mask=tokens_pad_mask)
    else:
        raise NotImplementedError(f"Unknown model type:{type(model)}.")
    return preds

def inference(model, feature, start_id, end_id, max_seqlen=62):
    if isinstance(model, RNNCSLR):
        pred_ids, _ = model.inference(feature,
                                      start_id,
                                      end_id,
                                      max_seqlen=max_seqlen)
    else:
        raise NotImplementedError(f"Unknown model type:{type(model)}.")
    return pred_ids

# Access check.
dataset_dir = Path("gafs_dataset_very_small")
files = list(dataset_dir.iterdir())
dictionary = [fin for fin in files if ".json" in fin.name][0]
hdf5_files = [fin for fin in files if ".hdf5" in fin.name]

print(dictionary)
print(hdf5_files)

# Load dictionary.
with open(dictionary, "r") as fread:
    key2token = json.load(fread)

VOCAB = len(key2token)
# Add keywords.
key2token["<sos>"] = VOCAB
key2token["<eos>"] = VOCAB + 1
key2token["<pad>"] = VOCAB + 2
# Reset.
VOCAB = len(key2token)

_, use_landmarks = get_fullbody_landmarks()
use_features = ["x", "y"]

trans_select_feature = SelectLandmarksAndFeature(landmarks=use_landmarks, features=use_features)
trans_repnan = ReplaceNan()
trans_norm = PartsBasedNormalization(align_mode="framewise", scale_mode="unique")
trans_insert_token = InsertTokensForS2S(sos_token=key2token["<sos>"], eos_token=key2token["<eos>"])

pre_transforms = Compose([
    trans_select_feature,
    trans_repnan,
    trans_insert_token,
    trans_norm
])

train_transforms = Compose([
    ToTensor()
])

val_transforms = Compose([
    ToTensor()
])

test_transforms = Compose([
    ToTensor()
])

batch_size = 2
feature_shape = (len(use_features), -1, len(use_landmarks))
token_shape = (-1,)
merge_fn = partial(merge_padded_batch,
                   feature_shape=feature_shape,
                   token_shape=token_shape,
                   feature_padding_val=0.0,
                   token_padding_val=key2token["<pad>"])

dataset = HDF5Dataset(hdf5_files, pre_transforms=pre_transforms, transforms=train_transforms)

dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=merge_fn)
try:
    data = next(iter(dataloader))
    feature_origin = data["feature"]
    tokens_origin = data["token"]

    print(feature_origin.shape)
    print(tokens_origin)
except Exception as inst:
    print(inst)

# Define model.
# in_channels: J * C (130*2=260)
#   J: use_landmarks (130)
#   C: use_channels (2)
# out_channels: 10
in_channels = len(use_landmarks) * len(use_features)
out_channels = VOCAB
enc_hidden_channels = 64
dec_hidden_channels = 64
dec_emb_channels = 4
dec_att_dim = 64
model = RNNCSLR(
    enc_in_channels=in_channels,
    enc_hidden_channels=enc_hidden_channels,
    enc_rnn_type="gru",
    enc_num_layers=2,
    enc_activation="relu",
    enc_bidir=True,
    enc_dropout=0.1,
    enc_apply_mask=True,
    enc_proj_size=0,
    dec_in_channels=enc_hidden_channels,
    dec_hidden_channels=dec_hidden_channels,
    dec_out_channels=out_channels,
    dec_emb_channels=dec_emb_channels,
    dec_att_dim=dec_att_dim,
    dec_att_add_bias=True,
    dec_rnn_type="gru",
    dec_num_layers=2,
    dec_activation="relu",
    dec_dropout=0.1,
    dec_padding_val=key2token["<pad>"],
    dec_proj_size=0)

print(model)

# Sanity check.
sample = next(iter(dataloader))
logit = model(sample["feature"],
              tokens=sample["token"],
              feature_pad_mask=sample["feature_pad_mask"])
print(logit.shape)

# Set common parameters.
batch_size = 32
load_into_ram = True
test_pid = 0
num_workers = os.cpu_count()
print(f"Using {num_workers} cores for data loading.")
lr = 3e-4
label_smoothing = 0.1
sos_token = key2token["<sos>"]
eos_token = key2token["<eos>"]
pad_token = key2token["<pad>"]
max_seqlen = 60

epochs = 50
eval_every_n_epochs = 1
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} for computation.")

train_hdf5files = [fin for fin in hdf5_files if str(test_pid) not in fin.name]
val_hdf5files = [fin for fin in hdf5_files if str(test_pid) in fin.name]
test_hdf5files = [fin for fin in hdf5_files if str(test_pid) in fin.name]

_, use_landmarks = get_fullbody_landmarks()
use_features = ["x", "y"]

# Build dataloaders.
train_dataset = HDF5Dataset(train_hdf5files,
    pre_transforms=pre_transforms, transforms=train_transforms, load_into_ram=load_into_ram)
val_dataset = HDF5Dataset(val_hdf5files,
    pre_transforms=pre_transforms, transforms=val_transforms, load_into_ram=load_into_ram)
test_dataset = HDF5Dataset(test_hdf5files,
    pre_transforms=pre_transforms, transforms=test_transforms, load_into_ram=load_into_ram)

train_dataloader = DataLoader(train_dataset, batch_size=batch_size, collate_fn=merge_fn, num_workers=num_workers, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=batch_size, collate_fn=merge_fn, num_workers=num_workers, shuffle=False)
test_dataloader = DataLoader(test_dataset, batch_size=1, collate_fn=merge_fn, num_workers=num_workers, shuffle=False)

model_rnn = RNNCSLR(
    enc_in_channels=in_channels,
    enc_hidden_channels=enc_hidden_channels,
    enc_rnn_type="gru",
    enc_num_layers=2,
    enc_activation="relu",
    enc_bidir=True,
    enc_dropout=0.1,
    enc_apply_mask=True,
    enc_proj_size=0,
    dec_in_channels=enc_hidden_channels,
    dec_hidden_channels=dec_hidden_channels,
    dec_out_channels=out_channels,
    dec_emb_channels=dec_emb_channels,
    dec_att_dim=dec_att_dim,
    dec_att_add_bias=True,
    dec_rnn_type="gru",
    dec_num_layers=2,
    dec_activation="relu",
    dec_dropout=0.1,
    dec_padding_val=pad_token,
    dec_proj_size=0)

print(model_rnn)

loss_fn = LabelSmoothingCrossEntropyLoss(
    ignore_indices=pad_token, reduction="mean_temporal_prior",
    label_smoothing=label_smoothing)
optimizer = torch.optim.Adam(model_rnn.parameters(), lr=lr)

# Train, validation, and evaluation.
model_rnn.to(device)

train_losses = []
val_losses = []
test_wers = []
print("Start training.")
for epoch in range(epochs):
    print("-" * 80)
    print(f"Epoch {epoch+1}")

    train_loss, train_times = train_loop_csir_s2s(
        train_dataloader, model_rnn, loss_fn, optimizer, device,
        sos_token, eos_token,
        return_pred_times=True)
    val_loss, val_times = val_loop_csir_s2s(
        val_dataloader, model_rnn, loss_fn, device,
        sos_token, eos_token,
        return_pred_times=True)
    val_losses.append(val_loss)

    if (epoch+1) % eval_every_n_epochs == 0:
        wer, test_times = test_loop_csir_s2s(
            test_dataloader, model_rnn, device,
            sos_token, eos_token,
            return_pred_times=True,
            max_seqlen=max_seqlen)
        test_wers.append(wer)
train_losses_rnn = np.array(train_losses)
val_losses_rnn = np.array(val_losses)
test_wers_rnn = np.array(test_wers)

val_losses_rnn = np.array(val_losses_rnn)
test_wers_rnn = np.array(test_wers_rnn)
print(f"Minimum validation loss:{val_losses_rnn.min()} at {np.argmin(val_losses_rnn)+1} epoch.")
print(f"Minimum WER:{test_wers_rnn.min()} at {np.argmin(test_wers_rnn)*eval_every_n_epochs+1} epoch.")