"""
@package tests.test_scan
unit tests for pmsco.scan.

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-21 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
from pathlib import Path
import unittest

from pmsco.data import holo_grid
from pmsco.scan import Scan, ScanKey, ScanLoader, ScanCreator


class TestScanCreator(unittest.TestCase):
    """
    test case for @ref pmsco.project.ScanCreator class

    """
    def test_load_1(self):
        """
        test the load method, case 1

        test for:
        - correct array expansion of an ['e', 'a'] scan.
        - correct file name expansion with place holders and pathlib.Path objects.
        """
        sc = ScanCreator()
        sc.filename = Path("${test_p}", "twoatom_energy_alpha.etpai")
        sc.positions = {
            "e": "np.arange(10, 400, 5)",
            "t": "0",
            "p": "0",
            "a": "np.linspace(-30, 30, 31)"
        }
        sc.emitter = "Cu"
        sc.initial_state = "2p3/2"

        p = Path(__file__).parent.parent / "pmsco" / "projects" / "twoatom"
        dirs = {"test_p": p,
                "test_s": str(p)}

        result = sc.load(dirs=dirs)

        self.assertEqual(result.mode, ['e', 'a'])
        self.assertEqual(result.emitter, sc.emitter)
        self.assertEqual(result.initial_state, sc.initial_state)

        e = np.arange(10, 400, 5)
        a = np.linspace(-30, 30, 31)
        t = p = np.asarray([0])
        np.testing.assert_array_equal(result.energies, e)
        np.testing.assert_array_equal(result.thetas, t)
        np.testing.assert_array_equal(result.phis, p)
        np.testing.assert_array_equal(result.alphas, a)

        self.assertTrue(Path(result.filename).is_file(), msg=f"file {result.filename} not found")


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, "..", "pmsco", "projects", "twoatom", "twoatom_energy_alpha.etpai")

        scan = 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 = 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)

    def test_generate_holo_scan(self):
        scan = Scan()
        scan.generate_holo_scan(generator=holo_grid,
                                generator_args={},
                                other_positions={"e": 250, "a": 5},
                                emitter="C", initial_state="1s")

        self.assertEqual(scan.thetas.shape, (16376,))
        self.assertEqual(scan.phis.shape, (16376,))
        np.testing.assert_array_almost_equal(scan.alphas, np.asarray((5,)))
        np.testing.assert_array_almost_equal(scan.energies, np.asarray((250,)))


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