"""Script to test the acquisition functions."""
from __future__ import print_function
from __future__ import absolute_import

import os
import unittest
import numpy as np

from ase.ga.data import DataConnection

from catlearn.api.ase_data_setup import get_unique, get_train
from catlearn.featurize.setup import FeatureGenerator
from catlearn.regression import GaussianProcess
from catlearn.active_learning.acquisition_functions import (
    rank, classify, probability_density)
from catlearn.active_learning.algorithm import ActiveLearning


wkdir = os.getcwd()

train_size, test_size = 30, 30


def classifier(atoms):
    """Simple function to classify atoms objects."""
    return atoms.get_chemical_formula()


class TestAcquisition(unittest.TestCase):
    """Test out the various acquisition routines."""

    def get_data(self):
        """Generate features from atoms objects."""
        # Connect database generated by a GA search.
        gadb = DataConnection('{}/data/gadb.db'.format(wkdir))

        # Get all relaxed candidates from the db file.
        print('Getting candidates from the database')
        all_cand = gadb.get_all_relaxed_candidates(use_extinct=False)

        # Setup the test and training datasets.
        testset = get_unique(atoms=all_cand, size=test_size, key='raw_score')

        trainset = get_train(atoms=all_cand, size=train_size,
                             taken=testset['taken'], key='raw_score')

        # Clear out some old saved data.
        for i in trainset['atoms']:
            del i.info['data']['nnmat']

        # Initiate the fingerprint generators with relevant input variables.
        print('Getting the fingerprints')
        f = FeatureGenerator()

        train_features = f.return_vec(
            trainset['atoms'], [f.nearestneighbour_vec])
        test_features = f.return_vec(
            testset['atoms'], [f.nearestneighbour_vec])

        train_targets = []
        for a in trainset['atoms']:
            train_targets.append(a.info['key_value_pairs']['raw_score'])
        test_targets = []
        for a in testset['atoms']:
            test_targets.append(a.info['key_value_pairs']['raw_score'])

        return train_features, train_targets, trainset['atoms'], \
            test_features, test_targets, testset['atoms']

    def test_acquisition(self):
        """Test acquisition functions."""
        train_features, train_targets, train_atoms, test_features, \
            test_targets, test_atoms = self.get_data()
        # Test prediction routine with gaussian kernel.
        kdict = [{'type': 'gaussian', 'width': 1., 'scaling': 1.}]
        gp = GaussianProcess(
            train_fp=train_features, train_target=train_targets,
            kernel_list=kdict, regularization=1e-3,
            optimize_hyperparameters=True, scale_data=True)
        pred = gp.predict(
            test_fp=test_features, test_target=test_targets,
            get_validation_error=True, get_training_error=True,
            uncertainty=True)

        print('gaussian prediction (rmse):',
              pred['validation_error']['rmse_average'])

        acq = rank(
            targets=train_targets, predictions=pred['prediction'],
            uncertainty=pred['uncertainty'], train_features=train_features,
            test_features=test_features, metrics=[
                'cdf', 'optimistic', 'pdf', 'UCB', 'EI', 'PI'])
        self.assertTrue(len(acq['optimistic']) == len(pred['prediction']))
        self.assertTrue(len(acq['pdf']) == len(pred['prediction']))
        self.assertTrue(len(acq['UCB']) == len(pred['prediction']))
        self.assertTrue(len(acq['EI']) == len(pred['prediction']))
        self.assertTrue(len(acq['PI']) == len(pred['prediction']))

        acq = classify(
            classifier, train_atoms, test_atoms, targets=train_targets,
            predictions=pred['prediction'], uncertainty=pred['uncertainty'],
            train_features=train_features, test_features=test_features,
            metrics=['optimistic', 'pdf', 'UCB', 'EI', 'PI'])
        self.assertTrue(len(acq['optimistic']) == len(pred['prediction']))
        self.assertTrue(len(acq['pdf']) == len(pred['prediction']))
        self.assertTrue(len(acq['UCB']) == len(pred['prediction']))
        self.assertTrue(len(acq['EI']) == len(pred['prediction']))
        self.assertTrue(len(acq['PI']) == len(pred['prediction']))

    def test_surrogate_model(self):
        train_features, train_targets, train_atoms, test_features, \
            test_targets, test_atoms = self.get_data()

        sg0 = ActiveLearning(_surrogate_model,
                             train_features, train_targets)
        batch_size = 10
        sg0.test_acquisition(batch_size=batch_size)
        to_acquire, _ = sg0.acquire(test_features, batch_size=batch_size)

        self.assertTrue(len(to_acquire) == batch_size)

        sg1 = ActiveLearning(_surrogate_model,
                             train_features, train_targets)
        batch_size = 3
        to_acquire, _ = sg1.acquire(test_features, batch_size=batch_size)

        self.assertTrue(len(to_acquire) == batch_size)

        sg2 = ActiveLearning(_surrogate_model,
                             train_features, train_targets)
        batch_size = 1
        to_acquire, _ = sg2.acquire(test_features, batch_size=batch_size)

        self.assertTrue(len(to_acquire) == batch_size)

        sg0.ensemble_test(size=2, batch_size=batch_size)


def _surrogate_model(train_features, train_targets,
                     test_features=None, test_targets=None):
    kdict = [{'type': 'gaussian', 'width': 0.5}]
    gp = GaussianProcess(
        train_fp=train_features, train_target=train_targets,
        kernel_list=kdict, regularization=1e-3,
        optimize_hyperparameters=False, scale_data=True)
    if test_targets is None:
        get_validation_error = False
    else:
        get_validation_error = True
    score = gp.predict(
        test_fp=test_features, test_target=test_targets,
        get_validation_error=get_validation_error,
        get_training_error=False,
        uncertainty=True)
    to_acquire = np.argsort(probability_density(0.,
                                                score['prediction'],
                                                score['uncertainty']))
    return to_acquire, score


if __name__ == '__main__':
    unittest.main()
