import argparse
import sys
import os
import time
import numpy as np
import torch
import torchvision.transforms as transforms

sys.path.append("./Deep Version")
from utils import init_random_seed, clear_old_logs, display, get_learning_rate
from utils import ValueMeter, AverageMeter, ProgressMeter, AveragePrecisionMeter
from dataset.MultiScaleCrop import MultiScaleCrop
from dataset.get_PML_datasets import get_deannotated_datasets, get_test_dataset
from dataset.get_dataset_cfg import get_dataset_cfg
from dataset.CoCoClassification import CoCoClassification


parser = argparse.ArgumentParser()
parser.add_argument('exp', type=str,
                    help='name of experiment')
### Model
parser.add_argument("--model", type=str, default="PASE",
                    help='model to train and evaluate')
parser.add_argument("--backbone", type=str, default="resnet18",
                    help='model to train and evaluate')
parser.add_argument('--no_pretrain', action='store_true',
                    help='learn from scratch')

### Data
parser.add_argument('--dataset', type=str, default="coco",
                    help='dataset on which experiment is conducted')
parser.add_argument('--workers', '-j', type=int, default=8,
                    help='number of data loading workers (default: 8)')
parser.add_argument('--extra_val', action='store_true',
                    help='split an extra validation set from training set')
parser.add_argument('--noise_level', type=str, default="low_level",
                    choices=["low_level", "high_level"],
                    help='level of noise to sythesize partial-multi label data set')
# Data augmentation
parser.add_argument('--img_size', type=int, default=448,
                    help='size of input images')

### Train
parser.add_argument('--seed', default=1, type=int,
                    help='seed for initializing training')
parser.add_argument('--start_epoch', type=int, default=0,
                    help='manual epoch number (useful on restarts)')
parser.add_argument('--epochs', type=int, default=80,
                    help='number of total epochs to run')
parser.add_argument('--batch_size', '-b', default=32, type=int,
                    help='mini-batch size (default: 32)')
parser.add_argument('--lr', '-lr', type=float, default=5e-5,
                    help='maximal learning rate')
parser.add_argument('--meta_lr', '-meta_lr', type=float, default=1.0,
                    help='meta learning rate')
parser.add_argument('--weight_decay', '-wd', type=float, default=1e-2,
                    help='weight decay (default: 1e-2)')
parser.add_argument('--early_stop', action='store_true',
                    help='apply early stop')
parser.add_argument('--cuda', action='store_true',
                    help='whether to run on GPU')
parser.add_argument('--print_freq', '-p', type=int, default=1,
                    help='print frequency (default: 100)')

# EMA
parser.add_argument('--ema_decay', type=float, default=0.9997,
                    help='decay of model ema')

### Validation
parser.add_argument('--threshold', '-t', type=float, default=0.5,
                    help='threshold value')

def load_default_cfgs(args):
    args.exp_dir = os.path.join('Exps', args.model, args.exp, args.dataset)
    args.log_file = os.path.join(args.exp_dir, 'log.txt')
    args.checkpoint_dir = os.path.join(args.exp_dir, 'checkpoint')
    get_dataset_cfg(args)

def get_data_augmentation(args):
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
        
    train_transform_list = [transforms.Resize((args.img_size+64, args.img_size+64)),
                            MultiScaleCrop(args.img_size, scales=(1.0, 0.875, 0.75, 0.625, 0.5), max_distort=2),
                            transforms.RandomHorizontalFlip(),
                            transforms.ToTensor(),
                            normalize]
        
    train_transform = transforms.Compose(train_transform_list)
    test_transform = transforms.Compose([transforms.Resize((args.img_size, args.img_size)),
                                         transforms.ToTensor(),
                                         normalize])
    
    return train_transform, test_transform

if __name__ == '__main__':
    args = parser.parse_args()
    load_default_cfgs(args)

    # Set random seed
    init_random_seed(seed=args.seed)
    
    train_transform, test_transform = get_data_augmentation(args)
    
    path = os.path.join(args.dataset_dir, 'PML_annotations')
    if not os.path.exists(path):
        os.makedirs(path)
    
    CoCoClassification(args.dataset_dir, 'train', test_transform)
    CoCoClassification(args.dataset_dir, 'val', test_transform)
    args.noise_level = "low_level"
    get_deannotated_datasets(train_transform, test_transform, args)
    args.noise_level = "high_level"
    get_deannotated_datasets(train_transform, test_transform, args)
    