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 collections
import copy
import json
import math
import os
import sys
from functools import partial
from pathlib import Path

# Third party's modules
import numpy as np

from pydantic import (
    BaseModel,
    ConfigDict,
    Field)

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_fullbody_swappairs
)
from modules_gislr.layers.misc import (
    GPoolRecognitionHeadSettings,
    Identity,
    Zero,
    apply_norm,
    create_norm
)
from modules_gislr.train_functions import (
    test_loop,
    val_loop,
    train_loop
)
from modules_gislr.transforms import (
    PartsBasedNormalization,
    RandomAffineTransform2D,
    RandomClip,
    RandomDropParts,
    RandomHorizontalFlip,
    RandomNoise,
    RandomTimeWarping,
    ReplaceNan,
    SelectLandmarksAndFeature,
    ToTensor
)
from modules_gislr.utils import (
    make_causal_mask,
    make_san_mask,
    select_reluwise_activation
)

class ConfiguredModel(BaseModel):
    model_config = ConfigDict(extra="forbid")

class LinearFeatureExtractorSettings(ConfiguredModel):
    fext_type: str = "linear"
    in_channels: int = 64
    out_channels: int = 64
    norm_type: str = Field(default="batch", pattern=r"batch|layer")
    norm_eps: float = 1e-5
    activation: str = Field(default="relu",
        pattern=r"relu|gelu|swish|silu|mish|geluacc|tanhexp")
    add_bias: bool = True
    dropout: float = 0.1
    channel_first: bool = True
    add_residual: bool = True

    def build_layer(self):
        return LinearFeatureExtractor(self)

class LinearFeatureExtractor(nn.Module):
    def __init__(self,
                 settings):
        super().__init__()
        assert isinstance(settings, LinearFeatureExtractorSettings)
        self.settings = settings
        self.channel_first = settings.channel_first

        self.linear = nn.Linear(
            in_features=settings.in_channels,
            out_features=settings.out_channels,
            bias=settings.add_bias)

        self.norm = create_norm(
            settings.norm_type, settings.out_channels, settings.norm_eps,
            settings.add_bias)

        self.activation = select_reluwise_activation(settings.activation)

        self.dropout = nn.Dropout(p=settings.dropout)

        if settings.add_residual and settings.in_channels == settings.out_channels:
            self.residual = Identity()
        else:
            self.residual = Zero()

    def forward(self,
                feature):
        shape = feature.shape
        if self.channel_first:
            # `[N, C, T] -> [N, T, C]`
            if len(shape) == 3:
                feature = feature.permute([0, 2, 1])
            # `[N, C, T, *] -> [N, T, C, *] -> [N, T, C']`
            elif len(shape) == 4:
                feature = feature.permute([0, 2, 1, 3])
                feature = feature.reshape([shape[0], shape[2], -1])
            # `[N, C, T, *, *] -> [N, T, C, *, *] -> [N, T, C']`
            elif len(shape) == 5:
                feature = feature.permute([0, 2, 1, 3, 4])
                feature = feature.reshape([shape[0], shape[2], -1])
            else:
                raise NotImplementedError(f"Unsupported feature shape:{shape}.")
        res = self.residual(feature)
        feature = self.linear(feature)
        feature = apply_norm(self.norm, feature, channel_first=False)
        feature = self.activation(feature)
        feature = self.dropout(feature)
        feature = feature + res
        if self.channel_first:
            # `[N, T, C] -> [N, C, T]`
            feature = feature.permute([0, 2, 1])
        return feature

class CNN1DFeatureExtractorSettings(ConfiguredModel):
    fext_type: str = "cnn1d"
    in_channels: int = 64
    out_channels: int = 64
    # Conv1D settings.
    kernel_size: int = Field(default=3, ge=3)
    stride: int = Field(default=1, ge=1)
    padding_mode: str = "zeros"
    norm_type: str = Field(default="batch",
        pattern=r"layer|batch")
    norm_eps: float = 1e-5
    activation: str = Field(default="swish",
        pattern=r"relu|gelu|swish|silu|mish|geluacc|tanhexp")
    causal: bool | None = False
    add_residual: bool = True
    add_bias: bool = True
    dropout: float = 0.1

    def model_post_init(self, __context):
        message = f"kernel_size:{self.kernel_size} must be the odd number."
        assert (self.kernel_size - 1) % 2 == 0, message

    def build_layer(self):
        return CNN1DFeatureExtractor(self)

class CNN1DFeatureExtractor(nn.Module):
    def __init__(self,
                 settings):
        super().__init__()
        assert isinstance(settings, CNN1DFeatureExtractorSettings)
        self.settings = settings
        self.causal = settings.causal

        if settings.causal:
            self.padding = (settings.kernel_size - 1)
        else:
            self.padding = (settings.kernel_size - 1) // 2

        self.conv_module = self._build_conv_module(settings, self.padding)
        self.norm = create_norm(settings.norm_type, settings.out_channels, settings.norm_eps)
        self.activation = select_reluwise_activation(settings.activation)
        self.dropout = nn.Dropout(p=settings.dropout)

        if settings.add_residual:
            if settings.in_channels == settings.out_channels:
                if settings.stride == 1:
                    self.residual = Identity()
                else:
                    self.residual = nn.MaxPool1d(
                        kernel_size=settings.kernel_size,
                        stride=settings.stride,
                        padding=self.padding)
            else:
                self.residual = nn.Conv1d(
                    in_channels=settings.in_channels,
                    out_channels=settings.out_channels,
                    kernel_size=1,
                    stride=settings.stride)
        else:
            self.residual = Zero()


    def _build_conv_module(self, settings, padding):
        dict_modules = collections.OrderedDict([
            # Point-wise.
            ("pconv",
             nn.Conv1d(
                in_channels=settings.in_channels,
                out_channels=settings.out_channels,
                kernel_size=1,
                stride=1,
                padding=0,
                bias=settings.add_bias)),
            # Depth-wise
            ("dconv",
             nn.Conv1d(
                in_channels=settings.out_channels,
                out_channels=settings.out_channels,
                kernel_size=settings.kernel_size,
                stride=settings.stride,
                padding=padding,
                groups=settings.out_channels,
                padding_mode=settings.padding_mode,
                bias=settings.add_bias))])

        conv_module = nn.Sequential(dict_modules)
        return conv_module

    def forward(self,
                feature):
        shape = feature.shape
        if len(shape) == 4:
            # `[N, C, T, *] -> [N, C, *, T] -> [N, C', T]`
            feature = feature.permute([0, 1, 3, 2])
            feature = feature.reshape([shape[0], -1, shape[2]])
        elif len(shape) == 5:
            # `[N, C, T, *, *] -> [N, C, *, *, T] -> [N, C', T]`
            feature = feature.permute([0, 1, 3, 4, 2])
            feature = feature.reshape([shape[0], -1, shape[2]])

        res = self.residual(feature)
        feature = self.conv_module(feature)
        feature = apply_norm(self.norm, feature, channel_first=True)
        feature = self.activation(feature)
        feature = self.dropout(feature)
        feature = feature + res
        return feature

def create_fext_module(settings):
    if settings is None:
        fext_module = nn.ModuleList(
            [Identity()])
    else:
        fext_module = nn.ModuleList(
            [setting.build_layer() for setting in settings])
    return fext_module

# 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"]
swap_pairs, _ = get_fullbody_swappairs()

trans_select_feature = SelectLandmarksAndFeature(landmarks=use_landmarks, features=use_features)
trans_repnan = ReplaceNan()
trans_norm = PartsBasedNormalization(align_mode="framewise", scale_mode="unique")

trans_hflip = RandomHorizontalFlip(
    apply_ratio=0.5,
    num_joints=len(use_landmarks),
    swap_pairs=swap_pairs,
    flip_center=0.5,
    feature_dim=len(use_features),
    include_conf=False
)
trans_saffine = RandomAffineTransform2D(
    apply_ratio=0.5,
    center_joints=[0, 2],
    target_joints=np.arange(0, len(use_landmarks)),
    trans_range=[-0.1, 0.1],
    scale_range=[1.0/1.5, 1.5],
    rot_range=[-30, 30],
    skew_range=[-30, 30]
)
trans_snoise = RandomNoise(
    apply_ratio=0.5,
    scale_range=[1e-3, 1e-2],
    scale_unit="asis",
    noise_type="uniform",
    target_joints=np.arange(0, len(use_landmarks)),
    feature_dim=len(use_features),
    include_conf=False
)
trans_clip =  RandomClip(
    apply_ratio=0.5,
    clip_range=[0.4, 0.6],
    min_apply_size=10
)
trans_twarp = RandomTimeWarping(
    apply_ratio=0.5,
    scale_range=[0.5, 2.0],
    min_apply_size=10
)
trans_drop = RandomDropParts(apply_ratio=0.5)

pre_transforms = Compose([trans_select_feature,
                          trans_repnan])

train_transforms = Compose([
    trans_hflip,
    trans_saffine,
    trans_snoise,
    trans_norm,
    trans_clip,
    trans_twarp,
    trans_drop,
    ToTensor()
])
val_transforms = Compose([trans_norm, ToTensor()])
test_transforms = Compose([trans_norm, 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=VOCAB)

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

    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
pffn_channels = inter_channels * 2
out_channels = VOCAB

dict_linear_fext = {
    "fext_type": "linear",
    "in_channels": in_channels,
    "out_channels": inter_channels,
    "norm_type": "layer",
    "norm_eps": 1e-05,
    "activation": "relu",
    "add_bias": True,
    "dropout": 0.1,
    "channel_first": True,
    "add_residual": False
}

linear_fext_settings = [LinearFeatureExtractorSettings.model_validate(dict_linear_fext)]

islr_settings = TransformerEnISLRSettings(
    in_channels=inter_channels,
    inter_channels=inter_channels,
    out_channels=out_channels)

# Overwrite.
islr_settings.encoder_settings.num_layers = 2
islr_settings.encoder_settings.pe_settings.dim_model = inter_channels
islr_settings.enlayer_settings.dim_pffn = pffn_channels
islr_settings = TransformerEnISLRSettings.model_validate(islr_settings.model_dump())

model_linear = islr_settings.build_layer(linear_fext_settings)

# Sanity check.
print(model_linear)
logit = model_linear(feature_origin)
print(logit.shape)

dict_conv_fext = {
    "fext_type": "cnn1d",
    "in_channels": in_channels,
    "out_channels": inter_channels,
    "kernel_size": 3,
    "stride": 1,
    "norm_type": "batch",
    "norm_eps": 1e-5,
    "activation": "relu",
    "causal": False,
    "add_residual": False,
    "add_bias": True,
    "dropout": 0.1
}

conv_fext_settings = [CNN1DFeatureExtractorSettings.model_validate(dict_conv_fext)]
model_cnn1d = islr_settings.build_layer(conv_fext_settings)

# Sanity check.
print(model_cnn1d)
logit = model_cnn1d(feature_origin)
print(logit.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
label_smoothing = 0.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"]

# 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_linear = islr_settings.build_layer(linear_fext_settings)
print(model_linear)

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

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

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

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

    if (epoch+1) % eval_every_n_epochs == 0:
        acc = test_loop(test_dataloader, model_linear, device)
        test_accs.append(acc)
train_losses_linear = np.array(train_losses)
val_losses_linear = np.array(val_losses)
test_accs_linear = np.array(test_accs)

print(f"Minimum validation loss:{val_losses_linear.min()} at {np.argmin(val_losses_linear)+1} epoch.")
print(f"Maximum accuracy:{test_accs_linear.max()} at {np.argmax(test_accs_linear)*eval_every_n_epochs+1} epoch.")