import sys
sys.path.append('.')
import yaml
import os
from shutil import rmtree

import torch
import torchvision.transforms as transforms

from src.data import get_datasets, get_dataloaders
from src.model import initialize_model
from src.optimizer import get_optimizer
from src.criterion import get_criterion
from src.scheduler import get_scheduler
import yaml



def get_transforms(conf):
    model_name = conf['model']['name']
    if 'efficientdet' in model_name:
        train_transform = transforms.Compose([transforms.Resize((512,512)),
                                    transforms.RandomVerticalFlip(p=0.5),
                                    transforms.RandomHorizontalFlip(p=0.5),
                                    transforms.ToTensor(),
                                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                    ])
        valid_transform = transforms.Compose([transforms.Resize((512,512)),
                                    transforms.ToTensor(),
                                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                    ])
        test_transform = transforms.Compose([transforms.Resize((512,512)),
                                    transforms.ToTensor(),
                                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                    ])

    elif 'resnet' in model_name:
        train_transform = transforms.Compose([transforms.Resize((224,224)),
                                    transforms.RandomVerticalFlip(p=0.5),
                                    transforms.RandomHorizontalFlip(p=0.5),
                                    transforms.ToTensor(),
                                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                    ])
        valid_transform = transforms.Compose([transforms.Resize((224,224)),
                                    transforms.ToTensor(),
                                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                    ])
        test_transform = transforms.Compose([transforms.Resize((224,224)),
                                    transforms.ToTensor(),
                                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                    ])

    return train_transform, valid_transform, test_transform


def get_experiment_dir(conf):
    experiments_dir = conf['experiments_dir']
    if not os.path.exists(experiments_dir):
        os.makedirs(experiments_dir)

    experiment_dir = os.path.join(conf['experiments_dir'],
                                  conf['experiment_code'])
    if conf['task'] == 'training':
        if not os.path.exists(experiment_dir):
            os.makedirs(experiment_dir)
            save_conf = os.path.join(experiment_dir, conf['task'] + '.yaml')
            with open(save_conf, 'w') as fp:
                yaml.dump(conf, fp)
        else:
            print(f'Experiment dir {experiment_dir} exists.')
            exit()
    return experiment_dir

def get_test_config(conf):

    # Check if results dir exists and create it
    results_dir = conf['results_dir']
    if os.path.exists(results_dir): rmtree(results_dir)
    os.makedirs(results_dir)

    # Get transforms
    transforms = get_transforms(conf)
    conf['train_transform'] = transforms[0]
    conf['valid_transform'] = transforms[1]
    conf['test_transform'] = transforms[2]

    # Get datasets
    dataset_name = conf['data']['name']
    print(f'Dataset: {dataset_name}')
    _, _, conf['test_dataset'] = get_datasets(conf)

    # Check if GPU is available
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    conf['device'] = device
    print(f'Device: {device}')

    # Initialize model
    model_name = conf['model']['name']
    print(f'Model: {model_name}')
    conf['model'] = initialize_model(conf)

    return conf


def get_config(yaml_path):
    # Load YAML conf
    conf = yaml.safe_load(open(yaml_path, 'r'))

    # Task
    task = conf['task']
    if task in ['training', 'evaluation']:
        print(f'Task: {task}')

    elif task == 'testing':
        return get_test_config(conf)

    else:
        print(f'Task {task} not supported.')
        exit()

    # Get experiment directory
    experiment_dir = get_experiment_dir(conf)
    conf['experiment_dir'] = experiment_dir
    print(f'Experiment directory: {experiment_dir}')

    # Get transforms
    transforms = get_transforms(conf)
    conf['train_transform'] = transforms[0]
    conf['valid_transform'] = transforms[1]
    conf['test_transform'] = transforms[2]

    # Get datasets
    dataset_name = conf['data']['name']
    print(f'Dataset: {dataset_name}')
    datasets = get_datasets(conf)
    conf['train_dataset'] = datasets[0]
    conf['valid_dataset'] = datasets[1]
    conf['test_dataset'] = datasets[2]
    conf['patients_dataset'] = datasets[3]

    # Get dataloaders
    dataloaders = get_dataloaders(conf)
    train_dataloader = dataloaders[0]
    valid_dataloader = dataloaders[1]
    test_dataloader = dataloaders[2]
    conf['dataloaders'] = {'train':train_dataloader,
                           'valid':valid_dataloader,
                           'test': test_dataloader}

    # Check if GPU is available
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    conf['device'] = device
    print(f'Device: {device}')

    # Initialize model
    model_name = conf['model']['name']
    print(f'Model: {model_name}')
    conf['model'] = initialize_model(conf)

    # Only in traning task
    if task == 'training':
        # Get optimizer
        optimizer_name = conf['optimizer']['name']
        print(f'Optimizer: {optimizer_name}')
        conf['optimizer'] = get_optimizer(conf)

        # Get criterion
        criterion_name = conf['criterion']['name']
        print(f'Criterion: {criterion_name}')
        conf['criterion'] = get_criterion(conf)

        # Get scheduler
        scheduler_name = conf['scheduler']['name']
        print(f'Scheduler: {scheduler_name}')
        conf['scheduler'] = get_scheduler(conf)

    # Only in evaluation task
    elif task == 'evaluation':
        path = os.path.join(conf['experiment_dir'], 'best_weights.pt')
        if os.path.exists(path):
            print(f'Loading weights from {path}')
            conf['best_weights'] = torch.load(path)
        else:
            print(f'Experiment weights {path} not found.')
            exit()

    return conf