import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config.config import cfg
from util.human_models import smpl_x
from util.preprocessing import (
    load_img, process_bbox, augmentation_instance_sample
    ,process_human_model_output_batch_simplify,process_db_coord_batch_no_valid)
from util.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset


class SynBody(HumanDataset):
    def __init__(self, transform, data_split):
        super(SynBody, self).__init__(transform, data_split)
        self.img_dir = 'data/datasets/synbody'
        self.annot_path = 'data/preprocessed_npz/multihuman_data/synbody_v1.1_multi_new.npz'
        self.annot_path_cache = 'data/preprocessed_npz/cache/synbody_v1.1_cache_new_10.npz'
        self.use_cache = getattr(cfg, 'use_cache', False)
        self.img_shape = (720, 1280)  # (h, w)
        self.cam_param = {
            'focal': (540, 540),  # (fx, fy)
            'princpt': (640, 360)  # (cx, cy)
        }

        # check image shape
        img_path = osp.join(self.img_dir,
                            np.load(self.annot_path)['image_path'][0])

        img_shape = cv2.imread(img_path).shape[:2]
        assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(
            self.img_shape, img_shape)

        # load data or cache
        if self.use_cache and osp.isfile(self.annot_path_cache):
            print(
                f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}'
            )
            self.datalist = self.load_cache(self.annot_path_cache)
        else:
            if self.use_cache:
                print(
                    f'[{self.__class__.__name__}] Cache not found, generating cache...'
                )
            self.datalist = self.load_data(train_sample_interval=getattr(
                cfg, f'{self.__class__.__name__}_train_sample_interval', 15))
            if self.use_cache:
                self.save_cache(self.annot_path_cache, self.datalist)