"""
@package tests.test_data
unit tests for pmsco.data

the purpose of these tests is to mainly to check the syntax, and correct data types,
i.e. anything that could cause a run-time error.
calculation results are sometimes checked for plausibility but not exact values,
depending on the level of debugging required for a specific part of 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 math
import numpy as np
from pathlib import Path
import unittest

import pmsco.data as md


class TestDataFunctions(unittest.TestCase):
    def setUp(self):
        # before each test method
        shape = (10, )
        self.e_scan = md.create_data(shape, dtype=md.DTYPE_EI)
        self.e_scan['e'] = np.linspace(100.0, 200.0, shape[0])
        self.e_scan['i'] = np.linspace(0.0, 10.0, shape[0])

        shape = (12, )
        self.ea_scan = md.create_data(shape, dtype=md.DTYPE_ETPAI)
        self.ea_scan[0] = (1.0, 0.0, 0.0, -1.0, 1.0)
        self.ea_scan[1] = (1.0, 0.0, 0.0, 0.0, 1.0)
        self.ea_scan[2] = (1.0, 0.0, 0.0, 1.0, 1.0)
        self.ea_scan[3] = (1.0, 0.0, 0.0, 2.0, 1.0)
        self.ea_scan[4] = (2.0, 0.0, 0.0, -1.0, 2.0)
        self.ea_scan[5] = (2.0, 0.0, 0.0, 0.0, 2.0)
        self.ea_scan[6] = (2.0, 0.0, 0.0, 1.0, 2.0)
        self.ea_scan[7] = (2.0, 0.0, 0.0, 2.0, 2.0)
        self.ea_scan[8] = (3.0, 0.0, 0.0, -1.0, 3.0)
        self.ea_scan[9] = (3.0, 0.0, 0.0, 0.0, 3.0)
        self.ea_scan[10] = (3.0, 0.0, 0.0, 1.0, 3.0)
        self.ea_scan[11] = (3.0, 0.0, 0.0, 2.0, 3.0)

        shape = (10, )
        self.holo_scan = md.create_data(shape, dtype=md.DTYPE_ETPI)
        self.holo_scan[0] = (1.0, 0.0, 0.0, 1.0)
        self.holo_scan[1] = (1.0, 1.0, 0.0, 2.0)
        self.holo_scan[2] = (1.0, 1.0, 120.0, 3.0)
        self.holo_scan[3] = (1.0, 1.0, 240.0, 4.0)
        self.holo_scan[4] = (1.0, 2.0, 0.0, 5.0)
        self.holo_scan[5] = (1.0, 2.0, 60.0, 6.0)
        self.holo_scan[6] = (1.0, 2.0, 120.0, 7.0)
        self.holo_scan[7] = (1.0, 2.0, 180.0, 8.0)
        self.holo_scan[8] = (1.0, 2.0, 240.0, 9.0)
        self.holo_scan[9] = (1.0, 2.0, 300.0, 10.0)

    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

    def test_create_data(self):
        shape = (10, )
        data = md.create_data(shape, dtype=md.DTYPE_ETPAIS)
        expected_names = ('e', 't', 'p', 'a', 'i', 's')
        self.assertEqual(data.dtype.names, expected_names)
        self.assertEqual(data.shape, shape)

    def test_detect_scan_mode_1d(self):
        scan_mode, scan_positions = md.detect_scan_mode(self.e_scan)
        
        expected_mode = ['e']
        expected_positions = {}
        expected_positions['e'] = np.linspace(100.0, 200.0,  10)
        
        self.assertEqual(scan_mode, expected_mode)
        for dim in expected_positions:
            np.testing.assert_almost_equal(scan_positions[dim], expected_positions[dim], decimal=3)
        
    def test_detect_scan_mode_2d(self):
        scan_mode, scan_positions = md.detect_scan_mode(self.ea_scan)
        
        expected_mode = ['e', 'a']
        expected_positions = {}
        expected_positions['e'] = np.asarray((1.0, 2.0, 3.0))
        expected_positions['t'] = np.zeros((1))
        expected_positions['p'] = np.zeros((1))
        expected_positions['a'] = np.asarray((-1.0, 0.0, 1.0, 2.0))
        
        self.assertEqual(scan_mode, expected_mode)
        for dim in expected_positions:
            np.testing.assert_almost_equal(scan_positions[dim], expected_positions[dim], decimal=3)
        
    def test_detect_scan_mode_holo(self):
        scan_mode, scan_positions = md.detect_scan_mode(self.holo_scan)

        expected_mode = ['t', 'p']
        expected_positions = {}
        expected_positions['e'] = np.ones((1))
        expected_positions['t'] = self.holo_scan['t']
        expected_positions['p'] = self.holo_scan['p']

        self.assertEqual(scan_mode, expected_mode)
        for dim in expected_positions:
            np.testing.assert_almost_equal(scan_positions[dim], expected_positions[dim], decimal=3)

    def test_detect_scan_mode_theta(self):
        scan = self.holo_scan
        scan['t'] = np.linspace(1.0, 2.0, scan.shape[0])
        scan['p'] = 3.3
        scan_mode, scan_positions = md.detect_scan_mode(scan)

        expected_mode = ['t']
        expected_positions = {}
        expected_positions['e'] = np.ones((1))
        expected_positions['t'] = np.linspace(1.0, 2.0, scan.shape[0])
        expected_positions['p'] = np.ones((1)) * 3.3

        self.assertEqual(scan_mode, expected_mode)
        for dim in expected_positions:
            np.testing.assert_almost_equal(scan_positions[dim], expected_positions[dim], decimal=3)

    def test_holo_array(self):
        args = {
            "theta_start": 90,
            "theta_step": 1,
            "theta_range": 90,
            "phi_start": 0,
            "phi_range": 360,
            "phi_refinement": 1
            }

        result = md.holo_array(generator=md.holo_grid, generator_args=args, datatype="ETPI")
        result['e'] = 250.
        md.sort_data(result)

        ref_path = Path(__file__).parent.parent / "pmsco" / "projects" / "twoatom" / "twoatom_hemi_250e.etpi"
        ref_array = md.load_data(ref_path)
        np.testing.assert_array_almost_equal(result['t'], ref_array['t'], decimal=1)
        np.testing.assert_array_almost_equal(result['p'], ref_array['p'], decimal=1)

    def test_analyse_holoscan_steps(self):
        args = {
            "theta_start": 90.0,
            "theta_step": 2.0,
            "theta_range": 90.0,
            "phi_start": 0.0,
            "phi_range": 120.0,
            "phi_refinement": 1.0
            }

        holo = md.holo_array(generator=md.holo_grid, generator_args=args, datatype="TPI")
        theta, dtheta, dphi = md.analyse_holoscan_steps(holo)

        expected_theta = np.arange(args["theta_start"] - args["theta_range"],
                                   args["theta_start"] + args["theta_step"],
                                   args["theta_step"])
        expected_dtheta = np.ones_like(expected_theta) * args["theta_step"]

        np.testing.assert_almost_equal(theta, expected_theta)
        np.testing.assert_almost_equal(dtheta, expected_dtheta)
        self.assertEqual(expected_theta.shape, dphi.shape)
        self.assertEqual(args["phi_range"], dphi[0])
        self.assertEqual(args["theta_step"], dphi[-1])
        np.testing.assert_array_less(np.ones_like(expected_theta) * args["theta_step"] * 0.999, dphi)

    def test_calc_modfunc_mean_1d(self):
        modf = md.calc_modfunc_mean(self.e_scan)

        exp_modf = self.e_scan.copy()
        exp_modf['i'] = (self.e_scan['i'] - 5.0) / 5.0

        np.testing.assert_allclose(modf['e'], exp_modf['e'])
        np.testing.assert_allclose(modf['i'], exp_modf['i'])

    def test_calc_modfunc_mean_2d(self):
        modf = md.calc_modfunc_mean(self.ea_scan)

        exp_modf = self.ea_scan.copy()
        exp_modf['i'][0] = -0.5
        exp_modf['i'][1] = -0.5
        exp_modf['i'][2] = -0.5
        exp_modf['i'][3] = -0.5
        exp_modf['i'][4] = 0.0
        exp_modf['i'][5] = 0.0
        exp_modf['i'][6] = 0.0
        exp_modf['i'][7] = 0.0
        exp_modf['i'][8] = 0.5
        exp_modf['i'][9] = 0.5
        exp_modf['i'][10] = 0.5
        exp_modf['i'][11] = 0.5

        np.testing.assert_allclose(modf['e'], exp_modf['e'])
        np.testing.assert_allclose(modf['t'], exp_modf['t'])
        np.testing.assert_allclose(modf['p'], exp_modf['p'])
        np.testing.assert_allclose(modf['a'], exp_modf['a'])
        np.testing.assert_allclose(modf['i'], exp_modf['i'])

    def test_calc_modfunc_loess_1d(self):
        """
        check that the result of msc_data.calc_modfunc_loess() is between -1 and 1.
        """

        # loess package not available
        if md.loess is None:
            return

        modf = md.calc_modfunc_loess(self.e_scan)
        self.assertEqual(self.e_scan.shape, modf.shape)
        exp_modf = self.e_scan.copy()
        np.testing.assert_allclose(modf['e'], exp_modf['e'])
        exp_modf['i'] = -1.000001
        np.testing.assert_array_less(exp_modf['i'],  modf['i'])
        exp_modf['i'] = +1.000001
        np.testing.assert_array_less(modf['i'], exp_modf['i'])

    @unittest.skip("test_calc_modfunc_loess_1d_nan fails and kills the whole test process")
    def test_calc_modfunc_loess_1d_nan(self):
        """
        check that data.calc_modfunc_loess() ignores NaNs gracefully.
        """

        # loess package not available
        if md.loess is None:
            return

        modified_index = 2
        self.e_scan['i'][modified_index] = np.nan
        modf = md.calc_modfunc_loess(self.e_scan)
        exp_modf = self.e_scan.copy()
        self.assertEqual(self.e_scan.shape, modf.shape)
        np.testing.assert_allclose(modf['e'], exp_modf['e'])
        self.assertTrue(math.isnan(modf['i'][modified_index]))
        modf['i'][modified_index] = 0.0
        exp_modf['i'] = -1.000001
        np.testing.assert_array_less(exp_modf['i'], modf['i'])
        exp_modf['i'] = +1.000001
        np.testing.assert_array_less(modf['i'], exp_modf['i'])

    def test_calc_modfunc_loess_2d(self):
        """
        check that the msc_data.calc_modfunc_loess() function does approximately what we want for a two-dimensional dataset.
        """

        # loess package not available
        if md.loess is None:
            return

        n_e = 10
        n_a = 15
        shape = (n_e * n_a, )
        scan = md.create_data(shape, dtype=md.DTYPE_ETPAI)

        e_range = np.linspace(100.0, 200.0, n_e)
        a_range = np.linspace(-15.0, 15.0, n_a)
        a_grid, e_grid = np.meshgrid(a_range, e_range)
        scan['e'] = np.ravel(e_grid)
        scan['a'] = np.ravel(a_grid)
        scan['t'] = 0.0
        scan['p'] = 90.0
        scan['i'] = 0.02 * scan['e'] + 0.03 * scan['a'] + np.sin((scan['e'] - 150) / 50 * math.pi) + np.sin(scan['a'] / 180 * math.pi)
        modf = md.calc_modfunc_loess(scan)
        self.assertEqual(scan.shape, modf.shape)
        exp_modf = scan.copy()

        np.testing.assert_allclose(modf['e'], exp_modf['e'])
        np.testing.assert_allclose(modf['t'], exp_modf['t'])
        np.testing.assert_allclose(modf['p'], exp_modf['p'])
        np.testing.assert_allclose(modf['a'], exp_modf['a'])
        
        exp_modf['i'] = -1.000001
        np.testing.assert_array_less(exp_modf['i'],  modf['i'])
        exp_modf['i'] = +1.000001
        np.testing.assert_array_less(modf['i'], exp_modf['i'])
        # this is rough estimate of the result, manually optimized by trial and error in Igor.
        # the R factor should be sensitive enough to detect mixed-up axes.
        exp_modf['i'] = 0.03 * np.sin((scan['e'] - 150) / 50 * math.pi)
        rf = md.square_diff_rfactor(modf, exp_modf)
        self.assertLessEqual(rf, 0.50)

    def test_alpha_mirror_average(self):
        n_e = 10
        n_a = 15
        shape = (n_e * n_a, )
        scan = md.create_data(shape, dtype=md.DTYPE_ETPAI)

        e_range = np.linspace(100.0, 200.0, n_e)
        a_range = np.linspace(-15.0, 15.0, n_a)
        a_grid, e_grid = np.meshgrid(a_range, e_range)
        scan['e'] = np.ravel(e_grid)
        scan['a'] = np.ravel(a_grid)
        scan['t'] = 0.0
        scan['p'] = 90.0
        scan['i'] = 0.02 * scan['e'] + 0.03 * scan['a']
        act_result = md.alpha_mirror_average(scan)
        exp_result = scan.copy()
        exp_result['i'] = 0.02 * scan['e']

        np.testing.assert_allclose(act_result['e'], exp_result['e'])
        np.testing.assert_allclose(act_result['t'], exp_result['t'])
        np.testing.assert_allclose(act_result['p'], exp_result['p'])
        np.testing.assert_allclose(act_result['a'], exp_result['a'])
        np.testing.assert_allclose(act_result['i'], exp_result['i'])

    def test_alpha_average(self):
        data = md.create_data((20), dtype=md.DTYPE_ETPAIS)
        data['e'][0:10] = 11
        data['e'][10:20] = 22
        data['a'][0:10] = np.arange(10)
        data['a'][10:20] = np.arange(10)
        data['i'] = np.arange(20)
        data['s'] = np.arange(20) / 10.0

        exp_result = md.create_data((2), dtype=md.DTYPE_ETPIS)
        exp_result['e'][0] = 11
        exp_result['e'][1] = 22
        exp_result['i'][0] = 4.5
        exp_result['i'][1] = 14.5
        exp_result['s'] = exp_result['i'] / 10.0

        act_result = md.alpha_average(data)

        np.testing.assert_allclose(act_result['e'], exp_result['e'])
        np.testing.assert_allclose(act_result['t'], exp_result['t'])
        np.testing.assert_allclose(act_result['p'], exp_result['p'])
        np.testing.assert_allclose(act_result['i'], exp_result['i'])
        np.testing.assert_allclose(act_result['s'], exp_result['s'])

    def test_phi_average(self):
        data = self.holo_scan.copy()

        exp_result = md.create_data((3), dtype=[('e', 'f4'), ('t', 'f4'), ('i', 'f4')])
        exp_result['e'] = np.asarray([1, 1, 1])
        exp_result['t'] = np.asarray([0, 1, 2])
        exp_result['i'] = np.asarray([1, 3, 7.5])

        act_result = md.phi_average(data)

        np.testing.assert_allclose(act_result['e'], exp_result['e'])
        np.testing.assert_allclose(act_result['t'], exp_result['t'])
        np.testing.assert_allclose(act_result['i'], exp_result['i'])

    def test_filter_tp(self):
        data = md.create_data((17), dtype=md.DTYPE_ETPAIS)
        data['e'] = 100.0
        data['t'] = np.array([0.0, 0.0, 0.0, 0.0, 1.0, 1.0,  1.0, 1.0, 2.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 3.0, 3.0])
        data['p'] = np.array([0.0, 1.0, 2.0, 3.0, 0.0, 1.01, 2.0, 3.0, 0.0, 1.0, 2.0, 3.0, 0.0, 1.0, 2.0, 3.0, 3.1])
        data['a'] = 0.0
        data['i'] = data['t'] * 10.0 + data['p']
        data['s'] = np.sqrt(data['i'])
        # duplicate
        data['p'][10] = 0.0

        filter = md.create_data((10), dtype=md.DTYPE_ETPI)
        filter['e'] = 100.0
        filter['t'] = np.array([0.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 3.0])
        filter['p'] = np.array([0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 0.0, 1.0, 2.0, 3.0])
        filter['i'] = -99.999

        exp_result = md.create_data((10), dtype=md.DTYPE_ETPAIS)
        exp_result['e'] = 100.0
        exp_result['t'] = np.array([0.0, 1.0, 1.00, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 3.0])
        exp_result['p'] = np.array([0.0, 0.0, 1.01, 0.0, 1.0, 2.0, 0.0, 1.0, 2.0, 3.0])
        exp_result['a'] = 0.0
        exp_result['i'] = exp_result['t'] * 10.0 + exp_result['p']
        exp_result['s'] = np.sqrt(exp_result['i'])
        exp_result['p'][5] = 0.0

        act_result = md.filter_tp(data, filter)

        np.testing.assert_allclose(act_result['e'], exp_result['e'])
        np.testing.assert_allclose(act_result['t'], exp_result['t'])
        np.testing.assert_allclose(act_result['p'], exp_result['p'])
        np.testing.assert_allclose(act_result['a'], exp_result['a'])
        np.testing.assert_allclose(act_result['i'], exp_result['i'])
        np.testing.assert_allclose(act_result['s'], exp_result['s'])

    def test_interpolate_hemi_scan(self):
        n_t = 91
        n_p = 11
        shape = (n_t * n_p, )
        calc_data = md.create_data(shape, dtype=md.DTYPE_TPI)

        t_range = np.linspace(0.0,  90.0, n_t)
        p_range = np.linspace(0.0, 360.0, n_p)
        t_grid, p_grid = np.meshgrid(t_range, p_range)
        calc_data['t'] = np.ravel(t_grid)
        calc_data['p'] = np.ravel(p_grid)
        calc_data['i'] = 100.0 * calc_data['t'] + calc_data['p']

        scan_data = md.create_data((10), dtype=md.DTYPE_TPI)
        scan_data['t'] = np.array([0.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 3.0])
        scan_data['p'] = np.array([0.0, 0.0, 1.0, 0.5, 1.7, 2.8, 0.0, 90.0, 180.0, 360.0])
        scan_data['i'] = -99.999

        exp_result = scan_data.copy()
        exp_result['i'] = 100.0 * scan_data['t'] + scan_data['p']

        act_result = md.interpolate_hemi_scan(calc_data, scan_data)

        np.testing.assert_allclose(act_result['t'], exp_result['t'])
        np.testing.assert_allclose(act_result['p'], exp_result['p'])
        np.testing.assert_allclose(act_result['i'], exp_result['i'])

    def test_scaled_rfactor(self):
        n = 20
        calc_modf = md.create_etpi((n,), False)
        calc_modf['e'] = 0.0
        calc_modf['t'] = 0.0
        calc_modf['p'] = np.linspace(0, np.pi, n)
        calc_modf['i'] = np.sin(calc_modf['p'])

        exp_modf = calc_modf.copy()
        exp_modf['i'] = 0.6 * np.sin(exp_modf['p'] + np.pi / 100.0)

        weights = np.ones_like(exp_modf['i'])

        r = md.scaled_rfactor_func(1.4, exp_modf['i'], weights, calc_modf['i'])
        self.assertGreater(r, 0.0)
        self.assertLess(r, 0.05)

    def test_square_diff_rfactor(self):
        n = 20
        calc_modf = md.create_data((n,), dtype=md.DTYPE_ETPI)
        calc_modf['e'] = 0.0
        calc_modf['t'] = 0.0
        calc_modf['p'] = np.linspace(0, np.pi, n)
        calc_modf['i'] = np.sin(calc_modf['p'])

        exp_modf = md.create_data((n,), dtype=md.DTYPE_ETPIS)
        exp_modf['i'] = 0.6 * np.sin(exp_modf['p'] + np.pi / 100.0)
        exp_modf['s'] = np.sqrt(np.abs(exp_modf['i']))

        r1 = md.square_diff_rfactor(exp_modf, calc_modf)
        self.assertAlmostEqual(r1, 0.95, delta=0.02)

        # one nan should not make a big difference
        calc_modf['i'][3] = np.nan
        r2 = md.square_diff_rfactor(exp_modf, calc_modf)
        self.assertAlmostEqual(r1, r2, delta=0.02)

        # all values nan should raise an exception
        with self.assertRaises(ValueError):
            calc_modf['i'] = np.nan
            md.square_diff_rfactor(exp_modf, calc_modf)

    def test_optimize_rfactor(self):
        n = 20
        calc_modf = md.create_data((n,), dtype=md.DTYPE_ETPI)
        calc_modf['e'] = 0.0
        calc_modf['t'] = 0.0
        calc_modf['p'] = np.linspace(0, np.pi, n)
        calc_modf['i'] = np.sin(calc_modf['p'])

        exp_modf = md.create_data((n,), dtype=md.DTYPE_ETPIS)
        exp_modf['i'] = 0.6 * np.sin(exp_modf['p'] + np.pi / 100.0)
        exp_modf['s'] = np.sqrt(np.abs(exp_modf['i']))

        r1 = md.optimize_rfactor(exp_modf, calc_modf)
        self.assertAlmostEqual(r1, 0.55, delta=0.02)

        # one nan should not make a big difference
        calc_modf['i'][3] = np.nan
        r2 = md.optimize_rfactor(exp_modf, calc_modf)
        self.assertAlmostEqual(r1, r2, delta=0.02)

        # all values nan should raise an exception
        with self.assertRaises(ValueError):
            calc_modf['i'] = np.nan
            md.optimize_rfactor(exp_modf, calc_modf)


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