from google.colab import drive

drive.mount("/content/drive")

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

import json
import os
import random
import sys
from functools import partial
from pathlib import Path
from typing import (
    Dict
)

# Third party's modules
import numpy as np

import torch
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_padded_batch)
from modules_gislr.defines import (
    get_fullbody_landmarks
)
from modules_gislr.layers import (
    TransformerEnISLR
)
from modules_gislr.train_functions import (
    test_loop,
    val_loop,
    train_loop
)
from modules_gislr.transforms import (
    PartsBasedNormalization,
    ReplaceNan,
    SelectLandmarksAndFeature,
    ToTensor,
    matrix_interp
)

class PartsBasedTemporalInterpolation():
    """Apply parts-based temporal interpolation.
    """

    def __init__(self,
                 apply_ratio=1.0,
                 face_head=0,
                 face_num=76,
                 lhand_head=76,
                 lhand_num=21,
                 pose_head=76+21,
                 pose_num=12,
                 rhand_head=76+21+12,
                 rhand_num=21,
                 accept_correct_ratio=0.1) -> None:
        self.apply_ratio = apply_ratio
        self.face_head = face_head
        self.face_num = face_num
        self.lhand_head = lhand_head
        self.lhand_num = lhand_num
        self.pose_head = pose_head
        self.pose_num = pose_num
        self.rhand_head = rhand_head
        self.rhand_num = rhand_num
        self.accept_correct_ratio = accept_correct_ratio

    def _gen_tmask(self, feature):
        tmask = feature == 0.0
        tmask = np.all(tmask, axis=(0, 2))
        tmask = np.logical_not(tmask)
        return tmask

    def _interp(self, feature):
        tmask = self._gen_tmask(feature)
        # `[C, T, J]`
        orig_shape = feature.shape
        tlength = orig_shape[1]
        # No failed tracking.
        if tmask.sum() == tlength:
            return feature
        # Too many failed.
        if tmask.sum() < self.accept_correct_ratio:
            return feature
        # All failed.
        if tmask.sum() == 0.0:
            return feature

        x = np.arange(tlength)
        xs = np.where(tmask != 0)[0]
        # `[C, T, J] -> [T, C*J]`
        feature = feature.transpose([1, 0, 2]).reshape([tlength, -1])
        ys = feature[xs, :]

        newfeature = matrix_interp(x, xs, ys)
        # `[T, C*J] -> [T, C, J] -> [C, T, J]`
        newfeature = newfeature.reshape([-1, orig_shape[0], orig_shape[2]])
        newfeature = newfeature.transpose([1, 0, 2])
        assert np.isnan(newfeature).any() == np.False_
        assert np.isinf(newfeature).any() == np.False_
        return newfeature

    def __call__(self,
                 data: Dict[str, Any]) -> Dict[str, Any]:
        """Execute interpolation.
        """
        if random.random() > self.apply_ratio:
            return data

        feature = data["feature"]
        if self.face_num > 0:
            face = feature[:, :, self.face_head: self.face_head+self.face_num]
            face = self._interp(face)
            feature[:, :, self.face_head: self.face_head+self.face_num] = face
        if self.lhand_num > 0:
            lhand = feature[:, :, self.lhand_head: self.lhand_head+self.lhand_num]
            lhand = self._interp(lhand)
            feature[:, :, self.lhand_head: self.lhand_head+self.lhand_num] = lhand
        if self.pose_num > 0:
            pose = feature[:, :, self.pose_head: self.pose_head+self.pose_num]
            pose = self._interp(pose)
            feature[:, :, self.pose_head: self.pose_head+self.pose_num] = pose
        if self.rhand_num > 0:
            rhand = feature[:, :, self.rhand_head: self.rhand_head+self.rhand_num]
            rhand = self._interp(rhand)
            feature[:, :, self.rhand_head: self.rhand_head+self.rhand_num] = rhand
        data["feature"] = feature
        return data

    def __str__(self):
        return f"{self.__class__.__name__}:{self.__dict__}"

# Access check.
dataset_dir = Path("dataset_top10")
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)

_, 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")

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

transforms_tinterp_00 = Compose([
    PartsBasedTemporalInterpolation(apply_ratio=0.0),
    ToTensor()])

transforms_tinterp_05 = Compose([
    PartsBasedTemporalInterpolation(apply_ratio=0.5),
    ToTensor()])

transforms_tinterp_10 = Compose([
    PartsBasedTemporalInterpolation(apply_ratio=1.0),
    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=0)

for trans in [transforms_tinterp_00, transforms_tinterp_05, transforms_tinterp_10]:
    dataset = HDF5Dataset(hdf5_files, pre_transforms=pre_transforms, transforms=trans)
    dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=merge_fn)
    try:
        data = next(iter(dataloader))
        feature_origin = data["feature"]

        print(feature_origin.shape)
    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)
inter_channels = 64
out_channels = VOCAB
activation = "relu"
tren_num_layers = 2
tren_num_heads = 2
tren_dim_ffw = 256
tren_dropout_pe = 0.1
tren_dropout = 0.1
tren_layer_norm_eps = 1e-5
tren_norm_first = True
tren_add_bias = True
tren_add_tailnorm = True

pooling_type = "none"

model = TransformerEnISLR(in_channels=in_channels,
                          inter_channels=inter_channels,
                          out_channels=out_channels,
                          activation=activation,
                          pooling_type=pooling_type,
                          tren_num_layers=tren_num_layers,
                          tren_num_heads=tren_num_heads,
                          tren_dim_ffw=tren_dim_ffw,
                          tren_dropout_pe=tren_dropout_pe,
                          tren_dropout=tren_dropout,
                          tren_layer_norm_eps=tren_layer_norm_eps,
                          tren_norm_first=tren_norm_first,
                          tren_add_bias=tren_add_bias,
                          tren_add_tailnorm=tren_add_tailnorm)
print(model)

# Sanity check.
logit = model(feature_origin)
print(logit.shape)
attw0 = model.tr_encoder.layers[0].attw.detach().cpu().numpy()
attw1 = model.tr_encoder.layers[0].attw.detach().cpu().numpy()
print(attw0.shape, attw1.shape)

# Set common parameters.
batch_size = 32
load_into_ram = True
test_pid = 16069
num_workers = os.cpu_count()
print(f"Using {num_workers} cores for data loading.")
lr = 3e-4

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_tinterp_00 = HDF5Dataset(
    train_hdf5files,
    pre_transforms=pre_transforms,
    transforms=transforms_tinterp_00,
    load_into_ram=load_into_ram)

train_dataset_tinterp_05 = HDF5Dataset(
    train_hdf5files,
    pre_transforms=pre_transforms,
    transforms=transforms_tinterp_05,
    load_into_ram=load_into_ram)

train_dataset_tinterp_10 = HDF5Dataset(
    train_hdf5files,
    pre_transforms=pre_transforms,
    transforms=transforms_tinterp_10,
    load_into_ram=load_into_ram)

val_dataset_tinterp_00 = HDF5Dataset(
    val_hdf5files,
    pre_transforms=pre_transforms,
    transforms=transforms_tinterp_00,
    load_into_ram=load_into_ram)
test_dataset_tinterp_00 = HDF5Dataset(
    test_hdf5files,
    pre_transforms=pre_transforms,
    transforms=transforms_tinterp_00,
    load_into_ram=load_into_ram)

val_dataset_tinterp_10 = HDF5Dataset(
    val_hdf5files,
    pre_transforms=pre_transforms,
    transforms=transforms_tinterp_10,
    load_into_ram=load_into_ram)
test_dataset_tinterp_10 = HDF5Dataset(
    test_hdf5files,
    pre_transforms=pre_transforms,
    transforms=transforms_tinterp_10,
    load_into_ram=load_into_ram)

train_dataloader_tinterp_00 = DataLoader(train_dataset_tinterp_00, batch_size=batch_size, collate_fn=merge_fn, num_workers=num_workers, shuffle=True)
train_dataloader_tinterp_05 = DataLoader(train_dataset_tinterp_05, batch_size=batch_size, collate_fn=merge_fn, num_workers=num_workers, shuffle=True)
train_dataloader_tinterp_10 = DataLoader(train_dataset_tinterp_10, batch_size=batch_size, collate_fn=merge_fn, num_workers=num_workers, shuffle=True)

val_dataloader_tinterp_00 = DataLoader(val_dataset_tinterp_00, batch_size=batch_size, collate_fn=merge_fn, num_workers=num_workers, shuffle=False)
test_dataloader_tinterp_00 = DataLoader(test_dataset_tinterp_00, batch_size=1, collate_fn=merge_fn, num_workers=num_workers, shuffle=False)

val_dataloader_tinterp_10 = DataLoader(val_dataset_tinterp_10, batch_size=batch_size, collate_fn=merge_fn, num_workers=num_workers, shuffle=False)
test_dataloader_tinterp_10 = DataLoader(test_dataset_tinterp_10, batch_size=1, collate_fn=merge_fn, num_workers=num_workers, shuffle=False)

tren_norm_first = True
tren_add_tailnorm = True
pooling_type = "none"

model_tinterp_00 = TransformerEnISLR(
    in_channels=in_channels,
    inter_channels=inter_channels,
    out_channels=out_channels,
    activation=activation,
    pooling_type=pooling_type,
    tren_num_layers=tren_num_layers,
    tren_num_heads=tren_num_heads,
    tren_dim_ffw=tren_dim_ffw,
    tren_dropout_pe=tren_dropout_pe,
    tren_dropout=tren_dropout,
    tren_layer_norm_eps=tren_layer_norm_eps,
    tren_norm_first=tren_norm_first,
    tren_add_bias=tren_add_bias,
    tren_add_tailnorm=tren_add_tailnorm)
print(model_tinterp_00)

loss_fn = nn.CrossEntropyLoss(reduction="mean")
optimizer = torch.optim.Adam(model_tinterp_00.parameters(), lr=lr)

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

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

    train_losses = train_loop(train_dataloader_tinterp_00, model_tinterp_00, loss_fn, optimizer, device)
    val_loss_tinterp_00 = val_loop(val_dataloader_tinterp_00, model_tinterp_00, loss_fn, device)
    val_losses_tinterp_00.append(val_loss_tinterp_00)
    val_loss_tinterp_10 = val_loop(val_dataloader_tinterp_10, model_tinterp_00, loss_fn, device)
    val_losses_tinterp_10.append(val_loss_tinterp_10)

    if (epoch+1) % eval_every_n_epochs == 0:
        acc_tinterp_00 = test_loop(test_dataloader_tinterp_00, model_tinterp_00, device)
        test_accs_tinterp_00.append(acc_tinterp_00)
        acc_tinterp_10 = test_loop(test_dataloader_tinterp_10, model_tinterp_00, device)
        test_accs_tinterp_10.append(acc_tinterp_10)
train_losses_tinterp_00 = np.array(train_losses)
val_losses_tinterp_00_vs_tinterp_00 = np.array(val_losses_tinterp_00)
test_accs_tinterp_00_vs_tinterp_00 = np.array(test_accs_tinterp_00)
val_losses_tinterp_00_vs_tinterp_10 = np.array(val_losses_tinterp_10)
test_accs_tinterp_00_vs_tinterp_10 = np.array(test_accs_tinterp_10)

print(f"Minimum validation loss (Train: W0-Tinterp, Test: WO-Tinterp):{val_losses_tinterp_00_vs_tinterp_00.min()} at {np.argmin(val_losses_tinterp_00_vs_tinterp_00)+1} epoch.")
print(f"Maximum accuracy (Train: W0-Tinterp, Test: WO-Tinterp):{test_accs_tinterp_00_vs_tinterp_00.max()} at {np.argmax(test_accs_tinterp_00_vs_tinterp_00)*eval_every_n_epochs+1} epoch.")

print(f"Minimum validation loss (Train: W0-Tinterp, Test: W-Tinterp):{val_losses_tinterp_00_vs_tinterp_10.min()} at {np.argmin(val_losses_tinterp_00_vs_tinterp_10)+1} epoch.")
print(f"Maximum accuracy (Train: W0-Tinterp, Test: W-Tinterp):{test_accs_tinterp_00_vs_tinterp_10.max()} at {np.argmax(test_accs_tinterp_00_vs_tinterp_10)*eval_every_n_epochs+1} epoch.")