"""
@package tests.test_population
unit tests for @ref pmsco.optimizers.population

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
"""

import numpy as np
import os
import os.path
import random
import shutil
import tempfile
import unittest

import pmsco.optimizers.population as population
import pmsco.project as project
from pmsco.helpers import BraceMessage as BMsg

POP_SIZE = 5


class TestPopulation(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 = population.Population()

        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 reorder_pop_array(self, x):
        """
        re-order the columns of a population-like array.

        this makes it easier to create constant arrays with arbitrary column order
        since population arrays are ordered alphabetically.

        @param x: numpy structured array
        @return: numpy structured array having the same dtype as the population arrays.
        """
        y = np.zeros(x.shape, dtype=self.pop.get_pop_dtype(self.pop.model_start))
        for name in y.dtype.names:
            y[name] = x[name]
        return y

    def assert_pop_array_equal(self, x, y, msg=""):
        """
        array equality test for population-like structured numpy arrays.

        - check for same column names and kinds
        - check for same shape
        - check float columns using np.testing.assert_array_almost_equal
        - check all other columns using np.testing.assert_array_equal

        @param x: structured numpy array.
        @param y: structured numpy array.
        @param msg: optional error message.
        @return: None
        @raise AssertionError if arrays are not equal.
        """
        self.assertEqual(x.dtype.names, y.dtype.names, msg=BMsg("{0} (dtype)", msg))
        self.assertEqual(x.shape, y.shape, msg=BMsg("{0} (shape)", msg))
        for n, t in x.dtype.fields.items():
            self.assertEqual(t[0].kind, y.dtype.fields[n][0].kind, msg=BMsg("{0} (column {1} kind)", msg, n))
            if t[0].kind in {'f', 'c'}:
                np.testing.assert_array_almost_equal(x[n], y[n], err_msg="{0} (column {1})".format(msg, n))
            else:
                np.testing.assert_array_equal(x[n], y[n], err_msg="{0} (column {1})".format(msg, n))

    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_setup(self):
        self.pop.setup(self.size, self.model_space)
        self.assertEqual(self.pop.pos.dtype.names, self.expected_names)
        self.assertEqual(self.pop.pos.shape, (POP_SIZE,))
        np.testing.assert_array_equal(np.arange(POP_SIZE), self.pop.pos['_particle'])
        np.testing.assert_array_equal(np.zeros(POP_SIZE), self.pop.pos['_gen'])
        np.testing.assert_array_equal(np.arange(POP_SIZE), self.pop.pos['_model'])
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][0], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][1], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][2], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][3], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][4], 3)
        self.assertEqual(0, self.pop.generation)
        self.assertEqual(POP_SIZE, self.pop.model_count)

    def test_setup_with_results(self):
        data_dir = os.path.dirname(os.path.abspath(__file__))
        data_file = os.path.join(data_dir, "test_swarm.setup_with_results.1.dat")
        self.pop.setup(self.size, self.model_space, seed_file=data_file, recalc_seed=False)

        self.assertEqual(self.pop.pos.dtype.names, self.expected_names)
        self.assertEqual(self.pop.pos.shape, (POP_SIZE,))
        self.assertEqual(self.pop.generation, 0)
        self.assertEqual(self.pop.model_count, POP_SIZE)
        np.testing.assert_array_equal(self.pop.pos['_particle'], np.arange(POP_SIZE))
        np.testing.assert_array_equal(self.pop.pos['_gen'], [0, 0, -1, 0, 0])
        np.testing.assert_array_equal(self.pop.pos['_model'], np.arange(POP_SIZE))
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][0], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][1], 3)
        self.assertAlmostEqual(0.6, self.pop.pos['_rfac'][2], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][3], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][4], 3)
        self.assertAlmostEqual(1.3, self.pop.pos['A'][1], 3)
        self.assertAlmostEqual(1.1, self.pop.pos['A'][2], 3)
        self.assertAlmostEqual(1.5, self.pop.pos['A'][0], 3)
        self.assertAlmostEqual(2.3, self.pop.pos['B'][1], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['B'][2], 3)
        self.assertAlmostEqual(2.5, self.pop.pos['B'][0], 3)
        self.assertGreaterEqual(4.0, self.pop.pos['C'][1], 3)
        self.assertAlmostEqual(3.1, self.pop.pos['C'][2], 3)
        self.assertAlmostEqual(3.5, self.pop.pos['C'][0], 3)

    def test_setup_with_results_recalc(self):
        data_dir = os.path.dirname(os.path.abspath(__file__))
        data_file = os.path.join(data_dir, "test_swarm.setup_with_results.1.dat")
        self.pop.setup(self.size, self.model_space, seed_file=data_file, recalc_seed=True)

        self.assertEqual(self.pop.pos.dtype.names, self.expected_names)
        self.assertEqual(self.pop.pos.shape, (POP_SIZE,))
        self.assertEqual(self.pop.generation, 0)
        self.assertEqual(self.pop.model_count, POP_SIZE)
        np.testing.assert_array_equal(self.pop.pos['_particle'], np.arange(POP_SIZE))
        np.testing.assert_array_equal(self.pop.pos['_gen'], [0, 0, 0, 0, 0])
        np.testing.assert_array_equal(self.pop.pos['_model'], np.arange(POP_SIZE))
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][0], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][1], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][2], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][3], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][4], 3)
        self.assertAlmostEqual(1.3, self.pop.pos['A'][1], 3)
        self.assertAlmostEqual(1.1, self.pop.pos['A'][2], 3)
        self.assertAlmostEqual(1.5, self.pop.pos['A'][0], 3)
        self.assertAlmostEqual(2.3, self.pop.pos['B'][1], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['B'][2], 3)
        self.assertAlmostEqual(2.5, self.pop.pos['B'][0], 3)
        self.assertGreaterEqual(4.0, self.pop.pos['C'][1], 3)
        self.assertAlmostEqual(3.1, self.pop.pos['C'][2], 3)
        self.assertAlmostEqual(3.5, self.pop.pos['C'][0], 3)

    def test_setup_with_partial_results(self):
        self.model_space.add_param('D', 4.5, 4.0, 5.0, 0.1)
        self.expected_names = ('_gen', '_model', '_particle', '_rfac', 'A', 'B', 'C', 'D')

        data_dir = os.path.dirname(os.path.abspath(__file__))
        data_file = os.path.join(data_dir, "test_swarm.setup_with_results.1.dat")
        self.pop.setup(self.size, self.model_space, seed_file=data_file, recalc_seed=False)

        self.assertEqual(self.pop.pos.dtype.names, self.expected_names)
        self.assertEqual(self.pop.pos.shape, (POP_SIZE,))
        self.assertEqual(self.pop.generation, 0)
        self.assertEqual(self.pop.model_count, POP_SIZE)
        np.testing.assert_array_equal(self.pop.pos['_particle'], np.arange(POP_SIZE))
        np.testing.assert_array_equal(self.pop.pos['_gen'], [0, 0, -1, 0, 0])
        np.testing.assert_array_equal(self.pop.pos['_model'], np.arange(POP_SIZE))
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][0], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][1], 3)
        self.assertAlmostEqual(0.6, self.pop.pos['_rfac'][2], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][3], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['_rfac'][4], 3)
        self.assertAlmostEqual(1.3, self.pop.pos['A'][1], 3)
        self.assertAlmostEqual(1.1, self.pop.pos['A'][2], 3)
        self.assertAlmostEqual(1.5, self.pop.pos['A'][0], 3)
        self.assertAlmostEqual(2.3, self.pop.pos['B'][1], 3)
        self.assertAlmostEqual(2.1, self.pop.pos['B'][2], 3)
        self.assertAlmostEqual(2.5, self.pop.pos['B'][0], 3)
        self.assertGreaterEqual(4.0, self.pop.pos['C'][1], 3)
        self.assertAlmostEqual(3.1, self.pop.pos['C'][2], 3)
        self.assertAlmostEqual(3.5, self.pop.pos['C'][0], 3)

    def test_pos_gen(self):
        self.pop.setup(self.size, self.model_space)
        for index, item in enumerate(self.pop.pos_gen()):
            self.assertIsInstance(item, dict)
            self.assertEqual(set(item.keys()), set(self.expected_names))
            self.assertEqual(item['_particle'], index)

    def test_get_pop_dtype(self):
        result = population.Population.get_pop_dtype({'A': 1.045, 'C': 3.873, '_D': 4.381, 'B': 2.346})
        expected = [('_gen', 'i8'), ('_model', 'i8'), ('_particle', 'i8'), ('_rfac', 'f8'),
                    ('A', 'f8'), ('B', 'f8'), ('C', 'f8')]
        self.assertEqual(result, expected)
        result = population.Population.get_pop_dtype(['A', 'C', '_D', 'B'])
        self.assertEqual(result, expected)

    def test_get_model_dtype(self):
        result = population.Population.get_model_dtype({'A': 1.045, 'C': 3.873, '_D': 4.381, 'B': 2.346})
        expected = [('A', 'f8'), ('B', 'f8'), ('C', 'f8')]
        self.assertEqual(result, expected)
        result = population.Population.get_model_dtype(['A', 'C', '_D', 'B'])
        self.assertEqual(result, expected)

    def test_get_model_array(self):
        result = population.Population.get_model_array({'A': 1.045, 'C': 3.873, '_D': 4.381, 'B': 2.346})
        dt = [('A', 'f8'), ('B', 'f8'), ('C', 'f8')]
        expected = np.array([(1.045, 2.346, 3.873)], dtype=dt)
        np.testing.assert_array_equal(result, expected)

    def test_randomize(self):
        self.pop.setup(self.size, self.model_space)
        self.pop.randomize()
        m = np.mean(self.pop.pos['A'])
        self.assertGreaterEqual(m, self.model_space.min['A'])
        self.assertLessEqual(m, self.model_space.max['A'])

    def test_seed(self):
        self.pop.setup(self.size, self.model_space)
        self.pop.seed(self.model_space.start)
        self.assertAlmostEqual(self.pop.pos['A'][0], self.model_space.start['A'], delta=0.001)

    def test_add_result(self):
        self.pop.setup(self.size, self.model_space)
        i_sample = 1
        i_result = 0
        result = self.pop.pos[i_sample]
        self.pop.add_result(result, 0.0)
        self.assertEqual(self.pop.results.shape[0], 1)
        self.assertEqual(self.pop.results[i_result], result)
        self.assertEqual(self.pop.best[i_sample], result)

    def test_save_population(self):
        self.pop.setup(self.size, self.model_space)
        filename = os.path.join(self.test_dir, "test_save_population.pop")
        self.pop.save_population(filename)

    def test_save_results(self):
        self.pop.setup(self.size, self.model_space)
        i_sample = 1
        result = self.pop.pos[i_sample]
        self.pop.add_result(result, 1.0)
        filename = os.path.join(self.test_dir, "test_save_results.dat")
        self.pop.save_results(filename)

    def test_save_array(self):
        self.pop.setup(self.size, self.model_space)
        filename = os.path.join(self.test_dir, "test_save_array.pos")
        self.pop.save_array(filename, self.pop.pos)

    def test_load_array(self):
        n = 3
        filename = os.path.join(self.test_dir, "test_load_array")
        self.pop.setup(self.size, self.model_space)

        # expected array
        dt_exp = self.pop.get_pop_dtype(self.model_space.start)
        a_exp = np.zeros((n,), dtype=dt_exp)
        a_exp['A'] = np.linspace(0, 1, n)
        a_exp['B'] = np.linspace(1, 2, n)
        a_exp['C'] = np.linspace(3, 4, n)
        a_exp['_rfac'] = np.linspace(5, 6, n)
        a_exp['_gen'] = np.array([3, 4, 7])
        a_exp['_particle'] = np.array([1, 0, 2])
        a_exp['_model'] = np.array([3, 6, 1])

        # test array is a expected array with different column order
        dt_test = [('A', 'f4'), ('_particle', 'i4'), ('_rfac', 'f4'), ('C', 'f4'), ('_gen', 'i4'), ('B', 'f4'),
                   ('_model', 'i4')]
        names_test = [a[0] for a in dt_test]
        a_test = np.zeros((n,), dtype=dt_test)
        for name in names_test:
            a_test[name] = a_exp[name]
        header = " ".join(names_test)
        np.savetxt(filename, a_test, fmt='%g', header=header)

        result = np.zeros((n,), dtype=dt_exp)
        result = self.pop.load_array(filename, result)
        self.assert_pop_array_equal(result, a_exp)

    def test_constrain_position(self):
        # upper
        pos1 = 11.0
        vel1 = 5.0
        min1 = 0.0
        max1 = 10.0
        pos2, vel2, min2, max2 = population.Population.constrain_position(pos1, vel1, min1, max1, 're-enter')
        self.assertAlmostEqual(pos2, 1.0)
        self.assertAlmostEqual(vel2, vel1)
        self.assertAlmostEqual(min1, min2)
        self.assertAlmostEqual(max1, max2)
        pos2, vel2, min2, max2 = population.Population.constrain_position(pos1, vel1, min1, max1, 'bounce')
        self.assertAlmostEqual(pos2, 9.0)
        self.assertAlmostEqual(vel2, -vel1)
        self.assertAlmostEqual(min1, min2)
        self.assertAlmostEqual(max1, max2)
        pos2, vel2, min2, max2 = population.Population.constrain_position(pos1, vel1, min1, max1, 'scatter')
        self.assertGreaterEqual(pos2, 6.0)
        self.assertLessEqual(pos2, 10.0)
        self.assertGreaterEqual(vel2, 0.0)
        self.assertLessEqual(vel2, vel1)
        self.assertAlmostEqual(min1, min2)
        self.assertAlmostEqual(max1, max2)
        pos2, vel2, min2, max2 = population.Population.constrain_position(pos1, vel1, min1, max1, 'stick')
        self.assertAlmostEqual(pos2, max1)
        self.assertAlmostEqual(vel2, max1 - (pos1 - vel1))
        self.assertAlmostEqual(min1, min2)
        self.assertAlmostEqual(max1, max2)
        pos2, vel2, min2, max2 = population.Population.constrain_position(pos1, vel1, min1, max1, 'random')
        self.assertGreaterEqual(pos2, 0.0)
        self.assertLessEqual(pos2, 10.0)
        self.assertGreaterEqual(vel2, -(max1 - min1))
        self.assertLessEqual(vel2, (max1 - min1))
        self.assertAlmostEqual(min1, min2)
        self.assertAlmostEqual(max1, max2)
        pos2, vel2, min2, max2 = population.Population.constrain_position(pos1, vel1, min1, max1, 'expand')
        self.assertAlmostEqual(pos2, pos1)
        self.assertAlmostEqual(vel2, vel1)
        self.assertAlmostEqual(min1, min2)
        self.assertAlmostEqual(max2, pos1)
        self.assertRaises(ValueError, population.Population.constrain_position, pos1, vel1, min1, max1, 'error')

        # lower
        pos1 = -1.0
        vel1 = -5.0
        min1 = 0.0
        max1 = 10.0
        pos2, vel2, min2, max2 = population.Population.constrain_position(pos1, vel1, min1, max1, 're-enter')
        self.assertAlmostEqual(pos2, 9.0)
        self.assertAlmostEqual(vel2, vel1)
        self.assertAlmostEqual(min1, min2)
        self.assertAlmostEqual(max1, max2)
        pos2, vel2, min2, max2 = population.Population.constrain_position(pos1, vel1, min1, max1, 'bounce')
        self.assertAlmostEqual(pos2, 1.0)
        self.assertAlmostEqual(vel2, -vel1)
        self.assertAlmostEqual(min1, min2)
        self.assertAlmostEqual(max1, max2)
        pos2, vel2, min2, max2 = population.Population.constrain_position(pos1, vel1, min1, max1, 'scatter')
        self.assertGreaterEqual(pos2, 0.0)
        self.assertLessEqual(pos2, 4.0)
        self.assertGreaterEqual(vel2, vel1)
        self.assertLessEqual(vel2, 0.0)
        self.assertAlmostEqual(min1, min2)
        self.assertAlmostEqual(max1, max2)
        pos2, vel2, min2, max2 = population.Population.constrain_position(pos1, vel1, min1, max1, 'stick')
        self.assertAlmostEqual(pos2, min1)
        self.assertAlmostEqual(vel2, min1 - (pos1 - vel1))
        self.assertAlmostEqual(min1, min2)
        self.assertAlmostEqual(max1, max2)
        pos2, vel2, min2, max2 = population.Population.constrain_position(pos1, vel1, min1, max1, 'random')
        self.assertGreaterEqual(pos2, 0.0)
        self.assertLessEqual(pos2, 10.0)
        self.assertGreaterEqual(vel2, -(max1 - min1))
        self.assertLessEqual(vel2, max1 - min1)
        self.assertAlmostEqual(min1, min2)
        self.assertAlmostEqual(max1, max2)
        pos2, vel2, min2, max2 = population.Population.constrain_position(pos1, vel1, min1, max1, 'expand')
        self.assertAlmostEqual(pos2, pos1)
        self.assertAlmostEqual(vel2, vel1)
        self.assertAlmostEqual(min2, pos1)
        self.assertAlmostEqual(max2, max1)
        self.assertRaises(ValueError, population.Population.constrain_position, pos1, vel1, min1, max1, 'error')

    def test_patch_from_file(self):
        self.pop.setup(self.size, self.model_space)

        data_dir = os.path.dirname(os.path.abspath(__file__))
        data_file = os.path.join(data_dir, "test_swarm.setup_with_results.1.dat")
        patch_size = 4
        expected_pos = {'A': [1.1, 1.2, 1.3, 1.4], 'B': [2.1, 2.2, 2.3, 2.4], 'C': [3.1, 3.2, 4.3, 3.4]}
        self.pop.patch_from_file(data_file)
        self.assertIsInstance(self.pop.pos_patch, np.ndarray)
        self.assertEqual(self.pop.pos_patch.shape, (patch_size,))
        np.testing.assert_array_equal(self.pop.pos_patch['_particle'], [1, 2, 3, 4])
        np.testing.assert_array_equal(self.pop.pos_patch['_gen'], [1, 2, 3, 4])
        np.testing.assert_array_almost_equal(self.pop.pos_patch['A'], expected_pos['A'])
        np.testing.assert_array_almost_equal(self.pop.pos_patch['B'], expected_pos['B'])
        np.testing.assert_array_almost_equal(self.pop.pos_patch['C'], expected_pos['C'])

    def test_apply_patch(self):
        self.pop.setup(self.size, self.model_space)
        expected_pos = self.pop.pos.copy()
        dt_test = [('A', 'f4'), ('_particle', 'i4'), ('_rfac', 'f4'), ('C', 'f4'), ('_model', 'i4')]
        patch_size = 3
        self.pop.pos_patch = np.zeros(patch_size, dtype=dt_test)
        self.pop.pos_patch['_particle'] = np.array([1, 2, 3])
        self.pop.pos_patch['_rfac'] = np.random.random_sample(patch_size)
        self.pop.pos_patch['_model'] = np.arange(patch_size)
        self.pop.pos_patch['A'] = np.array([1.1, 2.2, 3.3])
        self.pop.pos_patch['C'] = np.array([1.1, 2.2, 3.3])

        expected_pos['A'][1] = 1.1
        expected_pos['C'][3] = 3.3
        expected_pos['_particle'] = np.arange(POP_SIZE)
        expected_pos['_gen'] = 0
        expected_pos['_model'] = np.arange(POP_SIZE)
        expected_pos['_rfac'] = 2.1

        self.pop._apply_patch()
        self.assertIsNone(self.pop.pos_patch)
        self.assertEqual(self.pop.generation, 0)
        self.assertEqual(self.pop.model_count, POP_SIZE)
        self.assert_pop_array_equal(self.pop.pos, expected_pos)

    def test_find_result(self):
        self.model_space.min['A'] = -0.1
        self.model_space.max['A'] = 0.1
        self.model_space.min['B'] = 0.0
        self.model_space.max['B'] = 1000.
        self.model_space.min['C'] = 9.
        self.model_space.max['C'] = 9.001
        self.size = 100
        self.pop.setup(self.size, self.model_space)
        self.pop.results = self.pop.pos.copy()

        expected_index = 77
        names = ['A', 'B', 'C']
        target0 = {key: self.pop.results[key][expected_index] for key in names}
        target1 = self.pop.results[names].copy()
        target1 = target1[expected_index]
        target2 = {key: target0[key] for key in ['C', 'B']}
        target3 = target2.copy()
        target3['C'] += 9.e-6
        target4 = {'A': random.random(), 'B': random.random(), 'C': 9.002}

        result0 = self.pop.find_model(target0)
        self.assertEqual(result0, expected_index)
        result1 = self.pop.find_model(target1)
        self.assertEqual(result1, expected_index)
        result2 = self.pop.find_model(target2)
        self.assertEqual(result2, expected_index)
        result3 = self.pop.find_model(target3)
        self.assertEqual(result3, expected_index)
        self.assertRaises(ValueError, self.pop.find_model, target4)

    def test_import_positions(self):
        """
        test the population.Population.import_positions method.
        check the different type conversions.
        the main work is in test_import_positions_array.
        """
        self.pop.setup(self.size, self.model_space)

        source_type = [('A', 'f4'), ('B', 'f4'), ('C', 'f4'), ('D', 'f4'),
                       ('_model', 'i4'), ('_particle', 'i4'), ('_gen', 'i4'), ('_rfac', 'f4')]
        source = np.array([(1.5, 3.5, 3.5, 4.5, 51, 52, 53, 0.5),
                           (1.6, 2.6, 2.6, 4.6, 61, 62, 63, 0.6),
                           (1.7, 2.7, 3.7, 4.7, 71, 72, 73, 0.7),
                           (1.8, 2.8, 3.8, 4.8, 81, 82, 83, 0.8)], dtype=source_type)
        expected_type = [('_gen', 'i8'), ('_model', 'i8'), ('_particle', 'i8'), ('_rfac', 'f8'),
                         ('A', 'f8'), ('B', 'f8'), ('C', 'f8')]
        expected = np.array([(0, 0, 0, 0.0, 1.5, 3.5, 3.5),
                             (0, 0, 0, 0.0, 1.6, 2.6, 2.6),
                             (0, 0, 0, 0.0, 1.7, 2.7, 3.7),
                             (0, 0, 0, 0.0, 1.8, 2.8, 3.8)], dtype=expected_type)
        expected_single = np.array([expected[0]], dtype=expected_type)

        # None
        self.pop.clear_import()
        source_none = None
        retval = self.pop.import_positions(source_none)
        self.assertEqual(retval, 0)

        # np.ndarray
        self.pop.clear_import()
        source_ndarray = source
        retval = self.pop.import_positions(source_ndarray)
        self.assertEqual(retval, expected.shape[0])
        self.assert_pop_array_equal(self.pop.pos_import, expected, msg="numpy.ndarray")

        # np.void
        self.pop.clear_import()
        source_void = source[0]
        retval = self.pop.import_positions(source_void)
        self.assertEqual(retval, 1)
        self.assert_pop_array_equal(self.pop.pos_import, expected_single, msg="numpy.void")

        # dict
        self.pop.clear_import()
        source_dict = {k: source[0][k] for k in source.dtype.names}
        retval = self.pop.import_positions(source_dict)
        self.assertEqual(retval, 1)
        self.assert_pop_array_equal(self.pop.pos_import, expected_single, msg="dict")

        # sequence
        self.pop.clear_import()
        source_sequence = [row for row in source]
        retval = self.pop.import_positions(source_sequence)
        self.assertEqual(retval, expected.shape[0])
        self.assert_pop_array_equal(self.pop.pos_import, expected, msg="sequence")

        # file
        filename = os.path.join(self.test_dir, "test_import_positions")
        self.pop.save_array(filename, source)
        self.pop.clear_import()
        retval = self.pop.import_positions(filename)
        self.assertEqual(retval, expected.shape[0])
        self.assert_pop_array_equal(self.pop.pos_import, expected, msg="file")

        # callable
        def source_callable():
            return expected.copy()

        self.pop.clear_import()
        retval = self.pop.import_positions(source_callable)
        self.assertEqual(retval, expected.shape[0])
        self.assert_pop_array_equal(self.pop.pos_import, expected, msg="callable")

    def test_import_positions_array(self):
        """
        test the population.Population.import_positions_array method.
        - import a data array, check type and size.
        - add to previous import.
        - ignore control fields.
        - no range or duplicate checking.
        - missing parameter.
        """
        self.pop.setup(self.size, self.model_space)
        source_type = [('A', 'f4'), ('B', 'f4'), ('C', 'f4'), ('D', 'f4'),
                       ('_model', 'i4'), ('_particle', 'i4'), ('_gen', 'i4'), ('_rfac', 'f4')]
        source = np.array([(1.0, 0.0, 0.0, 0.0, 0, 0, 0, 0.0),
                           (1.5, 3.5, 3.5, 4.5, 51, 52, 53, 0.5),
                           (1.6, 2.6, 2.6, 4.6, 61, 62, 63, 0.6),
                           (1.7, 2.7, 3.7, 4.7, 71, 72, 73, 0.7),
                           (1.8, 2.8, 3.8, 4.8, 81, 82, 83, 0.8)], dtype=source_type)
        self.pop.pos_import.resize(1)
        self.pop.pos_import[0]['A'] = 1.0
        expected_type = [('_gen', 'i8'), ('_model', 'i8'), ('_particle', 'i8'), ('_rfac', 'f8'),
                         ('A', 'f8'), ('B', 'f8'), ('C', 'f8')]
        expected = np.array([(0, 0, 0, 0.0, 1.0, 0.0, 0.0),
                             (0, 0, 0, 0.0, 1.0, 0.0, 0.0),
                             (0, 0, 0, 0.0, 1.5, 3.5, 3.5),
                             (0, 0, 0, 0.0, 1.6, 2.6, 2.6),
                             (0, 0, 0, 0.0, 1.7, 2.7, 3.7),
                             (0, 0, 0, 0.0, 1.8, 2.8, 3.8)], dtype=expected_type)
        retval = self.pop.import_positions_array(source)
        self.assertEqual(retval, source.shape[0])
        self.assertEqual(self.pop.pos_import.dtype.names, self.pop.pos.dtype.names)
        self.assert_pop_array_equal(self.pop.pos_import, expected, msg="imported array")

        source_type = [('B', 'f4'), ('C', 'f4'), ('D', 'f4'), ('E', 'f4'),
                       ('_model', 'i4'), ('_particle', 'i4'), ('_gen', 'i4'), ('_rfac', 'f4')]
        source = np.array([(1.0, 0.0, 0.0, 0.0, 0, 0, 0, 0.0),
                           (1.8, 2.8, 3.8, 4.8, 81, 82, 83, 0.8)], dtype=source_type)
        self.assertRaises(ValueError, self.pop.import_positions_array, source)

    def test_advance_from_import(self):
        """
        test the population.Population.advance_from_import method.
        - array type and size.
        - range checks.
        - de-duplication.
        """
        self.pop.setup(self.size, self.model_space)
        self.pop.position_constrain_mode = 'error'
        source_type = [('A', 'f8'), ('B', 'f8'), ('C', 'f8'),
                       ('_model', 'i8'), ('_particle', 'i8'), ('_gen', 'i8'), ('_rfac', 'f8')]
        source = np.array([(1.3, 2.3, 3.3, 0, 0, 0, 0.0),
                           (1.3, 2.3, 3.3, 0, 0, 0, 0.0),
                           (1.4, 2.4, 3.4, 0, 0, 0, 0.0),
                           (1.5, 3.5, 3.5, 0, 0, 0, 0.0),
                           (1.6, 2.6, 2.6, 0, 0, 0, 0.0),
                           (1.7, 2.7, 3.7, 0, 0, 0, 0.0),
                           (1.8, 2.8, 3.8, 0, 0, 0, 0.0)], dtype=source_type)
        source = self.reorder_pop_array(source)
        results = np.array([(1.4, 2.4, 3.4, 0, 0, 0, 0.4)], dtype=source_type)
        results = self.reorder_pop_array(results)
        expected = np.array([(1.3, 2.3, 3.3, POP_SIZE + 0, 0, 0, 2.1),
                             (1.7, 2.7, 3.7, POP_SIZE + 1, 1, 0, 2.1),
                             (1.8, 2.8, 3.8, POP_SIZE + 2, 2, 0, 2.1)], dtype=source_type)
        expected = self.reorder_pop_array(expected)

        self.pop.pos_import = source
        self.pop.results = results
        retval = self.pop.advance_from_import()
        self.assertEqual(retval, expected.shape[0])
        self.assertEqual(self.pop.size_act, expected.shape[0])
        pos_act = self.pop.pos[0:self.pop.size_act]
        self.assert_pop_array_equal(pos_act, expected)


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