"""
训练MAE模型
MAE模型作为一个特征抽取器仅仅用作对于各个Patch的特征捕获 提前训练好模型即可
"""

import os, random
import cv2
import numpy as np
import torch
import torch.optim as optim
import torch.nn as nn
import tqdm

from torchvision.transforms import transforms
from torchvision.utils import make_grid

from torch.utils.data import dataloader

from lib.models.backbones import mae

MEAN = [0.376, 0.469, 0.293]
STD = [0.166, 0.155, 0.177]

class MAEJointDataset:
    """
    针对MAE训练的数据集
    """
    def __init__(
        self,
        data_root,
        data_file,
        patch_size,
        transform,
        augment,
    ) -> None:
        """
        Args:
            data_root: 各个数据集对应的根目录
            data_file: 数据集中每张图像对应的路径
            patch_size: 作为输入的每张Patch的大小
            transform: 图像变化transform
            augment: 数据增强
        """
        self.data_root = data_root
        self.data_file = data_file
        self.patch_size = patch_size
        self.transform = transform
        self.augment = augment

        self.images_path_list = []
        with open(data_file, encoding="utf-8", mode="r") as f:
            self.images_path_list = list(f.readlines())
        
    def __len__(self):
        return len(self.images_path_list)
    
    def __getitem__(self, index):
        
        img_path = self.images_path_list[index][:-1]
        img_path = os.path.join(self.data_root, img_path)
        img = cv2.imread(img_path)[:,:,:3]
        img = cv2.resize(img, (self.patch_size, self.patch_size), cv2.INTER_AREA)
        # bgr --> rgb
        img = np.ascontiguousarray(img[:, :, ::-1])
        if random.random() > 0.5 and self.augment:
            img = np.fliplr(img)
        # for continuous
        img = np.ascontiguousarray(img)
        # transform
        img = self.transform(img)

        return img

def train():

    global MEAN, STD

    data_root = "/home/featurize/workspace/main/FSoccerMOT/datasets"
    data_file = "./data/mae.train"

    batch_size = 128
    total_epoch = 400
    num_workers = 6

    lr = 1e-23

    patch_size = 128
    patch_num = 16
    embedding_dim = 128
    encoder_layers = 4
    decoder_layers = 2
    mask_ratio = 0.75
    device = "cuda"

    transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(MEAN, STD)
        ]
    )

    dataset = MAEJointDataset(data_root, data_file, patch_size, transform, True)

    # for i, sample in enumerate(dataset):
    #     print(i)
    #     print(sample.size())

    train_loader = dataloader.DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True, drop_last=True)
    model = mae.MAE(image_size=patch_size, patch_num=patch_num, embedding_dim=embedding_dim, encoder_layers=encoder_layers, decoder_layers=decoder_layers, mask_ratio=mask_ratio).to(device)
    crit = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr = lr, weight_decay=1e-5)

    total_loss = 0
    model.train()

    print("=" * 120)
    print(model)
    print("Total Batches:", len(train_loader))
    print("=" * 120)

    warmup_epoches = 20
    for e in range(0, total_epoch):
        for (_, patches) in tqdm.tqdm(enumerate(train_loader), desc="Batch"):

            patches = patches.to(device)
            origin_embeddings, output_embeddings, output_img, forward_indexes = model(patches)
            
            optimizer.zero_grad()
            masked_indexes = forward_indexes[:, int((1 - mask_ratio) * (patch_num **2)):]
            
            masked_origin_embeddings = torch.stack(
                [origin_embeddings[i, masked_indexes[i]] for i in range(batch_size)], dim=0
            )
            masked_output_embeddings = torch.stack(
                [output_embeddings[i, masked_indexes[i]] for i in range(batch_size)], dim=0
            )

            loss = crit(masked_origin_embeddings, masked_output_embeddings)
            loss.backward()
            optimizer.step()

            total_loss += loss

        print("Epoch: %d, MEAN MSE Loss %.6f" % (e, total_loss / len(train_loader)))
        total_loss = 0

        # lr step
        if e < warmup_epoches:
            lr *= 10
            for param_group in optimizer.param_groups:
                param_group['lr'] = lr
                print("Adjust lr to:", lr)

        # save model
        if e % 10 == 0:
            torch.save(model.state_dict(), "mae_pretrain_%d.pth" % (e, ))

            MEAN = torch.tensor(MEAN).view(3,1,1).to(device)
            STD = torch.tensor(STD).view(3,1,1).to(device)
            output_img = output_img * STD
            output_img = output_img + MEAN
            output_img = output_img * 255.0
            output_img = torch.clip(output_img, min = 0.0, max = 255.0)

            all_images = make_grid(output_img, nrow=16)
            all_images = all_images.permute(1,2,0)

            all_images = all_images.cpu().detach().numpy()
            
            # rgb --> bgr
            all_images = np.ascontiguousarray(all_images[:, :, ::-1])
            all_images = all_images.astype(np.uint8)
            cv2.imwrite("output_%d.jpg" % (e,), all_images)
            
def save_tensors_to_imgs(img, name):
    """
    将四维的张量变为图像格式并保存
    """
    mean = torch.tensor(MEAN, requires_grad=False).view(3,1,1).to(img.device)
    std = torch.tensor(STD, requires_grad=False).view(3,1,1).to(img.device)

    output_img = img * std
    output_img = output_img + mean
    output_img = output_img * 255.0
    output_img = torch.clip(output_img, min = 0.0, max = 255.0)

    all_images = make_grid(output_img, nrow=16)
    all_images = all_images.permute(1,2,0)

    all_images = all_images.cpu().detach().numpy()
    
    # rgb --> bgr
    all_images = np.ascontiguousarray(all_images[:, :, ::-1])
    all_images = all_images.astype(np.uint8)
    cv2.imwrite(name, all_images)

def eval():

    data_root = "/home/featurize/workspace/main/FSoccerMOT/datasets"
    data_file = "./data/mae.train"
    model_path = "./mae_pretrain_380.pth"

    batch_size = 128
    num_workers = 6

    patch_size = 128
    patch_num = 16
    embedding_dim = 128
    encoder_layers = 4
    decoder_layers = 2
    mask_ratio = 0.0
    device = "cuda"

    transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(MEAN, STD)
        ]
    )

    dataset = MAEJointDataset(data_root, data_file, patch_size, transform, False)

    train_loader = dataloader.DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True, drop_last=True)
    model = mae.MAE(image_size=patch_size, patch_num=patch_num, embedding_dim=embedding_dim, encoder_layers=encoder_layers, decoder_layers=decoder_layers, mask_ratio=mask_ratio)
    model.load_state_dict(torch.load(model_path))
    model = model.to(device)
    model.eval()

    print("=" * 120)
    print(model)
    print("=" * 120)
    
    with torch.no_grad():
        for (i, patches) in tqdm.tqdm(enumerate(train_loader), desc="Batch"):

            patches = patches.to(device)

            origin_embeddings, output_embeddings, output_img, forward_indexes = model(patches)

            save_tensors_to_imgs(patches, "origin_%d.jpg" % (i, ))
            save_tensors_to_imgs(output_img, "output_%d.jpg" % (i, ))

if __name__ == "__main__":

    # train()
    eval()