import time
import os
import logging

import ConfigSpace as CS
import ConfigSpace.hyperparameters as CSH


import logging
logging.basicConfig(level=logging.INFO)
import sys
sys.path.append('/home/rkohli/aml_project/src/fmin')
from fmin.bohamiann import bohamiann

from utils import get_config_dictionary, get_upper_lower


import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.utils
import torchvision

import torchvision.transforms as transforms

import hpbandster.core.result as hpres
# import hpbandster.visualization as hpvis
import classification.operations as O

from classification import utils
from classification.cnn_model_micro import MicroProbabilisticCNN
from classification.train import train, architecture_search
from classification.datasets import K49, KMNIST
import sys
sys.path.append('/Users/ravinkohli/automl_ss19_project_ssh/aml-project/src')
import setttings

import pickle

class BohamiannWorker(object):
    def __init__(self, channels_init, trained_theta, run_dir, model_class=MicroProbabilisticCNN, 
                n_classes=10, ops_server_class=O.ENASOpsServer, dataset=KMNIST, **kwargs):

        super().__init__(**kwargs)
        self.run_dir = run_dir
        if torch.cuda.is_available:
            self.gpu_id = 1
        else:
            self.gpu_id = 0
        self.model_class = model_class
        self.n_classes = n_classes
        self.channels_init = channels_init
        self.trained_theta = trained_theta
        self.ops_server_class = ops_server_class
        self.run_dir = run_dir
        data_augmentations = transforms.ToTensor()
        self.train_dataset = dataset('./data', True, data_augmentations)
        self.test_dataset = dataset('./data', False, data_augmentations)

    def compute(self, x, budget, config, **kwargs):
        """
        Get model with hyperparameters from config generated by get_configspace()
        """
        config = get_config_dictionary(x, config)
        print("config", config)
        if (len(config.keys())<len(x)):
            return 100

        print('Retraining...')
        nn_model = self.model_class(
                n_classes=self.n_classes,
                n_cells=20,
                channels_init=self.channels_init,
                affine=True,
                ops_server=self.ops_server_class(affine=True),
                trained_theta=self.trained_theta
            )
        nn_model.fix_arc()
        print('n_params: %fM' % (nn_model.get_params_mle()/1e6))

        with open(os.path.join(self.run_dir, 'description.txt'), 'a') as o:
            o.write('dim_theta: %d\n' % nn_model.p_model.C.sum())
            o.write('channels_init: %d\n' % self.channels_init)
            o.write('n_params: %fM\n' % (nn_model.get_params_mle()/1e6))

        if self.gpu_id >= 0:
            nn_model = nn_model.cuda()
 
        optimizer = settings.opti_dict[config['optimizer']](nn_model.parameters(), lr=config['initial_lr'], weight_decay=config['weight_decay'])
        
        if config['lr_schedular'] == 'Cosine':
            lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, int(budget))
        elif config['lr_schedular'] == 'Exponential':
            lr_schedular = torch.optim.lr_schedular.ExponnetialLR(optimizer, gamma=0.1)
        elif config['lr_schedular'] == 'Plateau':
            lr_schedular = torch.optim.lr_schedular.ReduceLROnPlateau(optimizer)
        
        nn_model, test_loss = train(
                nn_model=nn_model,
                loss_func=torch.nn.CrossEntropyLoss,
                optimizer=optimizer,
                lr_scheduler=lr_scheduler,
                n_epochs=int(budget),
                batch_size=config['batch_size_retrain'],
                max_droppath_rate=config['max_droppath_rate'],
                clip_value=config['grad_clip_value'],
                weight_auxiliary=config['weight_auxiliary'],
                train_data=self.train_dataset,
                test_data=self.test_dataset,
                gpu_id=self.gpu_id,
                log_file=os.path.join(self.run_dir, 'train_log.csv')
            )
        
        return test_loss  # Hyperband always minimizes, so we want to minimise the error, error = 1-accuracy
    
    @staticmethod
    def get_configspace():
        """
        Define all the hyperparameters that need to be optimised and store them in config
        """
        cs = CS.ConfigurationSpace()

        initial_lr = CSH.UniformFloatHyperparameter('initial_lr', lower=1e-6, upper=1e-1, default_value='1e-2', log=True)
        optimizer = CSH.CategoricalHyperparameter('optimizer', settings.opti_dict.keys())
        batch_size = CSH.UniformIntegerHyperparameter('batch_size', lower=16, upper=32, default_value=24)
        cs.add_hyperparameters([initial_lr, optimizer, batch_size])
        
        lr_scheduler = CSH.CategoricalHyperparameter('scheduler', ['Exponential', 'Cosine', 'Plateau'])
        weight_decay = CSH.UniformFloatHyperparameter('weight_decay', lower=1e-5, upper=1e-3, default_value=3e-4, log=True)
        drop_path_rate = CSH.UniformFloatHyperparameter('max_droppath_rate', lower=0, upper=0.4, default_value=0.3, log=False)
        weight_auxiliary = CSH.UniformFloatHyperparameter('weight_auxiliary', lower='0', upper='0.4', log=False)
        grad_clip_value = CSH.UniformIntegerHyperparameter('grad_clip_value', lower=4, upper=8, default_value=5)
        cs.add_hyperparameters([lr_scheduler, drop_path_rate, weight_auxiliary, weight_decay, grad_clip_value])

        return cs


    def run_bohamiann(self, iterations=20):
        cs = self.__class__.get_configspace()
        
        lower, upper = get_upper_lower(cs)

        results = bohamiann(self.compute, lower, upper, num_iterations=iterations, cs=cs)
        
        x_best = results["x_opt"]
        config = get_config_dictionary(x_best, cs)
        print(config)

