import os
import json
from PIL import Image
import torch
from torch.utils.data import Dataset
import random
import numpy as np

class VG200Classification_PML(Dataset):
    def __init__(self, image_dir, phase, input_transform=None, noise_level='low_level', seed=1):
        self.image_dir = image_dir
        self.phase = phase
        self.noise_level = noise_level
        self.seed = seed
        self.ann_files = os.path.join(self.image_dir, '{}_vg200_{}_seed_{}.json'.format(
            self.phase, self.noise_level, seed))
        self.noise_probs_file = os.path.join(self.image_dir, '{}_vg200_{}_noise_probs_seed_{}.npy'.format(
            self.phase, self.noise_level, seed))
        self.load_noisy_anns()

        self.input_transform = input_transform
        
    def load_noisy_anns(self):
        if not os.path.exists(self.ann_files):
            self.generate_noisy_labels()
        self.anns = []
        json_data = json.load(open(self.ann_files, "r"))
        self.anns += json_data
        self.noise_probs = torch.Tensor(np.load(self.noise_probs_file))
        
    def generate_noisy_labels(self):
        ann_files = os.path.join(self.image_dir, '{}_vg200.json'.format(self.phase))
        self.load_anns(ann_files)
        
        np.random.seed(self.seed)  # Numpy module.
        random.seed(self.seed)  # Python random module.
        
        targets = np.array([np.array(ann["target"]) for ann in self.anns])
        if self.noise_level == 'high_level':
            alpha = 0.1 * np.random.randint(5, 7, (1, targets.shape[1]), dtype=targets.dtype)
        else:
            alpha = 0.1 * np.random.randint(2, 6, (1, targets.shape[1]), dtype=targets.dtype)
            
        mask = np.random.random((targets.shape)) < alpha
        targets[mask & (targets==0)] = 1
        for i in range(len(self.anns)):
            self.anns[i]["target"] = targets[i].tolist()
        
        with open(self.ann_files, 'w') as f:
            json.dump(self.anns, f)
            
        np.save(self.noise_probs_file, alpha)
        
    def load_anns(self, ann_files):
        if not os.path.exists(ann_files):
            self.extract_anns(ann_files)
        self.anns = []
        json_data = json.load(open(ann_files, "r"))
        self.anns += json_data
        
    def extract_anns(self, ann_files):
        raise NotImplementedError("No annotation file found at %s" % ann_files)
            
    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"])
