"""
@package tests.test_project
unit tests for pmsco.project.

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 and mock must be installed.

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

@copyright (c) 2015-18 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 mock
import numpy as np
import os
import unittest

import pmsco.data as data
import pmsco.dispatch as dispatch
import pmsco.project as project


class TestScan(unittest.TestCase):
    """
    test case for @ref pmsco.project.Scan class

    """
    def test_import_scan_file(self):
        base_dir = os.path.dirname(os.path.abspath(__file__))
        test_file = os.path.join(base_dir, "..", "projects", "twoatom", "twoatom_energy_alpha.etpai")

        scan = project.Scan()
        scan.import_scan_file(test_file, "C", "1s")

        mode = ['e', 'a']
        self.assertEqual(scan.mode, mode)

        ae = np.arange(10, 1005, 5)
        at = np.asarray([0])
        ap = np.asarray([0])
        aa = np.arange(-90, 91, 1)

        np.testing.assert_array_almost_equal(scan.energies, ae)
        np.testing.assert_array_almost_equal(scan.thetas, at)
        np.testing.assert_array_almost_equal(scan.phis, ap)
        np.testing.assert_array_almost_equal(scan.alphas, aa)

    def test_define_scan(self):
        scan = project.Scan()
        p0 = np.asarray([20])
        p1 = np.linspace(1, 4, 4)
        p2 = np.linspace(11, 13, 3)
        d = {'t': p1, 'e': p0, 'p': p2}
        scan.define_scan(d, "C", "1s")

        ae = np.asarray([20])
        at = np.asarray([1, 2, 3, 4])
        ap = np.asarray([11, 12, 13])
        aa = np.asarray([0])

        np.testing.assert_array_almost_equal(scan.energies, ae)
        np.testing.assert_array_almost_equal(scan.thetas, at)
        np.testing.assert_array_almost_equal(scan.phis, ap)
        np.testing.assert_array_almost_equal(scan.alphas, aa)

        re = np.ones(12) * 20
        rt = np.asarray([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4])
        rp = np.asarray([11, 12, 13, 11, 12, 13, 11, 12, 13, 11, 12, 13])
        ra = np.ones(12) * 0

        np.testing.assert_array_almost_equal(scan.raw_data['e'], re)
        np.testing.assert_array_almost_equal(scan.raw_data['t'], rt)
        np.testing.assert_array_almost_equal(scan.raw_data['p'], rp)
        np.testing.assert_array_almost_equal(scan.raw_data['a'], ra)


class TestProject(unittest.TestCase):
    def setUp(self):
        # before each test method
        self.project = project.Project()

    def tearDown(self):
        # after each test method
        pass

    @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

    @mock.patch('pmsco.data.load_data')
    @mock.patch('pmsco.data.save_data')
    def test_combine_domains(self, save_data_mock, load_data_mock):
        self.project.scans.append(project.Scan())

        parent_task = dispatch.CalculationTask()
        parent_task.change_id(model=0, scan=0)
        parent_task.model['wdom1'] = 0.5

        child_tasks = [parent_task.copy()] * 2
        for idx, task in enumerate(child_tasks):
            task.change_id(domain=idx)

        data1 = data.create_data(5, datatype='EI')
        data1['e'] = np.arange(5)
        data1['i'] = 10.
        data2 = data1.copy()
        data2['i'] = 10.
        load_data_mock.side_effect = [data1, data2]
        data3 = data1.copy()
        data3['i'] = (10. + 0.5 * 10.) / 1.5

        self.project.combine_domains(parent_task, child_tasks)

        save_data_mock.assert_called()
        args, kwargs = save_data_mock.call_args
        result_data = args[1]
        np.testing.assert_array_almost_equal(result_data['e'], data3['e'])
        np.testing.assert_array_almost_equal(result_data['i'], data3['i'])


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