# hf-experiments
# @author Loreto Parisi (loretoparisi at gmail dot com)
# Copyright (c) 2021 Loreto Parisi (loretoparisi at gmail dot com)

import os,sys
import numpy as np
import natsort
from PIL import Image
import torch
import torchvision.transforms as T

# ResMLP
from res_mlp.res_mlp import ResMLP
# MLP-Mixer
from mlp_mixer.mlp_mixer import MLPMixer
# Perceiver, General Perception with Iterative Attention
from perceiver.perceiver import Perceiver

BASE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)))
sys.path.insert(0, os.path.join(BASE_PATH, '..'))
from lpdutils.lpimagedataset import LPImageDataSet

# Res MLP
res_model = ResMLP(
    image_size = 224,
    patch_size = 16,
    dim = 512,
    depth = 12,
    num_classes = 1000
)

parameters = filter(lambda p: p.requires_grad, res_model.parameters())
parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
print('ResMLP trainable Parameters: %.3fM' % parameters)

# input must have 256 channels
img = torch.randn(1, 3, 224, 224)
pred = res_model(img) # (1, 1000)
print("ResMLP:",pred.shape)

# MLP Mixer
mixer_model = MLPMixer(in_channels=3, 
                image_size=224, 
                patch_size=16, 
                num_classes=1000,
                dim=512, 
                depth=8, 
                token_dim=256, 
                channel_dim=2048)

parameters = filter(lambda p: p.requires_grad, mixer_model.parameters())
parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
print('MLPMixer trainable Parameters: %.3fM' % parameters)

img = torch.ones([1, 3, 224, 224])
pred = mixer_model(img)
print("MLPMixer:", pred.shape) # [B, in_channels, image_size, image_size]

# Perceiver
perceiver_model = Perceiver(
    input_channels = 3,          # number of channels for each token of the input
    input_axis = 2,              # number of axis for input data (2 for images, 3 for video)
    num_freq_bands = 6,          # number of freq bands, with original value (2 * K + 1)
    max_freq = 10.,              # maximum frequency, hyperparameter depending on how fine the data is
    depth = 6,                   # depth of net
    num_latents = 256,           # number of latents, or induced set points, or centroids. different papers giving it different names
    latent_dim = 512,            # latent dimension
    cross_heads = 1,             # number of heads for cross attention. paper said 1
    latent_heads = 8,            # number of heads for latent self attention, 8
    cross_dim_head = 64,
    latent_dim_head = 64,
    num_classes = 1000,          # output number of classes
    attn_dropout = 0.,
    ff_dropout = 0.,
    weight_tie_layers = False,   # whether to weight tie layers (optional, as indicated in the diagram)
    fourier_encode_data = True,  # whether to auto-fourier encode the data, using the input_axis given. defaults to True, but can be turned off if you are fourier encoding the data yourself
    self_per_cross_attn = 2      # number of self attention blocks per cross attention
)
parameters = filter(lambda p: p.requires_grad, perceiver_model.parameters())
parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
print('Perceiver trainable Parameters: %.3fM' % parameters)
img = torch.randn(1, 224, 224, 3) # 1 imagenet image, pixelized
pred = perceiver_model(img) # (1, 1000)
print("Perceiver:", pred.shape)

batch_size = 2
num_workers = 2

# load local dataset
my_dataset = LPImageDataSet(os.path.join(os.path.dirname(
    os.path.abspath(__file__)), '..', 'data', 'imagenet'), transform=LPImageDataSet.transform)
train_loader = torch.utils.data.DataLoader(my_dataset, 
                                batch_size=batch_size,
                                shuffle=True, 
                                num_workers=num_workers, 
                                drop_last=True,
                                collate_fn=LPImageDataSet.collate_fn)
# model predict
for idx, img in enumerate(train_loader):
    # [1, 3, 224, 224]
    print(idx, img.shape)
    pred = res_model(img)
    print(f"{res_model.__class__.__name__} pred:", pred.shape)
    pred = mixer_model(img)
    print(f"{mixer_model.__class__.__name__} pred:", pred.shape)
    img = np.transpose(img,(0,3,2,1)) # [1, 3, 224, 224] -> [1, 224, 224, 3]
    print(img.shape)
    pred = perceiver_model(img)
    print(f"{perceiver_model.__class__.__name__} pred:", pred.shape)