''' Script to check the correctness of the analysis. The analysis is done on raw data and all results are compared to a recorded analysis.
'''
import os

import unittest

import numpy as np

from beam_telescope_analysis import track_analysis
from beam_telescope_analysis.tools import test_tools


class TestTrackAnalysis(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        # virtual X server for plots under headless LINUX travis testing is needed
        if os.getenv('TRAVIS', False) and os.getenv('TRAVIS_OS_NAME', False) == 'linux':
            from xvfbwrapper import Xvfb  # virtual X server for plots under headless LINUX travis testing is needed
            cls.vdisplay = Xvfb()
            cls.vdisplay.start()

    @classmethod
    def tearDownClass(cls):  # Remove created files
        pass
        # os.remove(os.path.join(cls.output_folder, 'Tracks_merged.pdf'))

    def test_kalman(self):
        # pixel size of sensor
        pixel_size = np.array([(18.5, 18.5), (18.5, 18.5), (18.5, 18.5), (18.5, 18.5), (18.5, 18.5), (18.5, 18.5), (250., 50.)])
        pixel_resolution = pixel_size / np.sqrt(12)
        material_budget = np.array([100., 100., 100., 100., 100., 100., 250.]) / np.array([125390., 125390., 125390., 125390., 125390., 125390., 93700.])
        prealignment = {'z': [0., 29900., 60300., 82100., 118700., 160700., 197800.]}
        kwargs = {'track_hits': np.array([[[-1229.22372954, 2828.19616302, 0., pixel_resolution[0][0], pixel_resolution[0][1], 0.],
                                           [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],  # [-1254.51224282, 2827.4291421, 29900.],
                                           [-1285.6117892, 2822.34536687, 60300., pixel_resolution[2][0], pixel_resolution[2][1], 0.],
                                           [-1311.31083616, 2823.56121414, 82100., pixel_resolution[3][0], pixel_resolution[3][1], 0.],
                                           [-1335.8529645, 2828.43359043, 118700., pixel_resolution[4][0], pixel_resolution[4][1], 0.],
                                           [-1357.81872222, 2840.86947964, 160700., pixel_resolution[5][0], pixel_resolution[5][1], 0.],
                                           [-1396.35698339, 2843.76799577, 197800., pixel_resolution[6][0], pixel_resolution[6][1], 0.]]]),
                  'dut_fit_selection': 61,
                  'z_positions': [[0., 29900, 60300, 82100, 118700, 160700, 197800]],
                  'alignment': [prealignment],
                  'use_prealignment': True,
                  'pixel_size': pixel_size,
                  'n_pixels': ((576, 1152), (576, 1152), (576, 1152), (576, 1152), (576, 1152), (576, 1152), (80, 336)),
                  'beam_energy': 2500.,
                  'material_budget': material_budget,
                  'add_scattering_plane': False}

        # expected result array: (state estimates, chi, x error, y errors)
        result = [[[-1.23045812e+03, 2.82684464e+03, 0.00000000e+00, -9.54188957e-04, -5.78722777e-05, 9.99999543e-01],
                   [-1.25900270e+03, 2.82511339e+03, 2.99000000e+04, -9.54667558e-04, -5.79013065e-05, 9.99999543e-01],
                   [-1.28705254e+03, 2.82443254e+03, 6.03000000e+04, -9.22691847e-04, -2.23966180e-05, 9.99999574e-01],
                   [-1.30575083e+03, 2.82550588e+03, 8.21000000e+04, -8.57719095e-04, 4.92360053e-05, 9.99999631e-01],
                   [-1.33339390e+03, 2.83014572e+03, 1.18700000e+05, -7.55274948e-04, 1.26771487e-04, 9.99999707e-01],
                   [-1.36192826e+03, 2.83782855e+03, 1.60700000e+05, -6.79389377e-04, 1.82924497e-04, 9.99999752e-01],
                   [-1.38713361e+03, 2.84461505e+03, 1.97800000e+05, -6.79389377e-04, 1.82924497e-04, 9.99999752e-01]],
                  [79.59176738400244],
                  [3.62429044, 3.2884327, 3.21655702, 3.1539946, 3.23671172, 4.66501707, 8.62909928],
                  [3.62429044, 3.2884327, 3.21655702, 3.1539946, 3.23671172, 4.66501707, 8.62909928]]

        for i in range(4):  # test each return (state estimates, chi, x error, y errors) seperatly
            test = test_tools._call_function_with_args(function=track_analysis._fit_tracks_kalman_loop,
                                                       **kwargs)[0][i]
            data_equal = np.allclose(test, result[i])
            self.assertTrue(data_equal)


if __name__ == '__main__':
    import logging
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - [%(levelname)-8s] (%(threadName)-10s) %(message)s")
    suite = unittest.TestLoader().loadTestsFromTestCase(TestTrackAnalysis)
    unittest.TextTestRunner(verbosity=2).run(suite)

[((72, 4, 72, 128), 'logging.basicConfig', 'logging.basicConfig', (), '', False, 'import logging\n'), ((30, 21, 30, 128), 'numpy.array', 'np.array', ({(30, 30, 30, 127): '[(18.5, 18.5), (18.5, 18.5), (18.5, 18.5), (18.5, 18.5), (18.5, 18.5), (\n    18.5, 18.5), (250.0, 50.0)]'}, {}), '([(18.5, 18.5), (18.5, 18.5), (18.5, 18.5), (18.5, 18.5), (18.5, \n    18.5), (18.5, 18.5), (250.0, 50.0)])', True, 'import numpy as np\n'), ((18, 11, 18, 37), 'os.getenv', 'os.getenv', ({(18, 21, 18, 29): '"""TRAVIS"""', (18, 31, 18, 36): '(False)'}, {}), "('TRAVIS', False)", False, 'import os\n'), ((20, 27, 20, 33), 'xvfbwrapper.Xvfb', 'Xvfb', ({}, {}), '()', False, 'from xvfbwrapper import Xvfb\n'), ((31, 40, 31, 51), 'numpy.sqrt', 'np.sqrt', ({(31, 48, 31, 50): '(12)'}, {}), '(12)', True, 'import numpy as np\n'), ((32, 26, 32, 78), 'numpy.array', 'np.array', ({(32, 35, 32, 77): '[100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 250.0]'}, {}), '([100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 250.0])', True, 'import numpy as np\n'), ((32, 81, 32, 153), 'numpy.array', 'np.array', ({(32, 90, 32, 152): '[125390.0, 125390.0, 125390.0, 125390.0, 125390.0, 125390.0, 93700.0]'}, {}), '([125390.0, 125390.0, 125390.0, 125390.0, 125390.0, 125390.0, 93700.0])', True, 'import numpy as np\n'), ((34, 32, 40, 138), 'numpy.array', 'np.array', ({(34, 41, 40, 137): '[[[-1229.22372954, 2828.19616302, 0.0, pixel_resolution[0][0],\n    pixel_resolution[0][1], 0.0], [np.nan, np.nan, np.nan, np.nan, np.nan,\n    np.nan], [-1285.6117892, 2822.34536687, 60300.0, pixel_resolution[2][0],\n    pixel_resolution[2][1], 0.0], [-1311.31083616, 2823.56121414, 82100.0,\n    pixel_resolution[3][0], pixel_resolution[3][1], 0.0], [-1335.8529645, \n    2828.43359043, 118700.0, pixel_resolution[4][0], pixel_resolution[4][1],\n    0.0], [-1357.81872222, 2840.86947964, 160700.0, pixel_resolution[5][0],\n    pixel_resolution[5][1], 0.0], [-1396.35698339, 2843.76799577, 197800.0,\n    pixel_resolution[6][0], pixel_resolution[6][1], 0.0]]]'}, {}), '([[[-1229.22372954, 2828.19616302, 0.0, pixel_resolution[0][0],\n    pixel_resolution[0][1], 0.0], [np.nan, np.nan, np.nan, np.nan, np.nan,\n    np.nan], [-1285.6117892, 2822.34536687, 60300.0, pixel_resolution[2][0],\n    pixel_resolution[2][1], 0.0], [-1311.31083616, 2823.56121414, 82100.0,\n    pixel_resolution[3][0], pixel_resolution[3][1], 0.0], [-1335.8529645, \n    2828.43359043, 118700.0, pixel_resolution[4][0], pixel_resolution[4][1],\n    0.0], [-1357.81872222, 2840.86947964, 160700.0, pixel_resolution[5][0],\n    pixel_resolution[5][1], 0.0], [-1396.35698339, 2843.76799577, 197800.0,\n    pixel_resolution[6][0], pixel_resolution[6][1], 0.0]]])', True, 'import numpy as np\n'), ((66, 25, 66, 53), 'numpy.allclose', 'np.allclose', ({(66, 37, 66, 41): 'test', (66, 43, 66, 52): 'result[i]'}, {}), '(test, result[i])', True, 'import numpy as np\n'), ((73, 12, 73, 33), 'unittest.TestLoader', 'unittest.TestLoader', ({}, {}), '()', False, 'import unittest\n'), ((74, 4, 74, 40), 'unittest.TextTestRunner', 'unittest.TextTestRunner', (), '', False, 'import unittest\n'), ((18, 42, 18, 76), 'os.getenv', 'os.getenv', ({(18, 52, 18, 68): '"""TRAVIS_OS_NAME"""', (18, 70, 18, 75): '(False)'}, {}), "('TRAVIS_OS_NAME', False)", False, 'import os\n'), ((64, 19, 65, 64), 'beam_telescope_analysis.tools.test_tools._call_function_with_args', 'test_tools._call_function_with_args', (), '', False, 'from beam_telescope_analysis.tools import test_tools\n')]