# Filename: test_mc.py
# pylint: disable=C0111,R0904,C0103
# vim:set ts=4 sts=4 sw=4 et:
"""
Tests for MC Modules.
"""

import numpy as np
from numpy.testing import assert_array_equal, assert_allclose
import pytest

from km3pipe import Table, Blob, Pipeline, Module
from km3pipe.testing import TestCase

from km3modules.mc import (
    convert_mc_times_to_jte_times,
    MCTimeCorrector,
    GlobalRandomState,
)

__author__ = "Moritz Lotze, Michael Moser"
__copyright__ = "Copyright 2018, Tamas Gal and the KM3NeT collaboration."
__license__ = "MIT"
__maintainer__ = "Tamas Gal, Moritz Lotze"
__email__ = "tgal@km3net.de"
__status__ = "Development"


class TestMCConvert(TestCase):
    def setUp(self):
        self.event_info = Table(
            {"timestamp": 1, "nanoseconds": 700000000, "mc_time": 1.74999978e9,}
        )

        self.mc_tracks = Table({"time": 1,})

        self.mc_hits = Table({"time": 30.79,})

        self.blob = Blob(
            {
                "event_info": self.event_info,
                "mc_hits": self.mc_hits,
                "mc_tracks": self.mc_tracks,
            }
        )

    def test_convert_mc_times_to_jte_times(self):
        times_mc_tracks = convert_mc_times_to_jte_times(
            self.mc_tracks.time,
            self.event_info.timestamp * 1e9 + self.event_info.nanoseconds,
            self.event_info.mc_time,
        )
        times_mc_hits = convert_mc_times_to_jte_times(
            self.mc_hits.time,
            self.event_info.timestamp * 1e9 + self.event_info.nanoseconds,
            self.event_info.mc_time,
        )

        assert times_mc_tracks is not None
        assert times_mc_hits is not None
        print(times_mc_tracks, times_mc_hits)
        assert np.allclose(times_mc_tracks, 49999781)
        assert np.allclose(times_mc_hits, 49999810.79)

    def test_process(self):
        corr = MCTimeCorrector(
            mc_hits_key="mc_hits",
            mc_tracks_key="mc_tracks",
            event_info_key="event_info",
        )
        newblob = corr.process(self.blob)
        assert newblob["mc_hits"] is not None
        assert newblob["mc_tracks"] is not None
        assert np.allclose(newblob["mc_hits"].time, 49999810.79)
        assert np.allclose(newblob["mc_tracks"].time, 49999781)


class TestGlobalRandomState(TestCase):
    def test_default_random_state(self):
        assertAlmostEqual = self.assertAlmostEqual

        class Observer(Module):
            def configure(self):
                self.i = 0
                self.x = [0.3745401188, 0.950714306, 0.7319939418]

            def process(self, blob):
                assertAlmostEqual(self.x[self.i], np.random.rand())
                self.i += 1
                return blob

        pipe = Pipeline()
        pipe.attach(GlobalRandomState)
        pipe.attach(Observer)
        pipe.drain(3)

    def test_custom_random_state(self):
        assertAlmostEqual = self.assertAlmostEqual

        class Observer(Module):
            def configure(self):
                self.i = 0
                self.x = [0.221993171, 0.870732306, 0.206719155]

            def process(self, blob):
                assertAlmostEqual(self.x[self.i], np.random.rand())
                self.i += 1
                return blob

        pipe = Pipeline()
        pipe.attach(GlobalRandomState, seed=5)
        pipe.attach(Observer)
        pipe.drain(3)

    def test_without_pipeline_and_default_state(self):
        GlobalRandomState()
        numbers = np.arange(1, 50)
        np.random.shuffle(numbers)
        lotto_numbers = sorted(numbers[:6])
        self.assertListEqual([14, 18, 28, 45, 46, 48], lotto_numbers)

    def test_without_pipeline_with_custom_seed(self):
        GlobalRandomState(seed=23)
        numbers = np.arange(1, 50)
        np.random.shuffle(numbers)
        lotto_numbers = sorted(numbers[:6])
        self.assertListEqual([14, 15, 18, 19, 33, 44], lotto_numbers)
