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 math
import sys
import time
from functools import partial
from pathlib import Path
from typing import (
    Any,
    Dict
)

# Third party's modules
import numpy as np

import torch
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_padded_batch)
from modules_gislr.defines import (
    get_fullbody_landmarks,
    get_parts_connections
)
from modules_gislr.train_functions import (
    test_loop,
    train_loop,
    val_loop
)
from modules_gislr.draw_functions import (
    draw_tracking,
    showvideo_on_browser
)
from modules_gislr.layers import (
    SimpleISLR
)
from modules_gislr.transforms import (
    ReplaceNan,
    SelectLandmarksAndFeature,
    ToTensor
)

class PartsBasedNormalization():
    def __init__(self,
                 face_head=0, face_num=76, face_origin=[0, 2], face_unit1=[7], face_unit2=[42],
                 lhand_head=76, lhand_num=21, lhand_origin=[0, 2, 5, 9, 13, 17], lhand_unit1=[0], lhand_unit2=[2, 5, 9, 13, 17],
                 pose_head=76+21, pose_num=12, pose_origin=[0, 1], pose_unit1=[0], pose_unit2=[1],
                 rhand_head=76+21+12, rhand_num=21, rhand_origin=[0, 2, 3, 9, 13, 17], rhand_unit1=[0], rhand_unit2=[2, 5, 9, 13, 17],
                 align_mode="framewise", scale_mode="framewise") -> None:
        assert align_mode in ["framewise", "unique"]
        assert scale_mode in ["framewise", "unique", "none"]
        self.align_mode = align_mode
        self.scale_mode = scale_mode

        self.face_head = face_head
        self.face_num = face_num
        self.face_origin = face_origin
        self.face_unit1 = face_unit1
        self.face_unit2 = face_unit2

        self.lhand_head = lhand_head
        self.lhand_num = lhand_num
        self.lhand_origin = lhand_origin
        self.lhand_unit1 = lhand_unit1
        self.lhand_unit2 = lhand_unit2

        self.pose_head = pose_head
        self.pose_num = pose_num
        self.pose_origin = pose_origin
        self.pose_unit1 = pose_unit1
        self.pose_unit2 = pose_unit2

        self.rhand_head = rhand_head
        self.rhand_num = rhand_num
        self.rhand_origin = rhand_origin
        self.rhand_unit1 = rhand_unit1
        self.rhand_unit2 = rhand_unit2

    def _gen_tmask(self, feature):
        tmask = feature == 0.0
        tmask = np.all(tmask, axis=(0, 2))
        tmask = np.logical_not(tmask.reshape([1, -1, 1]))
        return tmask

    def _calc_origin(self, feature, origin_lm):
        # `[C, T, J] -> [C, T, 1]`
        origin = feature[:, :, origin_lm].mean(axis=-1, keepdims=True)
        if self.align_mode == "unique":
            # `[C, T, 1] -> [C, 1, 1]`
            mask = self._gen_tmask(origin)
            mask = mask.reshape([mask.shape[1]])
            if mask.any():
                origin = origin[:, mask, :].mean(axis=1, keepdims=True)
            else:
                origin = np.array([0.] * feature.shape[0]).reshape([-1, 1, 1])
        return origin

    def _calc_unit(self, feature, unit_lm1, unit_lm2, unit_range):
        if self.scale_mode == "none":
            return 1.0
        # The frame-wise unit lengths are unstable.
        # So, we calculate average unit length.
        # Extract.
        # `[C, T, J] -> [C, T, 1]`
        unit1 = feature[:, :, unit_lm1].mean(axis=-1)
        unit2 = feature[:, :, unit_lm2].mean(axis=-1)
        # Mean square between target points.
        unit = np.sqrt((unit1 - unit2) ** 2)
        # Norm.
        # `[C, T, J] -> [1, T, 1]`
        unit = np.linalg.norm(unit, axis=0)
        if self.scale_mode == "framewise":
            unit = unit.reshape([1, unit.shape[0], 1])
            unit[unit <= 0] = 1.0
            unit[np.isnan(unit)] = 1.0
        else:
            # Calculate average removing undetected frame.
            mask = unit > 0
            if mask.sum() > 0:
                unit = unit[unit > 0].mean()
            else:
                unit = 1.0
            unit = 1.0 if np.isnan(unit).any() else unit
        # Finally, clip extreme values.
        unit = np.clip(unit, a_min=unit_range[0], a_max=unit_range[1])
        return unit

    def _normalize(self, feature, origin_lm, unit_lm1, unit_lm2,
                   unit_range=[1.0e-3, 5.0]):
        tmask = self._gen_tmask(feature)
        origin = self._calc_origin(feature, origin_lm)
        unit = self._calc_unit(feature, unit_lm1, unit_lm2, unit_range)

        _feature = feature - origin
        _feature = _feature / unit
        _feature = _feature * tmask
        return _feature

    def __call__(self,
                 data: Dict[str, Any]) -> Dict[str, Any]:
        feature = data["feature"]
        if self.face_num > 0:
            face = feature[:, :, self.face_head: self.face_head+self.face_num]
            face = self._normalize(face, self.face_origin,
                                   self.face_unit1, self.face_unit2)
            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._normalize(lhand, self.lhand_origin,
                                    self.lhand_unit1, self.lhand_unit2)
            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._normalize(pose, self.pose_origin,
                                   self.pose_unit1, self.pose_unit2)
            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._normalize(rhand, self.rhand_origin,
                                    self.rhand_unit1, self.rhand_unit2)
            feature[:, :, self.rhand_head: self.rhand_head+self.rhand_num] = rhand
        data["feature"] = feature
        return data

# 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()
align_mode = "framewise"

trans_align_whole = PartsBasedNormalization(
                              face_num=0,
                              lhand_num=0,
                              pose_head=0, pose_num=len(use_landmarks), pose_origin=[97, 98], pose_unit1=[97], pose_unit2=[98],
                              rhand_num=0,
                              align_mode=align_mode,
                              scale_mode="none")

trans_align_parts = PartsBasedNormalization(align_mode=align_mode, scale_mode="none")

trans_align_scale_whole = PartsBasedNormalization(
                              face_num=0,
                              lhand_num=0,
                              pose_head=0, pose_num=len(use_landmarks), pose_origin=[97, 98], pose_unit1=[97], pose_unit2=[98],
                              rhand_num=0,
                              align_mode=align_mode, scale_mode="unique")

trans_align_scale_parts = PartsBasedNormalization(align_mode=align_mode, scale_mode="unique")

pre_transforms_origin = Compose([trans_select_feature,
                                 trans_repnan])

pre_transforms_align_whole = Compose([trans_select_feature,
                                      trans_repnan,
                                      trans_align_whole])

pre_transforms_align_parts = Compose([trans_select_feature,
                                      trans_repnan,
                                      trans_align_parts])

pre_transforms_align_scale_whole = Compose([trans_select_feature,
                                           trans_repnan,
                                           trans_align_scale_whole])

pre_transforms_align_scale_parts = Compose([trans_select_feature,
                                           trans_repnan,
                                           trans_align_scale_parts])

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=0)

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

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)

feature_origin = feature_origin.detach().cpu().numpy()
feature_align_w = feature_align_w.detach().cpu().numpy()
feature_align_p = feature_align_p.detach().cpu().numpy()
feature_align_scale_w = feature_align_scale_w.detach().cpu().numpy()
feature_align_scale_p = feature_align_scale_p.detach().cpu().numpy()

_, use_landmarks = get_fullbody_landmarks()
_, face_connections = get_parts_connections(use_landmarks, "face")
_, lhand_connections = get_parts_connections(use_landmarks, "lhand")
_, pose_connections = get_parts_connections(use_landmarks, "pose")
_, rhand_connections = get_parts_connections(use_landmarks, "rhand")

outpath = "./track_origin.mp4"
draw_tracking(feature_origin, outpath,
              face_connections=face_connections,
              lhand_connections=lhand_connections,
              pose_connections=pose_connections,
              rhand_connections=rhand_connections,
              channel_first=True)

# The MP4 file generated by OpenCV can not display correctly.
# So I convert it into WEBM file.
!ffmpeg -i track_origin.mp4 -vcodec vp9 -y track_origin.webm

showvideo_on_browser(outpath.replace(".mp4", ".webm"))

# Set common parameters.
batch_size = 32
load_into_ram = True
test_pid = 16069
num_workers = 1

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"]

# 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

model = SimpleISLR(in_channels, out_channels)
print(model)

# Sanity check.
logit = model(torch.from_numpy(feature_origin))
print(logit.shape)

loss_fn = nn.CrossEntropyLoss(reduction="mean")
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

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

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

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

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

    train_loop(train_dataloader, model, loss_fn, optimizer, device)
    val_loss = val_loop(val_dataloader, model, loss_fn, device)
    val_losses.append(val_loss)

    if (epoch+1) % eval_every_n_epochs == 0:
        acc = test_loop(test_dataloader, model, device)
        test_accs.append(acc)
val_losses_origin = np.array(val_losses)
test_accs_origin = np.array(test_accs)
print(f"Minimum validation loss:{val_losses_origin.min()} at {np.argmin(val_losses_origin)+1} epoch.")
print(f"Maximum accuracy:{test_accs_origin.max()} at {np.argmax(test_accs_origin)*eval_every_n_epochs+1} epoch.")