"""
@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-24 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 unittest

import pmsco.data as md
from pmsco.transforms.multipoles import MultipoleExpansion


class TestMultipoleExpansion(unittest.TestCase):
    def setUp(self):
        # before each test method
        pass

    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_generate_expand(self):
        orig_data = md.holo_array(md.holo_grid, {"theta_step": 2.}, datatype="ITP")
        orig_data["i"] = np.cos(np.deg2rad(orig_data["t"])) * np.sin(np.deg2rad(orig_data["p"]) * 3) ** 2

        lmax = 12
        me = MultipoleExpansion()
        me.holoscan = orig_data
        me.lmax = lmax
        alm = me.generate()

        self.assertEqual((lmax / 2 + 1, lmax * 2 + 1), alm.shape)
        self.assertTrue(np.any(np.real(alm) > 0.), "real part non-zero?")
        self.assertTrue(np.any(np.imag(alm) > 0.), "imaginary part non-zero?")

        me.expand()
        expanded_data = me.expansion

        self.assertEqual(expanded_data.shape, orig_data.shape)
        self.assertTrue(np.any(np.abs(expanded_data['i']) > 0.), "output array non-zero?")

        rf = md.square_diff_rfactor(expanded_data, orig_data)
        print(rf)
        self.assertLess(rf, 0.25)


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