from unittest import TestCase

import numpy as np
import os
import tensorflow as tf
from keras import backend as K
from sklearn.preprocessing import MinMaxScaler

import darwin
from darwin.zoo.networks import MLP
from darwin.training import pipeline, experiment
from darwin.zoo import metrics
from darwin.database.molecules import DBInterface


class MLPTrainer:

    def setUp(self):
        # Set up training enviroment
        darwin.gpu(0)
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # to disable nasty tf logging
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        self.tf_session = tf.InteractiveSession(config=config)
        K.set_session(self.tf_session)

        # Connect to the molecular database
        db = DBInterface(db_name='db_proba', username='user_proba', password='proba', port=27017)

        # Get dataset
        self.data = db.datasets.get_dataset(dataset_name='bbb_etkins', target_name='activity',
                                            mol_repr='fp', test_ratio=0.15)

    def tearDown(self):
        self.tf_session.close()

    def run_experiment(self):
        @pipeline
        class preprocessor:
            data = self.data
            minmax = MinMaxScaler()
            pipe = minmax

        @experiment
        class mnist_experiment:
            class classifier:
                pipe = MLP(input_shape=self.data.X.shape[1], output_dim=1, output_nonlinearity='sigmoid')
                pipe_batch_size = 32
                pipe_epochs = 5
                target = metrics.f1
                hyperparameter_space = {'hidden_sizes': [[200, 100], [100]]}

            experiment = preprocessor, classifier

        mnist_experiment.make()


pipe_test = MLPTrainer()
pipe_test.setUp()
pipe_test.run_experiment()