import os
import json
from PIL import Image
import torch
from torch.utils.data import Dataset
from pycocotools.coco import COCO

class CoCoClassification(Dataset):
    def __init__(self, image_dir, phase, input_transform=None):
        self.image_dir = image_dir
        self.phase = phase
        self.ann_files = os.path.join(self.image_dir, 'PML_annotations', '{}_coco2014.json'.format(self.phase))
        self.load_anns()
        
        self.input_transform = input_transform
        
    def load_anns(self):
        if not os.path.exists(self.ann_files):
            if self.phase in ['train', 'val']:
                self.extract_anns()
            else:
                print("Run 'prepare_datasets.py' first!!")
                exit(-1)
        self.anns = []
        json_data = json.load(open(self.ann_files, "r"))
        self.anns += json_data
    
    def extract_anns(self):
        annFile = os.path.join(self.image_dir, "annotations/instances_{}2014.json".format(self.phase))
        coco = COCO(annFile)
    
        img_id = coco.getImgIds()
        cat_id = coco.getCatIds()
        img_id = list(sorted(img_id))
        cat_trans = {}
        cat_names = []
        for i in range(len(cat_id)):
            cat_trans[cat_id[i]] = i
            cat_name = coco.loadCats(cat_id[i])[0]['name']
            cat_names.append(cat_name)
        with open(os.path.join(self.image_dir, "categories.json"), 'w') as f:
            json.dump(cat_names, f)
    
        message = []
        for i in img_id:
            data = {}
            target = [0] * 80
            path = ""
            img_info = coco.loadImgs(i)[0]
            ann_ids = coco.getAnnIds(imgIds = i)
            anns = coco.loadAnns(ann_ids)
            if len(anns) == 0:
                continue
            else:
                for i in range(len(anns)):
                    cls = anns[i]['category_id']
                    cls = cat_trans[cls]
                    target[cls] = 1
                path = img_info['file_name']
            data['target'] = target
            data['img_path'] = os.path.join(os.path.join(self.image_dir, "{}2014/".format(self.phase)), path)
            message.append(data)
    
        with open(self.ann_files, 'w') as f:
            json.dump(message, f)
            
    def __len__(self):
        return len(self.anns)
    
    def __getitem__(self, idx):
        idx = idx % len(self)
        ann = self.anns[idx]
        img = Image.open(ann["img_path"]).convert("RGB")
        if self.input_transform:
            img = self.input_transform(img)
        return img, torch.Tensor(ann["target"])
