""" TEST FILE
Some tests are run to check the behaviour of the classes
    - 'Problem'
    - 'Learning'
in 'popex.popex_objects'. A toy problem is set up and tested together with a
simple learning scheme. There are no 'asserts' involved, but a confusion matrix
is computed.
"""

# -------------------------------------------------------------------------
#   Authors: Christoph Jaeggli, Julien Straubhaar and Philippe Renard
#   Year: 2018
#   Institut: University of Neuchatel
#
#   Copyright (c) 2018 Christoph Jaeggli
#
#   This program is distributed in the hope that it will be useful, but
#   WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

# General imports
import unittest
import numpy as np
from sklearn import tree
from pathlib import Path
import pickle

# Package imports
from popex.popex_objects import PoPEx, Problem, Learning, CatParam, ContParam

PATH_TEST = Path(Path(__file__).parent, 'test_dir/')


# Define a toy subclass of 'Learning'
class LearningTree(Learning):
    # 'Learning' is an abstract class with abstract methods
    #   - train
    #   - compute_p_eval_for
    #   - learn_value_of
    def __init__(self, thresh):
        self.__ls = None
        self.tresh = thresh

    # 'train' handles a popex object in order to predict 'popex.log_p_lik'
    def train(self, popex):
        attr, ans = [], []
        l_ind = 1*(popex.log_p_lik > self.tresh)

        # The models are unpickled in popex.model
        for imod, mod in enumerate(popex.model):
            attr_i = [np.sum(mod[0].param_val) + np.sum(mod[1].param_val)]
            if popex.cmp_log_p_lik[imod]:
                attr.append(attr_i)
                ans.append(l_ind[imod])

        # Train the tree from sklearn
        clf = tree.DecisionTreeClassifier()
        self.__ls = clf.fit(attr, ans)

    def compute_p_eval_for(self, model):
        val = np.sum(model[0].param_val) + np.sum(model[1].param_val)
        if abs(val - self.tresh * 11 / 15) > 0.5:
            return 1.
        else:
            return 0.0

    def learn_value_of(self, model):
        X_pred = [[np.sum(model[0].param_val) + np.sum(model[1].param_val)]]
        return self.__ls.predict(X_pred)[0]


# Set up toy functions for 'Problem'
def generate_m():
    categories = [[(0, 0.5)], [(0.5, 1.)]]
    m1 = CatParam(param_val=np.random.rand(5), categories=categories)
    m2 = CatParam(param_val=np.random.rand(6), categories=categories)
    m3 = ContParam(param_val=np.random.rand(4))
    return tuple([m1, m2, m3])


def compute_log_p_lik(model):
    return sum([np.sum(mod.param_val) for mod in model])


def compute_log_p_pri(model):
    return -max([np.sum(mod.param_val) for mod in model])


def compute_log_p_gen(model):
    return -min([np.sum(mod.param_val) for mod in model])


def get_hd_pri():
    hd_prior_param_ind = tuple(None for _ in range(3))
    hd_prior_param_val = tuple(None for _ in range(3))
    return hd_prior_param_ind, hd_prior_param_val


# Unit test class
class TestLearning(unittest.TestCase):

    def setUp(self):
        self.nmod = 2000
        self.npred = 100
        self.tresh = 7.5
        self.lrn = LearningTree(self.tresh)

    def test_Learning(self):
        # Main PoPEx object (generated in 'algorithm.run_popex_mp'
        popex = PoPEx(ncmax=(10, 20, 0),
                      nmtype=3,
                      path_res=PATH_TEST)

        # Definition of the problem (generated by the user)
        prbl = Problem(generate_m=generate_m,
                       compute_log_p_lik=compute_log_p_lik,
                       compute_log_p_pri=compute_log_p_pri,
                       compute_log_p_gen=compute_log_p_gen,
                       get_hd_pri=get_hd_pri)

        # Run pseudo chain
        for imod in range(self.nmod):
            mod = prbl.generate_m()
            log_p_lik = prbl.compute_log_p_lik(mod)
            log_p_pri = prbl.compute_log_p_pri(mod)
            log_p_gen = prbl.compute_log_p_gen(mod)

            # This method also pickles the model
            popex.add_model(imod, mod, log_p_lik, True,
                            log_p_pri, log_p_gen, (2, 5, 0))

            # In order to use it, we unpickle it again
            with open(Path(popex.path_res, popex.model[imod]), 'rb') as file:
                popex.model[imod] = pickle.load(file)

        self.lrn.train(popex)

        log_lik_ex = []
        log_lik_pred = []
        for i in range(self.npred):
            mod = prbl.generate_m()
            p = self.lrn.compute_p_eval_for(mod)
            rnd = np.random.rand(1)
            if rnd < p:
                log_lik_ex.append(1*(prbl.compute_log_p_lik(mod) > self.tresh))
                log_lik_pred.append(self.lrn.learn_value_of(mod))
        log_lik_ex = np.array(log_lik_ex)
        log_lik_pred = np.array(log_lik_pred)

        TP = np.sum(np.logical_and(log_lik_ex == 1, log_lik_pred == 1))
        FP = np.sum(np.logical_and(log_lik_ex == 0, log_lik_pred == 1))
        FN = np.sum(np.logical_and(log_lik_ex == 1, log_lik_pred == 0))
        TN = np.sum(np.logical_and(log_lik_ex == 0, log_lik_pred == 0))

        print('\nConfusion Matrix:\n\n  {:5d}  |  '
              '{:5d}\n---------|---------\n  '
              '{:5d}  |  {:5d}\n'.format(TP, FN, FP, TN))
        print('')
        print('   Secificity      = {:.4f}'.format(TN / (TN + FP)))
        print('   Pred. value (-) = {:.4f}'.format(TN / (TN + FN)))
        print('   R- / N          = {:.4f} (% of saved forwards)'
              .format((TN + FN) / (TP + TN + FP + FN)))
        print('')
        print('FN: We decided to NOT evaluate L(m) but we should have')
        print('FP: We decided to evaluate L(m) but we could have not')
        print('so FN is worse than FP')
        print('')
        print('')


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