"""
@package tests.test_swarm
unit tests for @ref pmsco.optimizers.swarm

the purpose of these tests is to help debugging the code.

to run the tests, change to the directory which contains the tests directory, and execute =nosetests=.

@pre nose must be installed (python-nose package on Debian).

@author Matthias Muntwiler, matthias.muntwiler@psi.ch

@copyright (c) 2015 by Paul Scherrer Institut @n
Licensed under the Apache License, Version 2.0 (the "License"); @n
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  http://www.apache.org/licenses/LICENSE-2.0
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import os
import os.path
import random
import shutil
import tempfile
import unittest

import pmsco.optimizers.swarm as swarm
import pmsco.project as project

POP_SIZE = 5


class TestSwarmPopulation(unittest.TestCase):
    def setUp(self):
        random.seed(0)
        self.test_dir = tempfile.mkdtemp()
        self.model_space = project.ModelSpace()
        
        self.model_space.add_param('A', 1.5, 1.0, 2.0, 0.1)
        self.model_space.add_param('B', 2.5, 2.0, 3.0, 0.1)
        self.model_space.add_param('C', 3.5, 3.0, 4.0, 0.1)
        self.expected_names = ('_gen', '_model', '_particle', '_rfac', 'A', 'B', 'C')

        self.size = POP_SIZE
        self.pop = swarm.SwarmPopulation()

        self.optimum1 = {'A': 1.045351, 'B': 2.346212, 'C': 3.873627}

    def tearDown(self):
        # after each test method
        self.pop = None
        shutil.rmtree(self.test_dir)

    @classmethod
    def setup_class(cls):
        # before any methods in this class
        pass
 
    @classmethod
    def teardown_class(cls):
        # teardown_class() after any methods in this class
        pass

    def rfactor1(self, pos):
        r = (pos['A'] - self.optimum1['A'])**2 \
            + (pos['B'] - self.optimum1['B'])**2 \
            + (pos['C'] - self.optimum1['C'])**2
        r /= 3.0
        return r

    def test_best_friend(self):
        self.pop.setup(self.size, self.model_space)
        self.pop.best['_rfac'] = np.arange(self.size)
        friend = self.pop.best_friend(0)
        self.assertNotIsInstance(friend, np.ndarray)
        self.assertEqual(friend.dtype.names, self.expected_names)
        
    def test_advance_particle(self):
        self.pop.setup(self.size, self.model_space)
        
        self.pop.pos['A'] = np.linspace(1.0, 2.0, POP_SIZE)
        self.pop.pos['B'] = np.linspace(2.0, 3.0, POP_SIZE)
        self.pop.pos['C'] = np.linspace(3.0, 4.0, POP_SIZE)
        self.pop.pos['_rfac'] = np.linspace(2.0, 1.0, POP_SIZE)
        
        self.pop.vel['A'] = np.linspace(-0.1, 0.1, POP_SIZE)
        self.pop.vel['B'] = np.linspace(-0.1, 0.1, POP_SIZE)
        self.pop.vel['C'] = np.linspace(-0.1, 0.1, POP_SIZE)
        
        pos0 = self.pop.pos['A'][0]
        self.pop.advance_particle(0)
        pos1 = self.pop.pos['A'][0]
        self.assertNotAlmostEqual(pos0, pos1, delta=0.001)
        
        for key in ['A','B','C']:
            for pos in self.pop.pos[key]:
                self.assertGreaterEqual(pos, self.model_space.min[key])
                self.assertLessEqual(pos, self.model_space.max[key])

    def test_is_converged(self):
        self.pop.setup(self.size, self.model_space)
        self.assertFalse(self.pop.is_converged())
        i_sample = 0
        result = self.pop.pos[i_sample]
        for i in range(POP_SIZE):
            rfac = 1.0 - float(i)/POP_SIZE
            self.pop.add_result(result, rfac)
        self.assertFalse(self.pop.is_converged())
        for i in range(POP_SIZE):
            rfac = (1.0 - float(i)/POP_SIZE) / 1000.0
            self.pop.add_result(result, rfac)
        self.assertTrue(self.pop.is_converged())
        
    def test_convergence_1(self):
        self.pop.setup(self.size, self.model_space)

        self.pop.pos['A'] = np.linspace(1.0, 2.0, POP_SIZE)
        self.pop.pos['B'] = np.linspace(2.0, 3.0, POP_SIZE)
        self.pop.pos['C'] = np.linspace(3.0, 4.0, POP_SIZE)
        self.pop.pos['_rfac'] = np.linspace(2.0, 1.0, POP_SIZE)

        self.pop.vel['A'] = np.linspace(-0.1, 0.1, POP_SIZE)
        self.pop.vel['B'] = np.linspace(-0.1, 0.1, POP_SIZE)
        self.pop.vel['C'] = np.linspace(-0.1, 0.1, POP_SIZE)

        for i in range(10):
            self.pop.advance_population()
            for pos in self.pop.pos:
                self.pop.add_result(pos, self.rfactor1(pos))

        for pos in self.pop.pos:
            self.assertLess(pos['_rfac'], 0.2)


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