import os
import json
import random
import numpy as np
import math

from dataset.CoCoClassification import CoCoClassification
from dataset.CoCoClassification_PML import CoCoClassification_PML
from dataset.VG200Classification import VG200Classification
from dataset.VG200Classification_PML import VG200Classification_PML

def get_deannotated_datasets(train_transform, test_transform, args):
    if args.dataset == 'coco' or args.dataset == 'coco14':
        if args.extra_val:
            ann_files = os.path.join(args.dataset_dir, 'PML_annotations', 'train-train_90p_coco2014.json')
            if not os.path.exists(ann_files):
                train2trainval(os.path.join(args.dataset_dir, 'PML_annotations'), 'coco2014', val_ratio=0.1)
            train_dataset = CoCoClassification_PML(os.path.join(args.dataset_dir, 'PML_annotations'), 'train-train_90p', train_transform, args.noise_level, args.seed)
            val_dataset = CoCoClassification(args.dataset_dir, 'train-val_10p', test_transform)
        else:
            train_dataset = CoCoClassification_PML(os.path.join(args.dataset_dir, 'PML_annotations'), 'train', train_transform, args.noise_level, args.seed)
            val_dataset = CoCoClassification(args.dataset_dir, 'val', test_transform)
    elif args.dataset == 'vg200':
        if args.extra_val:
            ann_files = os.path.join(args.dataset_dir, 'train-train_90p_vg200.json')
            if not os.path.exists(ann_files):
                train2trainval(args.dataset_dir, 'vg200', val_ratio=0.1)
            train_dataset = VG200Classification_PML(args.dataset_dir, 'train-train_90p', train_transform, args.noise_level, args.seed)
            val_dataset = VG200Classification(args.dataset_dir, 'train-val_10p', test_transform)
        else:
            train_dataset = VG200Classification_PML(args.dataset_dir, 'train', train_transform, args.noise_level, args.seed)
            val_dataset = VG200Classification(args.dataset_dir, 'test', test_transform)
    else:
        raise NotImplementedError("Unknown dataset %s" % args.dataset)

    return train_dataset, val_dataset

def get_test_dataset(test_transform, args):
    if args.dataset == 'coco' or args.dataset == 'coco14':
        test_dataset = CoCoClassification(args.dataset_dir, 'val', test_transform)
    elif args.dataset == 'vg200':
        test_dataset = VG200Classification(args.dataset_dir, 'test', test_transform)
    else:
        raise NotImplementedError("Unknown dataset %s" % args.dataset)
        
    return test_dataset
    
def train2trainval(image_dir, dataset, val_ratio=0.1, seed=1, prefix='train'):
    np.random.seed(seed)  # Numpy module.
    random.seed(seed)  # Python random module.
        
    ann_file = os.path.join(image_dir, '{}_{}.json'.format(prefix, dataset))
    if os.path.exists(ann_file):
        anns = json.load(open(ann_file, "r"))
    else:
        print("no annotation file is found at '{}'".format(ann_file))
        exit(-1)
    
    val_size = math.ceil(len(anns) * val_ratio)
    val_idx = np.random.choice(len(anns), val_size, replace=False)
    val_idx = np.sort(val_idx)
    
    train_anns, val_anns = [], []
    cur = 0
    for i in range(len(anns)):
        if cur < val_size and i == val_idx[cur]:
            val_anns.append(anns[i])
            cur += 1
        else:
            train_anns.append(anns[i])
    print(len(anns), len(train_anns), len(val_anns))
    
    train_ann_file = os.path.join(image_dir, '{}-train_{:d}p_{}.json'.format(prefix, int(100 - 100*val_ratio), dataset))
    with open(train_ann_file, 'w') as f:
        json.dump(train_anns, f)
        
    val_ann_file = os.path.join(image_dir, '{}-val_{:d}p_{}.json'.format(prefix, int(100*val_ratio), dataset))
    with open(val_ann_file, 'w') as f:
        json.dump(val_anns, f)