import argparse
import datetime
import os
import traceback

import kornia
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader
from tqdm.autonotebook import tqdm

import models
from datasets import LowLightFDataset, LowLightFDatasetEval
from models import PSNR, SSIM, CosineLR
from tools import SingleSummaryWriter
from tools import saver, mutils

def get_args():
    parser = argparse.ArgumentParser('Breaking Downing the Darkness')
    parser.add_argument('--num_gpus', type=int, default=1, help='number of gpus being used')
    parser.add_argument('--num_workers', type=int, default=12, help='num_workers of dataloader')
    parser.add_argument('--batch_size', type=int, default=1, help='The number of images per batch among all devices')
    parser.add_argument('-m1', '--model1', type=str, default='INet',
                        help='Model Name')
    parser.add_argument('-m3', '--model3', type=str, default='INet',
                        help='Model Name')
    parser.add_argument('-m1w', '--model1_weight', type=str, default=None,
                        help='Model Name')
    parser.add_argument('-m3w', '--model3_weight', type=str, default=None,
                        help='Model Name')
    parser.add_argument('-ts', '--targets_split', type=str, default='targets',
                        help='dir of targets')
    parser.add_argument('--comment', type=str, default='default',
                        help='Project comment')
    parser.add_argument('--graph', action='store_true')
    parser.add_argument('--scratch', action='store_true')
    parser.add_argument('--sampling', action='store_true')
    parser.add_argument('--test_on_start', action='store_true')

    parser.add_argument('--lr', type=float, default=0.01)
    parser.add_argument('--no_sche', action='store_true')

    parser.add_argument('--optim', type=str, default='adam', help='select optimizer for training, '
                                                                  'suggest using \'admaw\' until the'
                                                                  ' very final stage then switch to \'sgd\'')
    parser.add_argument('--num_epochs', type=int, default=500)
    parser.add_argument('--val_interval', type=int, default=1, help='Number of epoches between valing phases')
    parser.add_argument('--save_interval', type=int, default=500, help='Number of steps between saving')
    parser.add_argument('--data_path', type=str, default='./data/LOL',
                        help='the root folder of dataset')
    parser.add_argument('--log_path', type=str, default='logs/')
    parser.add_argument('--saved_path', type=str, default='logs/')
    args = parser.parse_args()
    return args

def compute_gradient(img):
    gradx = img[..., 1:, :] - img[..., :-1, :]
    grady = img[..., 1:] - img[..., :-1]
    return gradx, grady

class ModelCANet(nn.Module):
    def __init__(self, model1, model3):
        super().__init__()
        self.color_loss = models.L1Loss()
        self.restor_loss = models.MSSSIML1Loss(channels=3)
        self.model_ianet = model1(in_channels=1, out_channels=1)
        self.model_canet = model3(in_channels=6, out_channels=2)
        self.eps = 1e-2
        self.load_weight(self.model_ianet, opt.model1_weight)
        if opt.model3_weight is not None:
            self.load_weight(self.model_canet, opt.model3_weight)
        self.model_ianet.eval()

    def load_weight(self, model, weight_pth):
        state_dict = torch.load(weight_pth)
        ret = model.load_state_dict(state_dict, strict=True)
        print(ret)

    def forward(self, image, image_gt, training=True):
        if training:
            image = image.squeeze(0)
            image_gt = image_gt.repeat(8, 1, 1, 1)

        texture_in, cb_in, cr_in = torch.split(kornia.color.rgb_to_ycbcr(image), 1, dim=1)

        texture_in_down = F.interpolate(texture_in, scale_factor=0.5, mode='bicubic', align_corners=True)
        texture_illumi = self.model_ianet(texture_in_down)
        texture_illumi = F.interpolate(texture_illumi, scale_factor=2, mode='bicubic', align_corners=True)

        texture_en, cb_en, cr_en = torch.split(kornia.color.rgb_to_ycbcr(image / torch.clamp_min(texture_illumi, self.eps)),
                                               1, dim=1)
        texture_gt, cb_gt, cr_gt = torch.split(kornia.color.rgb_to_ycbcr(image_gt), 1, dim=1)

        colors = self.model_canet(torch.cat([texture_in, cb_in, cr_in, texture_gt, cb_en, cr_en], dim=1))

        cb, cr = torch.split(colors, 1, dim=1)

        color_loss1 = self.color_loss(cb, cb_gt)
        color_loss2 = self.color_loss(cr, cr_gt)

        image_out = kornia.color.ycbcr_to_rgb(torch.cat([texture_gt, cb, cr], dim=1))
        restor_loss = self.restor_loss(image_out, image_gt) * 1.0

        psnr = PSNR(image_out, image_gt)
        ssim = SSIM(image_out, image_gt).item()
        return image_out, color_loss1, color_loss2, restor_loss, psnr, ssim

def train(opt):
    if torch.cuda.is_available():
        torch.cuda.manual_seed(42)
    else:
        torch.manual_seed(42)

    timestamp = mutils.get_formatted_time()
    opt.saved_path = opt.saved_path + f'/{opt.comment}/{timestamp}'
    opt.log_path = opt.log_path + f'/{opt.comment}/{timestamp}/tensorboard/'
    os.makedirs(opt.log_path, exist_ok=True)
    os.makedirs(opt.saved_path, exist_ok=True)

    training_params = {'batch_size': opt.batch_size,
                       'shuffle': True,
                       'drop_last': True,
                       'num_workers': opt.num_workers}

    val_params = {'batch_size': 1,
                  'shuffle': False,
                  'drop_last': False,
                  'num_workers': opt.num_workers}

    training_set = LowLightFDataset(os.path.join(opt.data_path, 'train'), targets_split=opt.targets_split,
                                    training=True)
    training_generator = DataLoader(training_set, **training_params)

    val_set = LowLightFDatasetEval(os.path.join(opt.data_path, 'eval'), training=False)
    val_generator = DataLoader(val_set, **val_params)

    model1 = getattr(models, opt.model1)
    model3 = getattr(models, opt.model3)

    model = ModelCANet(model1, model3)
    print(model)

    writer = SingleSummaryWriter(opt.log_path + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/')

    if opt.num_gpus > 0:
        model = model.cuda()
        if opt.num_gpus > 1:
            model = nn.DataParallel(model)

    if opt.optim == 'adam':
        optimizer = torch.optim.Adam(model.model_canet.parameters(), opt.lr)
    else:
        optimizer = torch.optim.SGD(model.model_canet.parameters(), opt.lr, momentum=0.9, nesterov=True)

    scheduler = CosineLR(optimizer, opt.lr, opt.num_epochs)
    epoch = 0
    step = 0
    model.model_canet.train()

    num_iter_per_epoch = len(training_generator)

    try:
        for epoch in range(opt.num_epochs):
            last_epoch = step // num_iter_per_epoch
            if epoch < last_epoch:
                continue

            epoch_loss = []
            progress_bar = tqdm(training_generator)
            if not opt.sampling and not opt.test_on_start:
                for iter, (data, target, name) in enumerate(progress_bar):
                    if iter < step - last_epoch * num_iter_per_epoch:
                        progress_bar.update()
                        continue
                    try:
                        if opt.num_gpus == 1:
                            data, target = data.cuda(), target.cuda()
                        optimizer.zero_grad()

                        image_out, color_loss1, color_loss2, \
                        restor_loss, psnr, ssim = model(data, target, training=True)
                        loss = color_loss1 + color_loss2 + restor_loss
                        loss.backward()
                        optimizer.step()

                        epoch_loss.append(float(loss))

                        progress_bar.set_description(
                            'Step: {}. Epoch: {}/{}. Iteration: {}/{}. color_loss1: {:1.5f}, color_loss2: {:1.5f}, restor_loss: {:1.5f}, psnr: {:.5f}, ssim: {:.5f}'.format(
                                step, epoch, opt.num_epochs, iter + 1, num_iter_per_epoch,
                                color_loss1.item(), color_loss2.item(),
                                restor_loss.item(), psnr, ssim))
                        writer.add_scalar('Loss/train', loss, step)
                        writer.add_scalar('PSNR/train', psnr, step)
                        writer.add_scalar('SSIM/train', ssim, step)

                        # log learning_rate
                        current_lr = optimizer.param_groups[0]['lr']
                        writer.add_scalar('learning_rate', current_lr, step)

                        step += 1

                    except Exception as e:
                        print('[Error]', traceback.format_exc())
                        print(e)
                        continue
                    # scheduler.step(np.mean(epoch_loss))

            if opt.no_sche:
                scheduler.step()

            saver.base_url = os.path.join(opt.saved_path, 'results', '%03d' % epoch)

            if epoch % opt.val_interval == 0:
                model.model_canet.eval()
                loss_ls = []
                psnrs = []
                ssims = []

                for iter, (data, target, name) in enumerate(val_generator):
                    with torch.no_grad():
                        if opt.num_gpus == 1:
                            data = data.squeeze(0).cuda()
                            target = target.cuda()

                        image_out, color_loss1, color_loss2, restor_loss, \
                        psnr, ssim = model(data, target, training=False)
                        saver.save_image(image_out, name=os.path.splitext(name[0])[0] + '_out')
                        saver.save_image(data, name=os.path.splitext(name[0])[0] + '_in')
                        saver.save_image(target, name=os.path.splitext(name[0])[0] + '_gt')

                        loss = restor_loss + color_loss1 + color_loss2
                        loss_ls.append(loss.item())
                        psnrs.append(psnr)
                        ssims.append(ssim)

                loss = np.mean(np.array(loss_ls))
                psnr = np.mean(np.array(psnrs))
                ssim = np.mean(np.array(ssims))

                print(
                    'Val. Epoch: {}/{}. Loss: {:1.5f}, psnr: {:.5f}, ssim: {:.5f}'.format(
                        epoch, opt.num_epochs, loss, psnr, ssim))
                writer.add_scalar('Loss/val', loss, step)
                writer.add_scalar('PSNR/val', psnr, step)
                writer.add_scalar('SSIM/val', ssim, step)

                save_checkpoint(model, f'{opt.model3}_{"%03d" % epoch}_{psnr}_{ssim}_{step}.pth')

                model.model_canet.train()

            opt.test_on_start = False
            if opt.sampling:
                exit(0)
    except KeyboardInterrupt:
        save_checkpoint(model, f'{opt.model3}_{epoch}_{step}_keyboardInterrupt.pth')
        writer.close()
    writer.close()

def save_checkpoint(model, name):
    if isinstance(model, nn.DataParallel):
        torch.save(model.module.model_canet.state_dict(), os.path.join(opt.saved_path, name))
    else:
        torch.save(model.model_canet.state_dict(), os.path.join(opt.saved_path, name))

if __name__ == '__main__':
    opt = get_args()
    train(opt)


import face_alignment
from pytorch_lightning import loggers
from pytorch_lightning import callbacks
import torch
from torch._C import device
from torch.autograd import grad
from torch.utils.data import DataLoader, Dataset, random_split
from .Loss import CoarseLoss
import torch.nn.functional as F
from ..models.encoders import ResnetEncoder
from ..models.FLAME import FLAME
import pytorch_lightning as pl
from ..utils.config import cfg
from ..utils.renderer import SRenderY
from ..utils import util
import numpy as np
from torchvision import transforms
from skimage.io import imread
import pandas as pd
import wandb
from pytorch_lightning.loggers import WandbLogger
import os
from skimage.transform import warp, estimate_transform
from pytorch_lightning.callbacks import ModelCheckpoint
import torch.optim as optim
import cv2
import re

from torch._six import container_abcs, string_classes, int_classes

from ..datasets import detectors
# It is used whenever the size of the inputs is same
torch.backends.cudnn.benchmark = True

_use_shared_memory = False
np_str_obj_array_pattern = re.compile(r'[SaUO]')

error_msg_fmt = "batch must contain tensors, numbers, dicts or lists; found {}"

numpy_type_map = {
    'float64': torch.DoubleTensor,
    'float32': torch.FloatTensor,
    'float16': torch.HalfTensor,
    'int64': torch.LongTensor,
    'int32': torch.IntTensor,
    'int16': torch.ShortTensor,
    'int8': torch.CharTensor,
    'uint8': torch.ByteTensor,
}

class LitAutoEncoder(pl.LightningModule):
    def __init__(self):
        super().__init__()

        self.cfg = cfg
        # self.device = cfg.device
        self.image_size = self.cfg.dataset.image_size

        self.coarseLoss = CoarseLoss()

        # log hyperparameters
        # self.save_hyperparameters()
        # self.train_acc = pl.metrics.Accuracy()
        # self.valid_acc = pl.metrics.Accuracy()
        # self.test_acc = pl.metrics.Accuracy()
        self.landmarkDetector = face_alignment.FaceAlignment(
            face_alignment.LandmarksType._2D, face_detector='sfd', device='cpu')
        # initializing encoder and decoder
        self._create_model(self.cfg.model)
        self._setup_renderer(self.cfg.model)

    def _create_model(self, model_cfg):
        self.n_param = model_cfg.n_shape+model_cfg.n_tex+model_cfg.n_exp + \
            model_cfg.n_pose+model_cfg.n_cam+model_cfg.n_light
        self.num_list = [model_cfg.n_shape, model_cfg.n_tex, model_cfg.n_exp,
                         model_cfg.n_pose, model_cfg.n_cam, model_cfg.n_light]
        self.param_dict = {i: model_cfg.get(
            'n_' + i) for i in model_cfg.param_list}

        # encoders
        self.E_flame = ResnetEncoder(outsize=self.n_param).to(self.device)
        self.initialParams = list(self.E_flame.parameters())[0].clone()

        # decoder
        self.flame = FLAME(model_cfg).to(self.device)

    def _setup_renderer(self, model_cfg):
        self.render = SRenderY(self.image_size, obj_filename=model_cfg.topology_path,
                               uv_size=model_cfg.uv_size).to(self.device)
        # face mask for rendering details
        mask = imread(model_cfg.face_eye_mask_path).astype(np.float32)/255.
        mask = torch.from_numpy(mask[:, :, 0])[None, None, :, :].contiguous()
        self.uv_face_eye_mask = F.interpolate(
            mask, [model_cfg.uv_size, model_cfg.uv_size]).to(self.device)
        mask = imread(model_cfg.face_mask_path).astype(np.float32)/255.
        mask = torch.from_numpy(mask[:, :, 0])[None, None, :, :].contiguous()
        self.uv_face_mask = F.interpolate(
            mask, [model_cfg.uv_size, model_cfg.uv_size]).to(self.device)
        # displacement correction
        fixed_dis = np.load(model_cfg.fixed_displacement_path)
        self.fixed_uv_dis = torch.tensor(fixed_dis).float().to(self.device)
        # mean texture
        mean_texture = imread(model_cfg.mean_tex_path).astype(np.float32)/255.
        mean_texture = torch.from_numpy(mean_texture.transpose(2, 0, 1))[
            None, :, :, :].contiguous()
        self.mean_texture = F.interpolate(
            mean_texture, [model_cfg.uv_size, model_cfg.uv_size]).to(self.device)
        # dense mesh template, for save detail mesh
        self.dense_template = np.load(
            model_cfg.dense_template_path, allow_pickle=True, encoding='latin1').item()

    def decompose_code(self, code, num_dict):
        ''' Convert a flattened parameter vector to a dictionary of parameters
        code_dict.keys() = ['shape', 'tex', 'exp', 'pose', 'cam', 'light']
        '''
        code_dict = {}
        start = 0
        for key in num_dict:
            end = start+int(num_dict[key])
            code_dict[key] = code[:, start:end]
            start = end
            if key == 'light':
                code_dict[key] = code_dict[key].reshape(
                    code_dict[key].shape[0], 9, 3)
        return code_dict

    def encode(self, images):
        parameters = self.E_flame(images)

        codedict = self.decompose_code(parameters, self.param_dict)
        codedict['images'] = images
        return codedict

    def decode(self, codedict):
        vertices, landmarks2d, landmarks3d = self.flame(
            shape_params=codedict['shape'], expression_params=codedict['exp'], pose_params=codedict['pose'])

        # projection
        landmarks2d = util.batch_orth_proj(
            landmarks2d, codedict['cam'])[:, :, :2]
        landmarks2d[:, :, 1:] = -landmarks2d[:, :, 1:]
        landmarks2d = landmarks2d*self.image_size/2 + self.image_size/2
        landmarks3d = util.batch_orth_proj(landmarks3d, codedict['cam'])
        landmarks3d[:, :, 1:] = -landmarks3d[:, :, 1:]
        landmarks3d = landmarks3d*self.image_size/2 + self.image_size/2
        trans_vertices = util.batch_orth_proj(vertices, codedict['cam'])
        trans_vertices[:, :, 1:] = -trans_vertices[:, :, 1:]

        # render shape
        shape_image = self.render.render_shape(vertices, trans_vertices)

        outputDict = {
            'landmarks2d': landmarks2d,
            'landmarks3d': landmarks3d,
            'trans_vertices': trans_vertices,
            'vertices': vertices,
            'shape_image': shape_image,
            'image': codedict['images']
        }
        return outputDict

    def forward(self, images):
        codedict = self.encode(images)
        decodedDict = self.decode(codedict)
        # print('---------------------Parameters---------------------------')
        # print(self.parameters())
        # print(self._current_dataloader_idx)
        # util.show_landmarks(torch.squeeze(codedict["images"]),torch.squeeze(decodedDict['landmarks2d']))
        # util.showImage(decodedDict['shape_image'])

        return codedict, decodedDict

    def loss(self, codedict, decodedDict, landmarksOrig):

        # landmark regularization loss
        # print(landmarksOrig.shape)
        # print(decodedDict['landmarks3d'].shape)
        lmk = self.coarseLoss.lmkLoss(
            landmarksOrig, decodedDict['landmarks2d'])

        # eye loss
        eyeLoss = self.coarseLoss.eyeLoss(
            landmarksOrig, decodedDict['landmarks2d'])

        # photometric loss
        phLoss = self.coarseLoss.photometricLoss(
            decodedDict['image'], decodedDict['shape_image'], landmarksOrig)

        #  identity loss
        idLoss = self.coarseLoss.identityLoss(
            decodedDict['image'], decodedDict['shape_image'])

        """
            Add shape consistency
        """
        reg = self.coarseLoss.regularization(codedict)

        lossTotal = lmk+eyeLoss+idLoss+phLoss+reg
        return lossTotal

    def training_step(self, batch, batch_idx):
        images = batch['image']
        # print(len(images))
        self.E_flame.train()
        trainLoss = 0
        for i in range(len(images)):
            codedict, decodedDict = self.forward(
                images[i].view(-1, 3, 224, 224).float())
            img = np.transpose(torch.squeeze(
                images[i]).cpu().detach().numpy(), (1, 2, 0))
            try:
                grndLandmarks = self.landmarkDetector.get_landmarks(img*255)
                grndLandmarks = grndLandmarks[0]
                grndLandmarks = torch.tensor(grndLandmarks).to(device='cuda')
                grndLandmarks = grndLandmarks.view(-1, 68, 2)
            except Exception as e:
                print(e)
                continue
            trainLoss += self.loss(codedict, decodedDict, grndLandmarks)
        # trainLoss = trainLoss/len(images)
        self.log('train_loss', trainLoss, on_epoch=True)
        return trainLoss

    def validation_step(self, batch, batch_idx):
        images = batch['image']
        self.E_flame.eval()
        validLoss = 0
        for i in range(len(images)):
            codedict, decodedDict = self.forward(
                images[i].view(-1, 3, 224, 224).float())
            img = np.transpose(torch.squeeze(
                images[i]).cpu().detach().numpy(), (1, 2, 0))
            try:
                grndLandmarks = self.landmarkDetector.get_landmarks(img*255)
                grndLandmarks = grndLandmarks[0]
                grndLandmarks = torch.tensor(grndLandmarks).to(device='cuda')
                grndLandmarks = grndLandmarks.view(-1, 68, 2)
            except Exception as e:
                print(e)
                continue
            validLoss += self.loss(codedict, decodedDict, grndLandmarks)
        # validLoss = validLoss/len(images)
        self.E_flame.train()
        self.log('valid_loss', validLoss, on_epoch=True)

    def test_step(self, batch, batch_idx):
        pass

    def configure_optimizers(self):
        # print(self.parameters)
        return optim.Adam(self.E_flame.parameters(), lr=1e-4, amsgrad=False)

    # def optimizer_step(self, *args, **kwargs):
    #     old_state = list(self.parameters())[0].clone()
    #     super().optimizer_step(*args, **kwargs)
    #     new_state = list(self.parameters())[0].clone()
    #     print(torch.equal(old_state.data, new_state.data))

class MyDataModule(pl.LightningDataModule):

    def __init__(self, batch_size: int = 4):
        super().__init__()
        self.dataset = Dataset_3D(csv_file=os.getcwd()+'/decalib/datasets/data.csv',
                                  root_dir=os.getcwd()+'/decalib/datasets/300W_LP',
                                  transform=transforms.Compose([
                                      Preprocess()
                                  ]))
        self.batch_size = batch_size
        # update size of random split
        l = len(self.dataset)
        valLen = l-int(0.3*l)
        # print(valLen)
        self.train, self.val = random_split(self.dataset, [valLen, l-valLen])

    def collate_fn(self, batch):
        '''
        collate_fn (callable, optional): merges a list of samples to form a mini-batch.
            This function refers to touch's default_collate function, which is also the default proofreading method of DataLoader. When the batch contains data such as None,
            The default default_collate squad method will get an error
            One solution is:
            Determine whether the image in the batch is None. If it is None, it is cleared in the original batch, so that it can avoid errors in the iteration.
        :param batch:
        :return:
        '''
        r"""Puts each data field into a tensor with outer dimension batch size"""
        # Add here: Determine whether image is None. If it is None, it will be cleared in the original batch, so you can avoid errors in iteration.
        # print(type(batch))
        # print(batch)
        if isinstance(batch, list):
            batch = [(image) for (image) in batch if image is not None]
        if batch == []:
            return (None)

        elem_type = type(batch[0])
        if isinstance(batch[0], torch.Tensor):
            out = None
            if _use_shared_memory:
                # If we're in a background process, concatenate directly into a
                # shared memory tensor to avoid an extra copy
                numel = sum([x.numel() for x in batch])
                storage = batch[0].storage()._new_shared(numel)
                out = batch[0].new(storage)
            return torch.stack(batch, 0, out=out)
        elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
                and elem_type.__name__ != 'string_':
            elem = batch[0]
            if elem_type.__name__ == 'ndarray':
                # array of string classes and object
                if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
                    raise TypeError(error_msg_fmt.format(elem.dtype))

                return self.collate_fn([torch.from_numpy(b) for b in batch])
            if elem.shape == ():  # scalars
                py_type = float if elem.dtype.name.startswith('float') else int
                return numpy_type_map[elem.dtype.name](list(map(py_type, batch)))
        elif isinstance(batch[0], float):
            return torch.tensor(batch, dtype=torch.float64)
        elif isinstance(batch[0], int_classes):
            return torch.tensor(batch)
        elif isinstance(batch[0], string_classes):
            return batch
        elif isinstance(batch[0], container_abcs.Mapping):
            return {key: self.collate_fn([d[key] for d in batch]) for key in batch[0]}
        elif isinstance(batch[0], tuple) and hasattr(batch[0], '_fields'):  # namedtuple
            return type(batch[0])(*(self.collate_fn(samples) for samples in zip(*batch)))
        elif isinstance(batch[0], container_abcs.Sequence):
            transposed = zip(*batch)  # ok
            return [self.collate_fn(samples) for samples in transposed]

        raise TypeError((error_msg_fmt.format(type(batch[0]))))

    def train_dataloader(self):
        return DataLoader(dataset=self.train, batch_size=self.batch_size, num_workers=1, collate_fn=self.collate_fn)

    def val_dataloader(self):
        return DataLoader(dataset=self.val, batch_size=self.batch_size, num_workers=1, collate_fn=self.collate_fn)

    def test_dataloader(self):
        pass

class Preprocess(object):
    def __init__(self, iscrop=True, crop_size=224, scale=1.25, face_detector='fan'):
        self.crop_size = crop_size
        self.scale = scale
        self.iscrop = iscrop
        self.resolution_inp = crop_size
        self.face_detector = face_detector
        if(face_detector == 'fan'):
            self.face_detector = detectors.FAN()

    def __call__(self, sample):
        image = sample['image']
        if len(image.shape) == 2:
            image = image[:, :, None].repeat(1, 1, 3)
        if len(image.shape) == 3 and image.shape[2] > 3:
            image = image[:, :, :3]

        h, w, _ = image.shape

        if self.iscrop:
            bbox, bbox_type = self.face_detector.run(image)
            # print('---------------------------------Bounding box and It type-------------------------')
            # print(bbox,bbox_type)
            if len(bbox) < 4:
                print('no face detected! run original image')
                left = 0
                right = h-1
                top = 0
                bottom = w-1
            else:
                left = bbox[0]
                right = bbox[2]
                top = bbox[1]
                bottom = bbox[3]
            old_size, center = self.bbox2point(
                left, right, top, bottom, type=bbox_type)
            size = int(old_size*self.scale)
            src_pts = np.array([[center[0]-size/2, center[1]-size/2], [center[0] -
                               size/2, center[1]+size/2], [center[0]+size/2, center[1]-size/2]])
        else:
            src_pts = np.array([[0, 0], [0, h-1], [w-1, 0]])

        DST_PTS = np.array(
            [[0, 0], [0, self.resolution_inp - 1], [self.resolution_inp - 1, 0]])
        tform = estimate_transform('similarity', src_pts, DST_PTS)

        image = image/255.

        dst_image = warp(image, tform.inverse, output_shape=(
            self.resolution_inp, self.resolution_inp))
        dst_image = dst_image.transpose(2, 0, 1)

        return {'image': torch.tensor(dst_image).float()}

    def bbox2point(self, left, right, top, bottom, type='bbox'):
        """
        bbox from detector and landmarks are different
        """
        if type == 'kpt68':
            old_size = (right - left + bottom - top)/2*1.1
            center = np.array([right - (right - left) / 2.0,
                              bottom - (bottom - top) / 2.0])
        elif type == 'bbox':
            old_size = (right - left + bottom - top)/2
            center = np.array([right - (right - left) / 2.0,
                              bottom - (bottom - top) / 2.0 + old_size*0.12])
        else:
            raise NotImplementedError
        return old_size, center

class Dataset_3D(Dataset):
    def __init__(self, csv_file, root_dir, transform=None):
        self.landmarks_frame = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.to_list()
        img_name = os.path.join(
            self.root_dir, self.landmarks_frame.iloc[idx, 0])
        image = cv2.imread(img_name)
        sample = {'image': image}

        if self.transform:
            try:
                sample = self.transform(sample)
                return sample
            except Exception as e:
                print(e)

    def __len__(self):
        return len(self.landmarks_frame)

if __name__ == "__main__":
    # ---------------Uncomment this to run training------------#
    # torch.multiprocessing.set_start_method('spawn')
    model = LitAutoEncoder()
    dm = MyDataModule()
    checkpoint_callback = ModelCheckpoint(
        monitor='valid_loss',
        dirpath='savedModel/',
        filename='decaCoarse_{epoch}',
        save_top_k=3,
        mode='min',
    )
    wandb_logger = WandbLogger()
    wandb.init()

    """Uncomment and specify checkpoint when you want to reuse previously trained model"""
    # trainer = pl.Trainer(gpus=1, logger=wandb_logger,
    #                      resume_from_checkpoint="/home/nandwalritik/3d-face/savedModel/decaCoarse_epoch=2.ckpt  ", callbacks=[checkpoint_callback])

    """Train from scratch"""
    trainer = pl.Trainer(gpus=1, logger=wandb_logger,
                         callbacks=[checkpoint_callback])

    trainer.fit(model, dm)

    # face_dataset = Dataset_3D(csv_file='/home/nandwalritik/3d-face/decalib/datasets/data.csv',
    #                           root_dir='/home/nandwalritik/3d-face/decalib/datasets/300W_LP',
    #                           transform=transforms.Compose([
    #                             Preprocess()
    #                           ]))

    # for i in range(len(face_dataset)):
    #     if i == 10:
    #         break
    #     sample = face_dataset[i]
    #     print(i, sample['image'].size())
    #     util.showImage(sample['image'])


# Copyright (c) 2020-2022, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch as th
import math
import numpy as np
import os
from scipy import ndimage

from checkpoint_quantization import checkpoint_quantization

L_ROOT   = 'Transformer/encoderblock_{}'
ATT_Q    = 'MultiHeadDotProductAttention_1/query'
ATT_K    = 'MultiHeadDotProductAttention_1/key'
ATT_V    = 'MultiHeadDotProductAttention_1/value'
ATT_OUT  = 'MultiHeadDotProductAttention_1/out'
ATT_NORM = 'LayerNorm_0'
FFN_NORM = 'LayerNorm_2'
FFN_IN   = 'MlpBlock_3/Dense_0'
FFN_OUT  = 'MlpBlock_3/Dense_1'

def np2th(weights, is_conv=False):
    if is_conv:
        """ convert HWIO to OIHW."""
        weights = weights.transpose([3, 2, 0, 1])
    return th.from_numpy(weights).contiguous()

class ViTINT8WeightLoader(object):
    def __init__(self, layer_num, img_size, patch_size, weight_dict=None, classifier='token'):
        """weights need be a pytorch state_dict of swin transformer model"""

        pre_layer_weight_names = [
                                 'transformer.embeddings.patch_embeddings.weight',
                                 'transformer.embeddings.patch_embeddings.bias',
                                 'transformer.embeddings.cls_token',
                                 'transformer.embeddings.position_embeddings'
        ]
        self.layer_num = layer_num
        self.weights = []
        self.int8 = False
        if weight_dict is None:
            print("[ERROR][SwinTransformerWeights::__init__] weights should not be empty!")
            exit(-1)

        self._generated_weights = False

        for name in pre_layer_weight_names:
            if name not in weight_dict.keys():
                print("Unsupport weight file: Missing weights %s" % name)

            th_weight = weight_dict[name]
            if name.split('.')[-1] == "pos_embedding":
                posemb_new_size = pow(img_size//patch_size, 2) + 1
                if th_weight.size(1) != posemb_new_size:
                    print("load_pretrained: resized variant: %s to %s" % (th_weight.size(1), posemb_new_size))
                    ntok_new = posemb_new_size

                    if classifier == "token":
                        posemb_tok, posemb_grid = th_weight[:, :1], th_weight[0, 1:]
                        ntok_new -= 1
                    else:
                        posemb_tok, posemb_grid = th_weight[:, :0], th_weight[0]

                    gs_old = int(np.sqrt(len(posemb_grid)))
                    gs_new = int(np.sqrt(ntok_new))
                    print('load_pretrained: grid-size from %s to %s' % (gs_old, gs_new))
                    posemb_grid = posemb_grid.reshape(gs_old, gs_old, -1)

                    zoom = (gs_new / gs_old, gs_new / gs_old, 1)
                    posemb_grid = ndimage.zoom(posemb_grid, zoom, order=1)
                    posemb_grid = posemb_grid.reshape(1, gs_new * gs_new, -1)
                    posemb = np.concatenate([posemb_tok, posemb_grid], axis=1)
                    th_weight = np2th(posemb)
                    weight_dict[name] = th_weight

        self.weights = weight_dict

    # def load_weights(self, weight_path:str):
    #     suffix = weight_path.split('.')[-1]
    #     if suffix != 'pth':
    #         print("Unsupport weight file: Unrecognized format %s " % suffix)
    #         exit(-1)
    #     return th.load(weight_path)

    def to_cuda(self):
        if not self.int8:
            for k, v in self.weights.items():
                self.weights[k] = v.cuda()
        else:
            h_scale_list = {}
            for k, v in self.weights.items():
                if "amaxList" in k:
                    k_h = k.replace("amaxList", "h_amaxList")
                    h_scale_list[k_h] = v
                self.weights[k] = v.cuda()
            for k, v in h_scale_list.items():
                self.weights[k] = v

    def to_half(self):
        for k, v in self.weights.items():
            self.weights[k] = v.half()

    def listed_weights(self):
        ret = []
        for k, v in self.weights.items():
            if k.split('.')[-1] == '_amax' or k.endswith('amaxList'):
                continue
            if k.split('.')[0] == 'head':
                continue
            # print(k, v.type())
            ret.append(v)

        for i in range(self.layer_num):
            name = 'transformer.encoder.layer.{}.amaxList'.format(i)
            ret.append(self.weights[name])
            # print(name, self.weights[name].type())
            name = 'transformer.encoder.layer.{}.h_amaxList'.format(i)
            ret.append(self.weights[name])
            # print(name, self.weights[name].type())

        return ret

    def to_int8(self, ths_path='../../../lib/libpyt_vit.so'):
        # print(self.weights.keys())
        if 'transformer.encoder.layer.0.attn.query._input_quantizer._amax' not in self.weights:
            raise RuntimeError("There is no quantization node in the checkpoint, cannot be quantized to int8.")
        if self.int8:
            return
        self.int8 = True
        for k, v in self.weights.items():
            if k.endswith('bias') or k.endswith('norm.weight') or 'embeddings' in k:
                self.weights[k] = v.half()
            elif k.endswith('weight'):
                self.weights[k] = v.float().cuda()
            else:
                self.weights[k] = v.float().cpu()
        self.weights = checkpoint_quantization(self.weights, ths_path, verbose=False)
        # print(self.weights.keys())


import os.path as osp

import mmcv
import numpy as np
from mmcv.parallel import DataContainer as DC
from torch.utils.data import Dataset

from .transforms import (ImageTransform, BboxTransform, MaskTransform,
                         SegMapTransform, Numpy2Tensor)
from .utils import to_tensor, random_scale, flip_kp, flip_pose, rot_aa
from .extra_aug import ExtraAugmentation
from .custom import CustomDataset
import os.path as osp
import tensorboardX
import math
import json
import pickle
import matplotlib.pyplot as plt
from mmdetection.mmdet.models.utils.smpl_utils import batch_rodrigues
import random
import cv2
from copy import deepcopy
import scipy
import scipy.misc
# import seaborn as sns
import torch
# import png

denormalize = lambda x: x * np.array([0.229, 0.224, 0.225])[None, None, :] + np.array([0.485, 0.456, 0.406])[None, None,
                                                                             :]
import pycocotools.mask as mask_util

def rot2DPts(x, y, rotMat):
    new_x = rotMat[0, 0] * x + rotMat[0, 1] * y + rotMat[0, 2]
    new_y = rotMat[1, 0] * x + rotMat[1, 1] * y + rotMat[1, 2]
    return new_x, new_y

class H36MDataset(Dataset):
    """Custom dataset for detection.

    Annotation format:
    [
        {
            'filename': 'a.jpg',
            'width': 1280,
            'height': 720,
            'ann': {
                'bboxes': <np.ndarray> (n, 4),
                'labels': <np.ndarray> (n, ),
                'bboxes_ignore': <np.ndarray> (k, 4),
                'labels_ignore': <np.ndarray> (k, 4) (optional field)
            }
        },
        ...
    ]

    The `ann` field is optional for testing.
    """

    # CLASSES = None
    CLASSES = ('Human',)

    def __init__(self,
                 ann_file,
                 img_prefix,
                 img_scale,
                 img_norm_cfg,
                 multiscale_mode='value',
                 size_divisor=None,
                 proposal_file=None,
                 num_max_proposals=1000,
                 flip_ratio=0,
                 with_mask=True,
                 with_crowd=True,
                 with_label=True,
                 with_semantic_seg=False,
                 seg_prefix=None,
                 seg_scale_factor=1,
                 extra_aug=None,
                 resize_keep_ratio=True,
                 test_mode=False,
                 with_shape=True,
                 with_kpts3d=True,
                 with_kpts2d=True,
                 with_pose=True,
                 with_trans=True,
                 max_samples=-1,  # Commonly used in validating
                 noise_factor=0.4,
                 square_bbox=True,
                 rot_factor=0,
                 sample_weight=1,
                 with_dp=False,
                 mosh_path=None,
                 ignore_3d=False,
                 ignore_smpl=False,
                 with_nr=False,
                 sample_by_persons=False,
                 use_poly=False,
                 **kwargs,
                 ):
        # prefix of images path
        self.img_prefix = img_prefix

        # load annotations (and proposals)
        self.img_infos = self.load_annotations(ann_file)
        self.max_samples = max_samples

        if proposal_file is not None:
            self.proposals = self.load_proposals(proposal_file)
        else:
            self.proposals = None
        # filter images with no annotation during training
        if not test_mode and False:
            valid_inds = self._filter_imgs()
            self.img_infos = [self.img_infos[i] for i in valid_inds]
            if self.proposals is not None:
                self.proposals = [self.proposals[i] for i in valid_inds]

        # Select a subset for quick validation
        if self.max_samples > 0:
            self.img_infos = random.sample(self.img_infos, max_samples)
            # self.img_infos = self.img_infos[:max_samples]

        # (long_edge, short_edge) or [(long1, short1), (long2, short2), ...]
        self.img_scales = img_scale if isinstance(img_scale,
                                                  list) else [img_scale]
        assert mmcv.is_list_of(self.img_scales, tuple)
        # normalization configs
        self.img_norm_cfg = img_norm_cfg

        # multi-scale mode (only applicable for multi-scale training)
        self.multiscale_mode = multiscale_mode
        assert multiscale_mode in ['value', 'range']

        # max proposals per image
        self.num_max_proposals = num_max_proposals
        # flip ratio
        self.flip_ratio = flip_ratio
        assert flip_ratio >= 0 and flip_ratio <= 1
        # padding border to ensure the image size can be divided by
        # size_divisor (used for FPN)
        self.size_divisor = size_divisor

        # with mask or not (reserved field, takes no effect)
        self.with_mask = with_mask
        # some datasets provide bbox annotations as ignore/crowd/difficult,
        # if `with_crowd` is True, then these info is returned.
        self.with_crowd = with_crowd
        # with label is False for RPN
        self.with_label = with_label
        # with semantic segmentation (stuff) annotation or not
        self.with_seg = with_semantic_seg
        # prefix of semantic segmentation map path
        self.seg_prefix = seg_prefix
        # rescale factor for segmentation maps
        self.seg_scale_factor = seg_scale_factor
        # in test mode or not
        self.test_mode = test_mode
        # For dataset with SMPL shape parameters
        self.with_shape = with_shape
        # For dataset with joints rotation matrix in SMPL model
        self.with_kpts3d = with_kpts3d
        # For dataset with 2D pose
        self.with_kpts2d = with_kpts2d
        # For dataset with camera parameters in SMPL model
        self.with_trans = with_trans
        # For pose in axis angle of the joints
        self.with_pose = with_pose

        # Densepose annotations
        self.with_dp = with_dp

        # noise factor for color jittering
        self.noise_factor = noise_factor

        # Whether to adjust bbox to square manually..
        self.square_bbox = square_bbox

        # Rotation facotr
        self.rot_factor = rot_factor

        # Mosh dataset for generator
        self.mosh_path = None  # mosh_path

        # set group flag for the sampler
        if not self.test_mode and False:
            self._set_group_flag()
        # transforms
        self.img_transform = ImageTransform(
            size_divisor=self.size_divisor, **self.img_norm_cfg)
        self.bbox_transform = BboxTransform()
        self.mask_transform = MaskTransform()
        self.seg_transform = SegMapTransform(self.size_divisor)
        self.numpy2tensor = Numpy2Tensor()

        # if use extra augmentation
        if extra_aug is not None:
            self.extra_aug = ExtraAugmentation(**extra_aug)
        else:
            self.extra_aug = None

        # image rescale if keep ratio
        self.resize_keep_ratio = resize_keep_ratio

        self.sample_weight = sample_weight

        if sample_by_persons:
            persons_cnt = np.zeros(len(self.img_infos))
            for i in range(len(self.img_infos)):
                persons_cnt[i] = self.get_ann_info(i)['kpts2d'].shape[0]
            self.density = sample_weight * persons_cnt / persons_cnt.sum()
        else:
            self.density = sample_weight * np.ones(len(self.img_infos)) / len(self.img_infos)

        self.ignore_3d = ignore_3d
        self.ignore_smpl = ignore_smpl
        self.with_nr = with_nr
        self.use_poly = use_poly

        if self.mosh_path:
            mosh = np.load(mosh_path)
            self.mosh_shape = mosh['shape'].copy()
            self.mosh_pose = mosh['pose'].copy()
            self.mosh_sample_list = range(self.mosh_shape.shape[0])

    def __len__(self):
        return len(self.img_infos)

    def load_annotations(self, ann_file):
        """
        filename:
        height: 1000
        width: commonly 1002 in h36m
        :param ann_file:
        :return:
        """
        with open(ann_file, 'rb') as f:
            raw_infos = pickle.load(f)
        return raw_infos

    def load_proposals(self, proposal_file):
        return mmcv.load(proposal_file)

    def get_ann_info(self, idx):
        """
        :param idx:
        :return:A dict of the following iterms:
            bboxes: [x1, y1, x2, y2]
            labels: number
            kpts3d: (24, 4)
            kpts2d: (24, 3)
            pose: (72,)
            shape: (10,)
            cam: (3,) (The trans in SMPL model)
        """
        # Visualization needed
        raw_info = deepcopy(self.img_infos[idx])
        bbox = raw_info['bbox']
        if self.square_bbox:
            center = np.array([int((bbox[0] + bbox[2]) / 2), int((bbox[1] + bbox[3]) / 2)])
            bbox_size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
            half_size = int(math.floor(bbox_size / 2))
            square_bbox = np.array(
                [center[0] - half_size, center[1] - half_size, center[0] + half_size, center[1] + half_size])
        else:
            square_bbox = np.array(bbox)
        #  Here comes a problem, what if the bbox overflows on the corner?
        # They will be args passed to `model.train_forward` if we overwrite `CustomDataset`.
        return {'bboxes': square_bbox.reshape(1, -1).astype(np.float32),  # (1,4)
                # There will be only one person here.
                'labels': np.array([1]),
                'kpts3d': raw_info['S'][np.newaxis].astype(np.float32),  # (1, 24,4) extra chanel for visibility
                'kpts2d': raw_info['part'][np.newaxis].astype(np.float32),  # (1, 24,3) extra chanel for visibility
                'pose': raw_info['pose'].reshape(-1, 3)[np.newaxis].astype(np.float32),  # (1, 24, 3)
                'shape': raw_info['shape'][np.newaxis].astype(np.float32),  # (1,10)
                'trans': raw_info['trans'][np.newaxis].astype(np.float32),  # (1, 3)
                'has_smpl': np.array([1])
                }  # I think `1` represents the first class

    def _filter_imgs(self, min_size=32):
        """Filter images too small."""
        valid_inds = []
        for i, img_info in enumerate(self.img_infos):
            if min(img_info['width'], img_info['height']) >= min_size:
                valid_inds.append(i)
        return valid_inds

    def _set_group_flag(self):
        """Set flag according to image aspect ratio.

        Images with aspect ratio greater than 1 will be set as group 1,
        otherwise group 0.
        """
        self.flag = np.zeros(len(self), dtype=np.uint8)
        for i in range(len(self)):
            img_info = self.img_infos[i]
            if img_info['width'] / img_info['height'] > 1:
                self.flag[i] = 1

    def _rand_another(self, idx):
        pool = np.where(self.flag == self.flag[idx])[0]
        return np.random.choice(pool)

    def __getitem__(self, idx):
        # if self.test_mode:
        #     return self.prepare_test_img(idx)
        while True:
            data = self.prepare_train_img(idx)
            if data is None:
                idx = self._rand_another(idx)
                continue
            return data

    @staticmethod
    def val(runner, dataloader, **kwargs):
        from IPython import embed
        embed()
        pass

    @staticmethod
    def annToRLE(ann):
        h, w = ann['height'], ann['width']
        segm = ann['segmentation']
        if type(segm) == list:
            # polygon -- a single object might consist of multiple parts
            # we merge all parts into one mask rle code
            rles = mask_util.frPyObjects(segm, h, w)
            rle = mask_util.merge(rles)
        elif type(segm['counts']) == list:
            # uncompressed RLE
            rle = mask_util.frPyObjects(segm, h, w)
        else:
            # rle
            rle = ann['segmentation']
        return rle

    def prepare_train_img(self, idx):
        img_info = deepcopy(self.img_infos[idx])
        # load image
        img = mmcv.imread(osp.join(self.img_prefix, img_info['filename']))
        img_info['height'], img_info['width'] = img.shape[:2]
        raw_shape = img.shape
        # Color jittering:
        pn = np.random.uniform(1 - self.noise_factor, 1 + self.noise_factor, 3)
        img[:, :, 0] = np.minimum(255.0, np.maximum(0.0, img[:, :, 0] * pn[0]))
        img[:, :, 1] = np.minimum(255.0, np.maximum(0.0, img[:, :, 1] * pn[1]))
        img[:, :, 2] = np.minimum(255.0, np.maximum(0.0, img[:, :, 2] * pn[2]))

        # load proposals if necessary
        if self.proposals is not None:
            proposals = self.proposals[idx][:self.num_max_proposals]
            # no proposals are just ignored, but they can be used for
            # training in concept.
            if len(proposals) == 0:
                return None
            if not (proposals.shape[1] == 4 or proposals.shape[1] == 5):
                raise AssertionError(
                    'proposals should have shapes (n, 4) or (n, 5), '
                    'but found {}'.format(proposals.shape))
            if proposals.shape[1] == 5:
                scores = proposals[:, 4, None]
                proposals = proposals[:, :4]
            else:
                scores = None

        ann = self.get_ann_info(idx)
        if self.ignore_3d:
            ann['kpts3d'] = np.zeros_like(ann['kpts3d'])

        if self.ignore_smpl:
            ann['has_smpl'] = np.zeros_like(ann['has_smpl'])

        gt_bboxes = ann['bboxes']
        gt_labels = ann['labels']
        if self.with_crowd:
            gt_bboxes_ignore = ann['bboxes_ignore']

        # skip the image if there is no valid gt bbox
        if len(gt_bboxes) == 0:
            return None

        # extra augmentation
        if self.extra_aug is not None:
            img, gt_bboxes, gt_labels = self.extra_aug(img, gt_bboxes,
                                                       gt_labels)

        # apply transforms
        flip = True if np.random.rand() < self.flip_ratio else False
        # randomly sample a scale
        img_scale = random_scale(self.img_scales, self.multiscale_mode)

        img, img_shape, pad_shape, scale_factor = self.img_transform(img, img_scale, flip,
                                                                     keep_ratio=self.resize_keep_ratio)

        # Force padding for the issue of multi-GPU training
        padded_img = np.zeros((img.shape[0], img_scale[1], img_scale[0]), dtype=img.dtype)
        padded_img[:, :img.shape[-2], :img.shape[-1]] = img
        img = padded_img

        if self.with_nr:
            padded_scene = np.zeros((img.shape[-2], img.shape[-1]), dtype=np.uint8)

        if self.with_nr:
            has_mask = True
            # if 'dp_mask_path' in ann:
            #     raw_mask = cv2.imread(osp.join(self.img_prefix, ann['dp_mask_path']), cv2.IMREAD_UNCHANGED)
            # elif 'segmentation' in ann and self.use_poly:
            if 'segmentation' in ann and self.use_poly:
                assert 'COCO' in ann['filename'], "Only support coco segmentation now"
                raw_mask = np.zeros((ann['height'], ann['width']), dtype=np.uint8)
                for i, seg in enumerate(ann['segmentation']):
                    ori_mask = mask_util.decode(
                        self.annToRLE({'width': ann['width'], 'height': ann['height'], 'segmentation': seg}))
                    raw_mask[ori_mask > 0] = i + 1
            else:
                has_mask = False
            padded_scene = np.zeros((img.shape[-2], img.shape[-1]), dtype=np.uint8)
            if has_mask:
                target_shape = int(np.round(raw_shape[1] * scale_factor)), int(np.round(raw_shape[0] * scale_factor))
                resized_mask = cv2.resize(raw_mask, target_shape, interpolation=cv2.INTER_NEAREST)
                if flip:
                    resized_mask = np.flip(resized_mask, axis=1)
                padded_scene[:resized_mask.shape[-2], :resized_mask.shape[-1]] = resized_mask

        if self.with_seg:
            gt_seg = mmcv.imread(
                osp.join(self.seg_prefix, img_info['file_name'].replace(
                    'jpg', 'png')),
                flag='unchanged')
            gt_seg = self.seg_transform(gt_seg.squeeze(), img_scale, flip)
            gt_seg = mmcv.imrescale(
                gt_seg, self.seg_scale_factor, interpolation='nearest')
            gt_seg = gt_seg[None, ...]
        if self.proposals is not None:
            proposals = self.bbox_transform(proposals, img_shape, scale_factor,
                                            flip)
            proposals = np.hstack(
                [proposals, scores]) if scores is not None else proposals
        gt_bboxes = self.bbox_transform(gt_bboxes, img_shape, scale_factor,
                                        flip)
        if self.with_crowd:
            gt_bboxes_ignore = self.bbox_transform(gt_bboxes_ignore, img_shape,
                                                   scale_factor, flip)
        if self.with_mask:
            gt_masks = self.mask_transform(ann['masks'], pad_shape,
                                           scale_factor, flip)

        if self.with_shape:
            gt_shapes = ann['shape']
        if self.with_kpts2d:
            gt_kpts2d = ann['kpts2d']
            # Rescale the 2D keypoints now.
            s_kpts2d = np.zeros_like(gt_kpts2d)
            s_kpts2d[..., -1] = gt_kpts2d[..., -1]
            s_kpts2d[..., :-1] = gt_kpts2d[..., :-1] * scale_factor
            gt_kpts2d = s_kpts2d

            if flip:
                for i, kp in enumerate(gt_kpts2d):
                    gt_kpts2d[i] = flip_kp(kp, img_shape[1])  # img is (C, H, W)
                    # NOTE: I use the img_shape to avoid the influence of padding.

        if self.with_kpts3d:
            gt_kpts3d = ann['kpts3d']
            if flip:
                for i, kp in enumerate(gt_kpts3d):
                    gt_kpts3d[i] = flip_kp(kp, 0)  # Not the image width as the pose is centered by hip.
        if self.with_trans:
            gt_trans = ann['trans']
        if self.with_pose:
            gt_poses = ann['pose']
            if flip:
                for i, ps in enumerate(gt_poses):
                    gt_poses[i] = flip_pose(ps.reshape(-1)).reshape(-1, 3)

        if self.with_dp:
            dp_num_pts = ann['dp_num_pts']
            dp_x = ann['dp_x']
            dp_y = ann['dp_y']
            dp_U = ann['dp_U']
            dp_V = ann['dp_V']
            dp_I = ann['dp_I']
            dp_x = img_shape[1] - dp_x

        if not self.rot_factor == 0 and np.random.uniform() > 0.6:
            rot = min(2 * self.rot_factor,
                      max(-2 * self.rot_factor, np.random.randn() * self.rot_factor))
            rot_rad = -rot * np.pi / 180
            sn, cs = np.sin(rot_rad), np.cos(rot_rad)
            rot_mat = np.eye(3)
            rot_mat[0, :2] = [cs, -sn]
            rot_mat[1, :2] = [sn, cs]
            sn, cs = np.sin(rot_rad), np.cos(rot_rad)
            rot_mat[0, :2] = [cs, -sn]
            rot_mat[1, :2] = [sn, cs]

            # As we are rotating around the center.
            t_mat = np.eye(3)
            res = img.shape[1:]
            t_mat[0, 2] = -res[1] / 2
            t_mat[1, 2] = -res[0] / 2
            t_inv = t_mat.copy()
            t_inv[:2, 2] *= -1
            rot_mat2d = t_inv @ rot_mat @ t_mat
            # Things to change:
            # kpts2d and 3d
            # img and bbox

            # On GraphCMR, they only rotate on the first 3 entries. So I think it apply to ours.
            for i in range(gt_poses.shape[0]):
                gt_poses[i, 0] = rot_aa(gt_poses[i, 0], rot)
                gt_kpts3d[i, :, :-1] = (rot_mat @ gt_kpts3d[i, :, :-1].T).T
                gt_kpts2d[i, :, :-1] = (rot_mat2d @ gt_kpts2d[i, :, ].T).T[..., :-1]
                if sum(gt_kpts2d[i, ..., -1]) > 0:
                    x_min, y_min, _ = gt_kpts2d[i, gt_kpts2d[i, ..., -1] > 0].min(0)
                    x_max, y_max, _ = gt_kpts2d[i, gt_kpts2d[i, ..., -1] > 0].max(0)
                else:
                    # If there is no valida ktps, we will use gt_bboxes instead.
                    x_min, y_min = rot2DPts(gt_bboxes[i][0], gt_bboxes[i][1], rot_mat2d)
                    x_max, y_max = rot2DPts(gt_bboxes[i][2], gt_bboxes[i][3], rot_mat2d)
                bbox_w, bbox_h = x_max - x_min, y_max - y_min
                bbox_center = (x_min + x_max) / 2, (y_min + y_max) / 2
                _, im_w, im_h = img.shape
                x_min, y_min = max(0, bbox_center[0] - bbox_w * 0.55), max(0, bbox_center[1] - bbox_h * 0.55)
                x_max, y_max = min(im_w, bbox_center[0] + bbox_w * 0.55), min(im_h, bbox_center[1] + bbox_h * 0.55)
                bbox = np.array([x_min, y_min, x_max, y_max], dtype=np.float32)
                if self.square_bbox:
                    center = np.array([int((bbox[0] + bbox[2]) / 2), int((bbox[1] + bbox[3]) / 2)])
                    bbox_size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
                    half_size = int(math.floor(bbox_size / 2))
                    square_bbox = np.array(
                        [center[0] - half_size, center[1] - half_size, center[0] + half_size, center[1] + half_size])
                    bbox = square_bbox.astype(np.float32).copy()
                gt_bboxes[i] = bbox
                if self.with_dp:
                    dp_x, dp_y = rot2DPts(dp_x, dp_y, rot_mat2d)

            img_raw = cv2.cvtColor((denormalize(img.transpose([1, 2, 0])) * 255).astype(np.uint8).copy(),
                                   cv2.COLOR_BGR2RGB)
            img_rotated = scipy.misc.imrotate(img_raw, rot)

            # For debug

            img = mmcv.imnormalize(img_rotated, self.img_transform.mean, self.img_transform.std,
                                   self.img_transform.to_rgb).transpose([2, 0, 1])
        ori_shape = (img_info['height'], img_info['width'], 3)
        img_meta = dict(
            ori_shape=ori_shape,
            img_shape=img_shape,
            pad_shape=pad_shape,
            scale_factor=scale_factor,
            flip=flip,
            idx=idx,
            file_name=img_info['filename']
        )
        data = dict(
            img=DC(to_tensor(img), stack=True),
            img_meta=DC(img_meta, cpu_only=True),
            gt_bboxes=DC(to_tensor(gt_bboxes)))
        if self.proposals is not None:
            data['proposals'] = DC(to_tensor(proposals))
        if self.with_label:
            data['gt_labels'] = DC(to_tensor(gt_labels))
        if self.with_crowd:
            data['gt_bboxes_ignore'] = DC(to_tensor(gt_bboxes_ignore))
        if self.with_mask:
            data['gt_masks'] = DC(gt_masks, cpu_only=True)
        if self.with_seg:
            data['gt_semantic_seg'] = DC(to_tensor(gt_seg), stack=True)
        if self.with_pose:
            data['gt_poses'] = DC(to_tensor(gt_poses))
        if self.with_kpts2d:
            # Filter out un-used kpts.
            kpts_filter = gt_kpts2d[..., -1].sum(-1) < 8
            if gt_kpts2d.shape[0] == 1:
                gt_kpts2d[kpts_filter] = 0
            data['gt_kpts2d'] = DC(to_tensor(gt_kpts2d))
        if self.with_kpts3d:
            data['gt_kpts3d'] = DC(to_tensor(gt_kpts3d))
        if self.with_shape:
            data['gt_shapes'] = DC(to_tensor(gt_shapes))
        if self.with_trans:
            data['gt_trans'] = DC(to_tensor(gt_trans))
        if self.with_dp:
            data['dp_x'] = DC(to_tensor(dp_x))
            data['dp_y'] = DC(to_tensor(dp_y))
            data['dp_U'] = DC(to_tensor(dp_U))
            data['dp_V'] = DC(to_tensor(dp_V))
            data['dp_I'] = DC(to_tensor(dp_I))
            data['dp_num_pts'] = DC(to_tensor(dp_num_pts))
        # if self.with_shape = DC(to_tensor(gt_shape))
        if self.with_nr:
            data['scene'] = DC(to_tensor(padded_scene))
        data['has_smpl'] = DC(to_tensor(ann['has_smpl']))
        if self.mosh_path:
            sampled_idxs = np.array(random.sample(self.mosh_sample_list, 36))
            mosh_pose = torch.tensor(deepcopy(self.mosh_pose[sampled_idxs].astype(np.float32)))
            mosh_shape = torch.tensor(deepcopy(self.mosh_shape[sampled_idxs].astype(np.float32)))
            mosh_bs = mosh_shape.shape[0]
            mosh_pose_shape = torch.cat([batch_rodrigues(mosh_pose.view(-1, 3)).view(mosh_bs, -1), mosh_shape], dim=1)
            data['mosh'] = DC(mosh_pose_shape)
        return data

    def prepare_test_img(self, idx):
        """Prepare an image for testing (multi-scale and flipping)"""
        img_info = self.img_infos[idx]
        img = mmcv.imread(osp.join(self.img_prefix, img_info['filename']))
        if self.proposals is not None:
            proposal = self.proposals[idx][:self.num_max_proposals]
            if not (proposal.shape[1] == 4 or proposal.shape[1] == 5):
                raise AssertionError(
                    'proposals should have shapes (n, 4) or (n, 5), '
                    'but found {}'.format(proposal.shape))
        else:
            proposal = None

        def prepare_single(img, scale, flip, proposal=None):
            _img, img_shape, pad_shape, scale_factor = self.img_transform(
                img, scale, flip, keep_ratio=self.resize_keep_ratio)
            _img = to_tensor(_img)
            _img_meta = dict(
                ori_shape=(img_info['height'], img_info['width'], 3),
                img_shape=img_shape,
                pad_shape=pad_shape,
                scale_factor=scale_factor,
                flip=flip)
            if proposal is not None:
                if proposal.shape[1] == 5:
                    score = proposal[:, 4, None]
                    proposal = proposal[:, :4]
                else:
                    score = None
                _proposal = self.bbox_transform(proposal, img_shape,
                                                scale_factor, flip)
                _proposal = np.hstack(
                    [_proposal, score]) if score is not None else _proposal
                _proposal = to_tensor(_proposal)
            else:
                _proposal = None
            return _img, _img_meta, _proposal

        imgs = []
        img_metas = []
        proposals = []
        for scale in self.img_scales:
            _img, _img_meta, _proposal = prepare_single(
                img, scale, False, proposal)
            imgs.append(_img)
            img_metas.append(DC(_img_meta, cpu_only=True))
            proposals.append(_proposal)
            if self.flip_ratio > 0:
                _img, _img_meta, _proposal = prepare_single(
                    img, scale, True, proposal)
                imgs.append(_img)
                img_metas.append(DC(_img_meta, cpu_only=True))
                proposals.append(_proposal)
        data = dict(img=imgs, img_meta=img_metas)
        if self.proposals is not None:
            data['proposals'] = proposals
        return data

class MyH36MDataset(Dataset):
    """Custom dataset for detection.

    Annotation format:
    [
        {
            'filename': 'a.jpg',
            'width': 1280,
            'height': 720,
            'ann': {
                'bboxes': <np.ndarray> (n, 4),
                'labels': <np.ndarray> (n, ),
                'bboxes_ignore': <np.ndarray> (k, 4),
                'labels_ignore': <np.ndarray> (k, 4) (optional field)
            }
        },
        ...
    ]

    The `ann` field is optional for testing.
    """

    # CLASSES = None
    CLASSES = ('Human',)

    def __init__(self,
                 ann_file,
                 img_prefix,
                 img_scale,
                 img_norm_cfg,
                 multiscale_mode='value',
                 size_divisor=None,
                 proposal_file=None,
                 num_max_proposals=1000,
                 flip_ratio=0,
                 with_mask=True,
                 with_crowd=True,
                 with_label=True,
                 with_semantic_seg=False,
                 seg_prefix=None,
                 seg_scale_factor=1,
                 extra_aug=None,
                 resize_keep_ratio=True,
                 test_mode=False,
                 with_shape=True,
                 with_kpts3d=True,
                 with_kpts2d=True,
                 with_pose=True,
                 with_trans=True,
                 max_samples=-1,  # Commonly used in validating
                 noise_factor=0.4,
                 square_bbox=True,
                 rot_factor=0,
                 sample_weight=1,
                 with_dp=False,
                 mosh_path=None,
                 ignore_3d=False,
                 ignore_smpl=False,
                 with_nr=False,
                 sample_by_persons=False,
                 use_poly=False,
                 **kwargs,
                 ):
        # prefix of images path
        self.img_prefix = img_prefix

        # load annotations (and proposals)
        self.img_infos = self.load_annotations(ann_file)
        self.max_samples = max_samples

        if proposal_file is not None:
            self.proposals = self.load_proposals(proposal_file)
        else:
            self.proposals = None
        # filter images with no annotation during training
        if not test_mode and False:
            valid_inds = self._filter_imgs()
            self.img_infos = [self.img_infos[i] for i in valid_inds]
            if self.proposals is not None:
                self.proposals = [self.proposals[i] for i in valid_inds]

        # Select a subset for quick validation
        if self.max_samples > 0:
            self.img_infos = random.sample(self.img_infos, max_samples)
            # self.img_infos = self.img_infos[:max_samples]

        # (long_edge, short_edge) or [(long1, short1), (long2, short2), ...]
        self.img_scales = img_scale if isinstance(img_scale,
                                                  list) else [img_scale]
        assert mmcv.is_list_of(self.img_scales, tuple)
        # normalization configs
        self.img_norm_cfg = img_norm_cfg

        # multi-scale mode (only applicable for multi-scale training)
        self.multiscale_mode = multiscale_mode
        assert multiscale_mode in ['value', 'range']

        # max proposals per image
        self.num_max_proposals = num_max_proposals
        # flip ratio
        self.flip_ratio = flip_ratio
        assert flip_ratio >= 0 and flip_ratio <= 1
        # padding border to ensure the image size can be divided by
        # size_divisor (used for FPN)
        self.size_divisor = size_divisor

        # with mask or not (reserved field, takes no effect)
        self.with_mask = with_mask
        # some datasets provide bbox annotations as ignore/crowd/difficult,
        # if `with_crowd` is True, then these info is returned.
        self.with_crowd = with_crowd
        # with label is False for RPN
        self.with_label = with_label
        # with semantic segmentation (stuff) annotation or not
        self.with_seg = with_semantic_seg
        # prefix of semantic segmentation map path
        self.seg_prefix = seg_prefix
        # rescale factor for segmentation maps
        self.seg_scale_factor = seg_scale_factor
        # in test mode or not
        self.test_mode = test_mode
        # For dataset with SMPL shape parameters
        self.with_shape = with_shape
        # For dataset with joints rotation matrix in SMPL model
        self.with_kpts3d = with_kpts3d
        # For dataset with 2D pose
        self.with_kpts2d = with_kpts2d
        # For dataset with camera parameters in SMPL model
        self.with_trans = with_trans
        # For pose in axis angle of the joints
        self.with_pose = with_pose

        # Densepose annotations
        self.with_dp = with_dp

        # noise factor for color jittering
        self.noise_factor = noise_factor

        # Whether to adjust bbox to square manually..
        self.square_bbox = square_bbox

        # Rotation facotr
        self.rot_factor = rot_factor

        # Mosh dataset for generator
        self.mosh_path = None  # mosh_path

        # set group flag for the sampler
        if not self.test_mode and False:
            self._set_group_flag()
        # transforms
        self.img_transform = ImageTransform(
            size_divisor=self.size_divisor, **self.img_norm_cfg)
        self.bbox_transform = BboxTransform()
        self.mask_transform = MaskTransform()
        self.seg_transform = SegMapTransform(self.size_divisor)
        self.numpy2tensor = Numpy2Tensor()

        # if use extra augmentation
        if extra_aug is not None:
            self.extra_aug = ExtraAugmentation(**extra_aug)
        else:
            self.extra_aug = None

        # image rescale if keep ratio
        self.resize_keep_ratio = resize_keep_ratio

        self.sample_weight = sample_weight

        if sample_by_persons:
            persons_cnt = np.zeros(len(self.img_infos))
            for i in range(len(self.img_infos)):
                persons_cnt[i] = self.get_ann_info(i)['kpts2d'].shape[0]
            self.density = sample_weight * persons_cnt / persons_cnt.sum()
        else:
            self.density = sample_weight * np.ones(len(self.img_infos)) / len(self.img_infos)

        self.ignore_3d = ignore_3d
        self.ignore_smpl = ignore_smpl
        self.with_nr = with_nr
        self.use_poly = use_poly

        if self.mosh_path:
            mosh = np.load(mosh_path)
            self.mosh_shape = mosh['shape'].copy()
            self.mosh_pose = mosh['pose'].copy()
            self.mosh_sample_list = range(self.mosh_shape.shape[0])

    def __len__(self):
        return len(self.img_infos)

    def load_annotations(self, ann_file):
        """
        filename:
        height: 1000
        width: commonly 1002 in h36m
        :param ann_file:
        :return:
        """
        with open(ann_file, 'rb') as f:
            raw_infos = pickle.load(f)
        return raw_infos

    def load_proposals(self, proposal_file):
        return mmcv.load(proposal_file)

    def get_ann_info(self, idx):
        """
        :param idx:
        :return:A dict of the following iterms:
            bboxes: [x1, y1, x2, y2]
            labels: number
            kpts3d: (24, 4)
            kpts2d: (24, 3)
            pose: (72,)
            shape: (10,)
            cam: (3,) (The trans in SMPL model)
        """
        # Visualization needed
        raw_info = deepcopy(self.img_infos[idx])
        bbox = raw_info['bbox']
        if self.square_bbox:
            center = np.array([int((bbox[0] + bbox[2]) / 2), int((bbox[1] + bbox[3]) / 2)])
            bbox_size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
            half_size = int(math.floor(bbox_size / 2))
            square_bbox = np.array(
                [center[0] - half_size, center[1] - half_size, center[0] + half_size, center[1] + half_size])
        else:
            square_bbox = np.array(bbox)
        #  Here comes a problem, what if the bbox overflows on the corner?
        # They will be args passed to `model.train_forward` if we overwrite `CustomDataset`.
        return {'bboxes': square_bbox.reshape(1, -1).astype(np.float32),  # (1,4)
                # There will be only one person here.
                'labels': np.array([1]),
                'kpts3d': raw_info['S'][np.newaxis].astype(np.float32),  # (1, 24,4) extra chanel for visibility
                'kpts2d': raw_info['part'][np.newaxis].astype(np.float32),  # (1, 24,3) extra chanel for visibility
                'pose': raw_info['pose'].reshape(-1, 3)[np.newaxis].astype(np.float32),  # (1, 24, 3)
                'shape': raw_info['shape'][np.newaxis].astype(np.float32),  # (1,10)
                'trans': raw_info['trans'][np.newaxis].astype(np.float32),  # (1, 3)
                'has_smpl': np.array([1])
                }  # I think `1` represents the first class

    def _filter_imgs(self, min_size=32):
        """Filter images too small."""
        valid_inds = []
        for i, img_info in enumerate(self.img_infos):
            if min(img_info['width'], img_info['height']) >= min_size:
                valid_inds.append(i)
        return valid_inds

    def _set_group_flag(self):
        """Set flag according to image aspect ratio.

        Images with aspect ratio greater than 1 will be set as group 1,
        otherwise group 0.
        """
        self.flag = np.zeros(len(self), dtype=np.uint8)
        for i in range(len(self)):
            img_info = self.img_infos[i]
            if img_info['width'] / img_info['height'] > 1:
                self.flag[i] = 1

    def _rand_another(self, idx):
        pool = np.where(self.flag == self.flag[idx])[0]
        return np.random.choice(pool)

    def __getitem__(self, idx):
        # if self.test_mode:
        #     return self.prepare_test_img(idx)
        while True:
            data = self.prepare_train_img(idx)
            if data is None:
                idx = self._rand_another(idx)
                continue
            return data

    @staticmethod
    def val(runner, dataloader, **kwargs):
        from IPython import embed
        embed()
        pass

    @staticmethod
    def annToRLE(ann):
        h, w = ann['height'], ann['width']
        segm = ann['segmentation']
        if type(segm) == list:
            # polygon -- a single object might consist of multiple parts
            # we merge all parts into one mask rle code
            rles = mask_util.frPyObjects(segm, h, w)
            rle = mask_util.merge(rles)
        elif type(segm['counts']) == list:
            # uncompressed RLE
            rle = mask_util.frPyObjects(segm, h, w)
        else:
            # rle
            rle = ann['segmentation']
        return rle

    def prepare_train_img(self, idx):
        img_info = deepcopy(self.img_infos[idx])
        # load image
        img = mmcv.imread(osp.join(self.img_prefix, img_info['filename']))
        img_info['height'], img_info['width'] = img.shape[:2]
        raw_shape = img.shape
        # Color jittering:
        pn = np.random.uniform(1 - self.noise_factor, 1 + self.noise_factor, 3)
        img[:, :, 0] = np.minimum(255.0, np.maximum(0.0, img[:, :, 0] * pn[0]))
        img[:, :, 1] = np.minimum(255.0, np.maximum(0.0, img[:, :, 1] * pn[1]))
        img[:, :, 2] = np.minimum(255.0, np.maximum(0.0, img[:, :, 2] * pn[2]))

        # load proposals if necessary
        if self.proposals is not None:
            proposals = self.proposals[idx][:self.num_max_proposals]
            # no proposals are just ignored, but they can be used for
            # training in concept.
            if len(proposals) == 0:
                return None
            if not (proposals.shape[1] == 4 or proposals.shape[1] == 5):
                raise AssertionError(
                    'proposals should have shapes (n, 4) or (n, 5), '
                    'but found {}'.format(proposals.shape))
            if proposals.shape[1] == 5:
                scores = proposals[:, 4, None]
                proposals = proposals[:, :4]
            else:
                scores = None

        ann = self.get_ann_info(idx)
        if self.ignore_3d:
            ann['kpts3d'] = np.zeros_like(ann['kpts3d'])

        if self.ignore_smpl:
            ann['has_smpl'] = np.zeros_like(ann['has_smpl'])

        gt_bboxes = ann['bboxes']
        gt_labels = ann['labels']
        if self.with_crowd:
            gt_bboxes_ignore = ann['bboxes_ignore']

        # skip the image if there is no valid gt bbox
        if len(gt_bboxes) == 0:
            return None

        # extra augmentation
        if self.extra_aug is not None:
            img, gt_bboxes, gt_labels = self.extra_aug(img, gt_bboxes,
                                                       gt_labels)

        # apply transforms
        flip = True if np.random.rand() < self.flip_ratio else False
        # randomly sample a scale
        img_scale = random_scale(self.img_scales, self.multiscale_mode)

        img, img_shape, pad_shape, scale_factor = self.img_transform(img, img_scale, flip,
                                                                     keep_ratio=self.resize_keep_ratio)

        # Force padding for the issue of multi-GPU training
        padded_img = np.zeros((img.shape[0], img_scale[1], img_scale[0]), dtype=img.dtype)
        padded_img[:, :img.shape[-2], :img.shape[-1]] = img
        img = padded_img

        if self.with_nr:
            padded_scene = np.zeros((img.shape[-2], img.shape[-1]), dtype=np.uint8)

        if self.with_nr:
            has_mask = True
            # if 'dp_mask_path' in ann:
            #     raw_mask = cv2.imread(osp.join(self.img_prefix, ann['dp_mask_path']), cv2.IMREAD_UNCHANGED)
            # elif 'segmentation' in ann and self.use_poly:
            if 'segmentation' in ann and self.use_poly:
                assert 'COCO' in ann['filename'], "Only support coco segmentation now"
                raw_mask = np.zeros((ann['height'], ann['width']), dtype=np.uint8)
                for i, seg in enumerate(ann['segmentation']):
                    ori_mask = mask_util.decode(
                        self.annToRLE({'width': ann['width'], 'height': ann['height'], 'segmentation': seg}))
                    raw_mask[ori_mask > 0] = i + 1
            else:
                has_mask = False
            padded_scene = np.zeros((img.shape[-2], img.shape[-1]), dtype=np.uint8)
            if has_mask:
                target_shape = int(np.round(raw_shape[1] * scale_factor)), int(np.round(raw_shape[0] * scale_factor))
                resized_mask = cv2.resize(raw_mask, target_shape, interpolation=cv2.INTER_NEAREST)
                if flip:
                    resized_mask = np.flip(resized_mask, axis=1)
                padded_scene[:resized_mask.shape[-2], :resized_mask.shape[-1]] = resized_mask

        if self.with_seg:
            gt_seg = mmcv.imread(
                osp.join(self.seg_prefix, img_info['file_name'].replace(
                    'jpg', 'png')),
                flag='unchanged')
            gt_seg = self.seg_transform(gt_seg.squeeze(), img_scale, flip)
            gt_seg = mmcv.imrescale(
                gt_seg, self.seg_scale_factor, interpolation='nearest')
            gt_seg = gt_seg[None, ...]
        if self.proposals is not None:
            proposals = self.bbox_transform(proposals, img_shape, scale_factor,
                                            flip)
            proposals = np.hstack(
                [proposals, scores]) if scores is not None else proposals
        gt_bboxes = self.bbox_transform(gt_bboxes, img_shape, scale_factor,
                                        flip)
        if self.with_crowd:
            gt_bboxes_ignore = self.bbox_transform(gt_bboxes_ignore, img_shape,
                                                   scale_factor, flip)
        if self.with_mask:
            gt_masks = self.mask_transform(ann['masks'], pad_shape,
                                           scale_factor, flip)

        if self.with_shape:
            gt_shapes = ann['shape']
        if self.with_kpts2d:
            gt_kpts2d = ann['kpts2d']
            # Rescale the 2D keypoints now.
            s_kpts2d = np.zeros_like(gt_kpts2d)
            s_kpts2d[..., -1] = gt_kpts2d[..., -1]
            s_kpts2d[..., :-1] = gt_kpts2d[..., :-1] * scale_factor
            gt_kpts2d = s_kpts2d

            if flip:
                for i, kp in enumerate(gt_kpts2d):
                    gt_kpts2d[i] = flip_kp(kp, img_shape[1])  # img is (C, H, W)
                    # NOTE: I use the img_shape to avoid the influence of padding.

        if self.with_kpts3d:
            gt_kpts3d = ann['kpts3d']
            if flip:
                for i, kp in enumerate(gt_kpts3d):
                    gt_kpts3d[i] = flip_kp(kp, 0)  # Not the image width as the pose is centered by hip.
        if self.with_trans:
            gt_trans = ann['trans']
        if self.with_pose:
            gt_poses = ann['pose']
            if flip:
                for i, ps in enumerate(gt_poses):
                    gt_poses[i] = flip_pose(ps.reshape(-1)).reshape(-1, 3)

        if self.with_dp:
            dp_num_pts = ann['dp_num_pts']
            dp_x = ann['dp_x']
            dp_y = ann['dp_y']
            dp_U = ann['dp_U']
            dp_V = ann['dp_V']
            dp_I = ann['dp_I']
            dp_x = img_shape[1] - dp_x

        if not self.rot_factor == 0 and np.random.uniform() > 0.6:
            rot = min(2 * self.rot_factor,
                      max(-2 * self.rot_factor, np.random.randn() * self.rot_factor))
            rot_rad = -rot * np.pi / 180
            sn, cs = np.sin(rot_rad), np.cos(rot_rad)
            rot_mat = np.eye(3)
            rot_mat[0, :2] = [cs, -sn]
            rot_mat[1, :2] = [sn, cs]
            sn, cs = np.sin(rot_rad), np.cos(rot_rad)
            rot_mat[0, :2] = [cs, -sn]
            rot_mat[1, :2] = [sn, cs]

            # As we are rotating around the center.
            t_mat = np.eye(3)
            res = img.shape[1:]
            t_mat[0, 2] = -res[1] / 2
            t_mat[1, 2] = -res[0] / 2
            t_inv = t_mat.copy()
            t_inv[:2, 2] *= -1
            rot_mat2d = t_inv @ rot_mat @ t_mat
            # Things to change:
            # kpts2d and 3d
            # img and bbox

            # On GraphCMR, they only rotate on the first 3 entries. So I think it apply to ours.
            for i in range(gt_poses.shape[0]):
                gt_poses[i, 0] = rot_aa(gt_poses[i, 0], rot)
                gt_kpts3d[i, :, :-1] = (rot_mat @ gt_kpts3d[i, :, :-1].T).T
                gt_kpts2d[i, :, :-1] = (rot_mat2d @ gt_kpts2d[i, :, ].T).T[..., :-1]
                if sum(gt_kpts2d[i, ..., -1]) > 0:
                    x_min, y_min, _ = gt_kpts2d[i, gt_kpts2d[i, ..., -1] > 0].min(0)
                    x_max, y_max, _ = gt_kpts2d[i, gt_kpts2d[i, ..., -1] > 0].max(0)
                else:
                    # If there is no valida ktps, we will use gt_bboxes instead.
                    x_min, y_min = rot2DPts(gt_bboxes[i][0], gt_bboxes[i][1], rot_mat2d)
                    x_max, y_max = rot2DPts(gt_bboxes[i][2], gt_bboxes[i][3], rot_mat2d)
                bbox_w, bbox_h = x_max - x_min, y_max - y_min
                bbox_center = (x_min + x_max) / 2, (y_min + y_max) / 2
                _, im_w, im_h = img.shape
                x_min, y_min = max(0, bbox_center[0] - bbox_w * 0.55), max(0, bbox_center[1] - bbox_h * 0.55)
                x_max, y_max = min(im_w, bbox_center[0] + bbox_w * 0.55), min(im_h, bbox_center[1] + bbox_h * 0.55)
                bbox = np.array([x_min, y_min, x_max, y_max], dtype=np.float32)
                if self.square_bbox:
                    center = np.array([int((bbox[0] + bbox[2]) / 2), int((bbox[1] + bbox[3]) / 2)])
                    bbox_size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
                    half_size = int(math.floor(bbox_size / 2))
                    square_bbox = np.array(
                        [center[0] - half_size, center[1] - half_size, center[0] + half_size, center[1] + half_size])
                    bbox = square_bbox.astype(np.float32).copy()
                gt_bboxes[i] = bbox
                if self.with_dp:
                    dp_x, dp_y = rot2DPts(dp_x, dp_y, rot_mat2d)

            img_raw = cv2.cvtColor((denormalize(img.transpose([1, 2, 0])) * 255).astype(np.uint8).copy(),
                                   cv2.COLOR_BGR2RGB)
            img_rotated = scipy.misc.imrotate(img_raw, rot)

            # For debug

            img = mmcv.imnormalize(img_rotated, self.img_transform.mean, self.img_transform.std,
                                   self.img_transform.to_rgb).transpose([2, 0, 1])
        ori_shape = (img_info['height'], img_info['width'], 3)
        img_meta = dict(
            ori_shape=ori_shape,
            img_shape=img_shape,
            pad_shape=pad_shape,
            scale_factor=scale_factor,
            flip=flip,
            idx=idx,
            file_name=img_info['filename']
        )
        data = dict(
            img=DC(to_tensor(img), stack=True),
            img_meta=DC(img_meta, cpu_only=True),
            gt_bboxes=DC(to_tensor(gt_bboxes)))
        if self.proposals is not None:
            data['proposals'] = DC(to_tensor(proposals))
        if self.with_label:
            data['gt_labels'] = DC(to_tensor(gt_labels))
        if self.with_crowd:
            data['gt_bboxes_ignore'] = DC(to_tensor(gt_bboxes_ignore))
        if self.with_mask:
            data['gt_masks'] = DC(gt_masks, cpu_only=True)
        if self.with_seg:
            data['gt_semantic_seg'] = DC(to_tensor(gt_seg), stack=True)
        if self.with_pose:
            data['gt_poses'] = DC(to_tensor(gt_poses))
        if self.with_kpts2d:
            # Filter out un-used kpts.
            kpts_filter = gt_kpts2d[..., -1].sum(-1) < 8
            if gt_kpts2d.shape[0] == 1:
                gt_kpts2d[kpts_filter] = 0
            data['gt_kpts2d'] = DC(to_tensor(gt_kpts2d))
        if self.with_kpts3d:
            data['gt_kpts3d'] = DC(to_tensor(gt_kpts3d))
        if self.with_shape:
            data['gt_shapes'] = DC(to_tensor(gt_shapes))
        if self.with_trans:
            data['gt_trans'] = DC(to_tensor(gt_trans))
        if self.with_dp:
            data['dp_x'] = DC(to_tensor(dp_x))
            data['dp_y'] = DC(to_tensor(dp_y))
            data['dp_U'] = DC(to_tensor(dp_U))
            data['dp_V'] = DC(to_tensor(dp_V))
            data['dp_I'] = DC(to_tensor(dp_I))
            data['dp_num_pts'] = DC(to_tensor(dp_num_pts))
        # if self.with_shape = DC(to_tensor(gt_shape))
        if self.with_nr:
            data['scene'] = DC(to_tensor(padded_scene))
        data['has_smpl'] = DC(to_tensor(ann['has_smpl']))
        if self.mosh_path:
            sampled_idxs = np.array(random.sample(self.mosh_sample_list, 36))
            mosh_pose = torch.tensor(deepcopy(self.mosh_pose[sampled_idxs].astype(np.float32)))
            mosh_shape = torch.tensor(deepcopy(self.mosh_shape[sampled_idxs].astype(np.float32)))
            mosh_bs = mosh_shape.shape[0]
            mosh_pose_shape = torch.cat([batch_rodrigues(mosh_pose.view(-1, 3)).view(mosh_bs, -1), mosh_shape], dim=1)
            data['mosh'] = DC(mosh_pose_shape)
        return data

    def prepare_test_img(self, idx):
        """Prepare an image for testing (multi-scale and flipping)"""
        img_info = self.img_infos[idx]
        img = mmcv.imread(osp.join(self.img_prefix, img_info['filename']))
        if self.proposals is not None:
            proposal = self.proposals[idx][:self.num_max_proposals]
            if not (proposal.shape[1] == 4 or proposal.shape[1] == 5):
                raise AssertionError(
                    'proposals should have shapes (n, 4) or (n, 5), '
                    'but found {}'.format(proposal.shape))
        else:
            proposal = None

        def prepare_single(img, scale, flip, proposal=None):
            _img, img_shape, pad_shape, scale_factor = self.img_transform(
                img, scale, flip, keep_ratio=self.resize_keep_ratio)
            _img = to_tensor(_img)
            _img_meta = dict(
                ori_shape=(img_info['height'], img_info['width'], 3),
                img_shape=img_shape,
                pad_shape=pad_shape,
                scale_factor=scale_factor,
                flip=flip)
            if proposal is not None:
                if proposal.shape[1] == 5:
                    score = proposal[:, 4, None]
                    proposal = proposal[:, :4]
                else:
                    score = None
                _proposal = self.bbox_transform(proposal, img_shape,
                                                scale_factor, flip)
                _proposal = np.hstack(
                    [_proposal, score]) if score is not None else _proposal
                _proposal = to_tensor(_proposal)
            else:
                _proposal = None
            return _img, _img_meta, _proposal

        imgs = []
        img_metas = []
        proposals = []
        for scale in self.img_scales:
            _img, _img_meta, _proposal = prepare_single(
                img, scale, False, proposal)
            imgs.append(_img)
            img_metas.append(DC(_img_meta, cpu_only=True))
            proposals.append(_proposal)
            if self.flip_ratio > 0:
                _img, _img_meta, _proposal = prepare_single(
                    img, scale, True, proposal)
                imgs.append(_img)
                img_metas.append(DC(_img_meta, cpu_only=True))
                proposals.append(_proposal)
        data = dict(img=imgs, img_meta=img_metas)
        if self.proposals is not None:
            data['proposals'] = proposals
        return data


import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

import time
import random
import copy
from copy import deepcopy
import threading

import gym3
from procgen import ProcgenGym3Env

train_num_levels = 100_000 #500 #1500
train_start_level = 0

for _ in range(10):
    bs = 144 * 10

    env = ProcgenGym3Env(num=bs, env_name="testgame", num_levels=train_num_levels, start_level=train_start_level,
                    color_theme=[5], color_theme_road=[5], background_noise_level=0)

    s = np.array([[.0,.0] for _ in range(bs)], dtype=np.float32)

    seq_len = 30

    for i in range(seq_len):
        env.act(s)
        rew, obs, first = env.observe()
        img = obs['rgb']
        info = env.get_info()

    env.close()


from gensim.models.fasttext import FastText as FT_gensim
import pickle
import numpy as np
import warnings
import pandas as pd
warnings.filterwarnings("ignore", category=FutureWarning)

class RunServerClassifier(object):
    '''

    '''

    def __init__(self):
        self.models = '../../models/'
        self.data = '../../labeldata/'
        self.metacyc = 'LabeledData.pkl'
        self.kegg = 'df_kegg.pkl'
        self.model = FT_gensim.load(self.models+'tierT12_5_3_5_300.model')
        self.classifier = self.models+'RMmodel.pkl'
        self.categories = ['Detoxification', 'Activation',
                           'Biosynthesis', 'Degradation', 'Energy', 'Glycan',
                           'Macromolecule']

    def check_format(self, sample):
        sample_list = sample.lower().split(',')
        try:
            ec = [i.count('ec') for i in sample_list]

            if len(sample_list) != sum(ec):
                printed_text = "Input not in the right format, you might have forgotten using ec or the semicolon between ec and the numbers"
                return 1, printed_text
        except (SyntaxError, ValueError):
            printed_text = "Input not in the right format, you might have forgotten using ec or the semicolon between ec and the numbers"
            return 1, printed_text
        try:
            # check if correct number and type of ec classes are being used
            ec_class_ind = [i.split(':')[1].split('.') for i in sample_list]
            ec_class = set([j for i in ec_class_ind for j in i])
            ec_class_isal = [j.isalpha() for i in ec_class for j in i]
            if True in ec_class_isal:
                printed_text = "Input not in the right format, your ec number has strings"
                return 1, printed_text
        except (IndexError, SyntaxError, ValueError):
            printed_text = "Input not in the right format, you might have forgotten using ec or the semicolon between ec and the numbers"
            return 1, printed_text
        try:
            # check if correct number and type of ec classes are being used
            ec_class_1 = set([int(i.split(':')[1].split('.')[0])
                             for i in sample_list])
            ec_class_2 = [i for i in ec_class_1 if i >= 7]
            if len(list(ec_class_1)) >= 7:
                printed_text = "Input not in the right format, you are not using EC classe 1-7"
                return 1, printed_text
            if len(ec_class_2) > 0:
                printed_text = "Input not in the right format, you are not using EC classe 1-7"
                return 1, printed_text

        except (IndexError, SyntaxError, ValueError):
            printed_text = "Input not in the right format, you might have forgotten using ec or the semicolon between ec and the numbers"
            return 1, printed_text

        try:
            # check number of digits in ec numbers
            ec_nums = set([len(i.split(':')[1].split('.'))
                          for i in sample_list])
            ec_nums_T = [True for i in ec_nums if i not in [2, 3, 4]]
            if len(ec_nums_T) > 0:
                printed_text = "Input not in the right format, check your ec numbers"
                return 1, printed_text
        except (IndexError, SyntaxError, ValueError):
            printed_text = "Input not in the right format, check your ec numbers"
            return 1, printed_text

        return 0, "Input in the right format"

    def get_pathway_vectors(self, sample_list):
        return np.mean(self.model.wv[sample_list], axis=0)

    def run_classifier(self, sample):
        sample_list = sample.lower().split(',')
        X_test = self.get_pathway_vectors(sample_list).reshape(1, -1)
        with open(self.classifier, 'rb') as file:
            pickle_model = pickle.load(file)
        prediction_prob = pickle_model.predict_proba(X_test)
        Y_pred = pickle_model.predict(X_test)
        classes = [i[0]
                   for i in list(zip(self.categories, Y_pred[0])) if i[1] == 1]
        classes_prob = [str(round(i[1], 4)) for i in list(
            zip(self.categories, prediction_prob[0])) if i[0] in classes]
        if len(classes) > 0:
            output_text1 = 'The predicted class/classes: ' + ', '.join(classes)
            output_text2 = 'The predicted class probability/probabilities: ' + \
                ', '.join(classes_prob)

        else:
            output_text1 = 'The predicted class/classes: No or ambiguous prediction'
            output_text2 = 'Probabilities associated with the list of classes:  No or ambiguous prediction'

        output_text3 = 'List of all class/classes: ' + \
            ', '.join(self.categories)
        probs = [str(round(i, 4)) for i in prediction_prob[0]]
        output_text4 = 'Probabilities associated with the list of all classes: ' + \
            ', '.join(probs)

        return output_text1, output_text2, output_text3, output_text4

    def similarity_dict_metacyc(self):
        with open(self.data+self.metacyc, 'rb') as f:
            data_df_multi = pickle.load(f)
        data_Metacyc_part = data_df_multi[['Map', 'Name', 'EC']]
        EC_list = data_Metacyc_part['EC'].to_list()
        EC_dict = {tup[0]: tup[1] for tup in list(
            zip(data_Metacyc_part['Map'], data_Metacyc_part['EC']))}
        EC_name = {tup[0]: tup[1] for tup in list(
            zip(data_Metacyc_part['Map'], data_Metacyc_part['Name']))}
        return EC_list, EC_dict, EC_name

    def similarity_dict_kegg(self):
        with open(self.data+self.kegg, 'rb') as f:
            data_df_multi = pickle.load(f)
        data_Metacyc_part = data_df_multi[['Map', 'Name', 'EC']]
        EC_list = data_Metacyc_part['EC'].to_list()
        EC_dict = {tup[0]: tup[1] for tup in list(
            zip(data_Metacyc_part['Map'], data_Metacyc_part['EC']))}
        EC_name = {tup[0]: tup[1] for tup in list(
            zip(data_Metacyc_part['Map'], data_Metacyc_part['Name']))}
        return EC_list, EC_dict, EC_name

    def get_ranking_single(self, valid, metacyc=True, kegg=True):

        if metacyc:
            data, data_dict, data_dict_name = self.similarity_dict_metacyc()
            df_simM = pd.DataFrame(
                columns=['Pathway_test', 'Pathway_train', 'Similarity'])
            for k, l in enumerate(data):
                sim = self.model.wv.n_similarity(valid, data[k])
                pathway_name = data_dict_name[list(data_dict.keys())[k]]
                perc = round(sim*100, 2)
                if sim > 0:
                    lst = [['UNK', list(data_dict.keys())[
                        k], pathway_name, sim, perc]]
                    df_simM = df_simM.append(pd.DataFrame(lst, columns=[
                                             'Pathway_test', 'Pathway_train', 'Pathway_name', 'Similarity', 'Percentage']))

        if kegg:
            data, data_dict, data_dict_name = self.similarity_dict_kegg()
            df_simK = pd.DataFrame(
                columns=['Pathway_test', 'Pathway_train', 'Similarity'])
            for k, l in enumerate(data):
                sim = self.model.wv.n_similarity(valid, data[k])
                pathway_name = data_dict_name[list(data_dict.keys())[k]]
                perc = round(sim*100, 2)
                if sim > 0:
                    lst = [['UNK', list(data_dict.keys())[
                        k], pathway_name, sim, perc]]
                    df_simK = df_simK.append(pd.DataFrame(lst, columns=[
                                             'Pathway_test', 'Pathway_train', 'Pathway_name', 'Similarity', 'Percentage']))

        return df_simM, df_simK

    def run_similarity(self, sample):
        sample_list = sample.lower().split(',')
        df_simM, df_simK = self.get_ranking_single(valid=sample_list)
        df_simM['Pathway Link'] = df_simM['Pathway_train'].apply(lambda x: '' +
                                                                 'href=https://biocyc.org/META/new-image?object='+str(x)+'')
        Metacyc = df_simM[['Pathway_train', 'Percentage', 'Pathway Link']].sort_values(
            by=['Percentage'], ascending=False).head()

        MetacycF = Metacyc.rename(
            {'Pathway_train': 'Pathway', 'Percentage': 'Percentage Similarity(%)'}, axis=1)

        df_simK['Pathway Link'] = df_simK['Pathway_train'].apply(lambda x: '' +
                                                                 'https://www.genome.jp/dbget-bin/www_bget?'+str(x)+'')
        Kegg = df_simK[['Pathway_train', 'Percentage', 'Pathway Link']].sort_values(
            by=['Percentage'], ascending=False).head()

        KeggF = Kegg.rename({'Pathway_train': 'Pathway',
                            'Percentage': 'Percentage Similarity(%)'}, axis=1)

        return MetacycF, KeggF

if __name__ == '__main__':
    R = RunServerClassifier()
    # print(R.check_format('ec:1.1.1.1.8'))
    print(R.run_classifier('ec:1.1.2.3'))