import argparse
import os
from unittest import TestCase
from unittest.mock import MagicMock
import tempfile
import pytest

import numpy as np

import westpa
from westpa.core.binning.assign import RectilinearBinMapper
from westpa.core.binning.binless_manager import BinlessSimManager
from westpa.core.binning.mab_manager import MABSimManager
from westpa.core.segment import Segment
from westpa.core.states import BasisState
from westpa.core.sim_manager import PropagationError


def dummy_callback_one(self):
    pass


class TestSimManager(TestCase):
    def setUp(self):
        parser = argparse.ArgumentParser()
        westpa.rc.add_args(parser)

        here = os.path.dirname(__file__)
        os.environ['WEST_SIM_ROOT'] = os.path.join(here, 'refs', 'odld')

        config_file_name = os.path.join(here, 'refs', 'odld', 'west.cfg')
        args = parser.parse_args(['-r={}'.format(config_file_name)])
        westpa.rc.process_args(args)
        self.sim_manager = westpa.rc.get_sim_manager()
        self.test_dir = tempfile.mkdtemp()
        self.hdf5 = os.path.join(self.test_dir, "west.h5")
        self.basis_states = [BasisState(label="label", probability=1.0)]
        self.segments = [self.segment(0.0, 1.5, weight=0.125) for _i in range(4)] + [
            self.segment(1.5, 0.5, weight=0.125) for _i in range(4)
        ]
        self.sim_manager.we_driver.new_iteration()
        self.sim_manager.we_driver.assign(self.segments)
        self.sim_manager.we_driver.construct_next()
        self.sim_manager.segments = {segment.seg_id: segment for segment in self.segments}
        self.sim_manager.incomplete_segments = self.sim_manager.segments
        self.sim_manager.current_iter_istates = self.sim_manager.segments
        self.sim_manager.completed_segments = self.sim_manager.segments
        self.sim_manager.report_bin_statistics = MagicMock(return_value=True)

        self.work_manager = westpa.rc.get_work_manager()
        self.work_manager.running = True

        data = self.sim_manager.we_driver.rc.get_data_manager()
        data.we_h5filename = self.hdf5
        data.prepare_backing()
        data.create_ibstate_group([])
        data.create_initial_states(1)
        data.save_target_states([])
        data.update_segments = MagicMock(return_value=None)

        n_iter = 0
        it_name = data.iter_group_name(n_iter)
        for group in ["seg_index", "parents", "ibstates", "pcoord"]:
            data.we_h5file.create_group(it_name + "/" + group)
        data.get_new_weight_data = MagicMock(return_value=None)
        data.get_segments = MagicMock(return_value=self.segments)
        self.sim_manager.we_driver.rc.get_data_manager = MagicMock(return_value=data)
        self.sim_manager.n_iter = n_iter

    def tearDown(self):
        westpa.rc._sim_manager = None
        westpa.rc._system = None
        westpa.rc._data_manager = None
        del os.environ['WEST_SIM_ROOT']

    def dummy_callback_one(self):
        system = self.sim_manager.system
        bounds = [0.0, 1.0, 2.0, 3.0]
        system.bin_mapper = RectilinearBinMapper([bounds])

    def dummy_callback_two(self):
        system = self.sim_manager.system
        bounds = [0.0, 1.0, 2.0, 5.0]
        system.bin_mapper = RectilinearBinMapper([bounds])

    def segment(self, init_pcoord, final_pcoord, weight=1.0):
        segment = Segment(n_iter=1, seg_id=1123, pcoord=self.sim_manager.system.new_pcoord_array(), weight=weight)
        segment.pcoord[0] = init_pcoord
        segment.pcoord[1] = final_pcoord
        segment.parent_id = 1
        return segment

    def test_sim_manager(self):
        self.assertEqual(self.sim_manager.n_propagated, 0)
        self.assertEqual(len(self.sim_manager._callback_table), 0)

    def test_register_callback(self):
        hook = self.sim_manager.prepare_new_iteration

        self.sim_manager.register_callback(hook, self.dummy_callback_one, 3)
        self.sim_manager.register_callback(hook, self.dummy_callback_two, 0)
        self.sim_manager.register_callback(
            hook, dummy_callback_one, 3
        )  # Same name and priority, but different function, should be added
        self.sim_manager.register_callback(
            hook, self.dummy_callback_one, 2
        )  # Duplicate should never be added, even with different priority
        self.assertTrue(hook in self.sim_manager._callback_table)

        callbacks = self.sim_manager._callback_table.get(hook, [])

        assert len(callbacks) == 3  # Make sure only 3 added.

        self.assertTrue((3, self.dummy_callback_one.__name__, self.dummy_callback_one) in callbacks)  # noqa
        self.assertTrue((0, self.dummy_callback_two.__name__, self.dummy_callback_two) in callbacks)  # noqa
        self.assertTrue((3, dummy_callback_one.__name__, dummy_callback_one) in callbacks)  # noqa

    def test_invoke_callback(self):
        hook = self.sim_manager.prepare_new_iteration

        self.sim_manager.register_callback(hook, self.dummy_callback_one, 3)
        self.sim_manager.register_callback(hook, self.dummy_callback_two, 0)

        self.sim_manager.invoke_callbacks(hook)

        system = self.sim_manager.system
        self.assertTrue(np.all(system.bin_mapper.boundaries == np.array([0.0, 1.0, 2.0, 3.0])))  # noqa

    def test_process_config(self):
        self.sim_manager.process_config()
        self.assertTrue(self.sim_manager.do_gen_istates)
        self.assertEqual(self.sim_manager.propagator_block_size, 10000)
        self.assertFalse(self.sim_manager.save_transition_matrices)
        self.assertEqual(self.sim_manager.max_run_walltime, 10800)
        self.assertEqual(self.sim_manager.max_total_iterations, 100)

    def test_load_plugins(self):
        self.sim_manager.load_plugins()

    def test_report_bin_statistics(self):
        self.sim_manager.report_bin_statistics([0.0, 1.0, 2.0, 5.0])

    def test_get_bstate_pcoords(self):
        self.sim_manager.get_bstate_pcoords(self.basis_states)

    def test_report_basis_states(self):
        self.sim_manager.report_basis_states(self.basis_states)

    def test_report_target_states(self):
        self.sim_manager.report_target_states(self.basis_states)

    @pytest.mark.skip(reason="Cannot currently test WESimManager.initialize_simulation()")
    def test_initialize_simulation(self):
        # TODO: determine how to test self.simulation_manager.initialize_simulation()
        pass

    def test_prepare_iteration(self):
        self.sim_manager.prepare_new_iteration()
        self.sim_manager.prepare_iteration()

    def test_finalize_iteration(self):
        self.sim_manager.finalize_iteration()

    def test_get_istate_futures(self):
        self.sim_manager.get_istate_futures()

    def test_propagate(self):
        westpa.core.states.pare_basis_initial_states = MagicMock(return_value=([], []))
        self.sim_manager.propagate

    def test_save_bin_data(self):
        self.sim_manager.save_bin_data()

    def test_check_propagation(self):
        self.assertRaises(PropagationError, self.sim_manager.check_propagation)

    def test_run_we(self):
        self.sim_manager.run_we()

    def test_run(self):
        self.sim_manager.run()

    def test_prepare_run(self):
        self.sim_manager.prepare_run()

    def test_finalize_run(self):
        self.sim_manager.finalize_run()

    def test_pre_propagation(self):
        self.sim_manager.pre_propagation()

    def test_post_propagation(self):
        self.sim_manager.post_propagation()

    def test_pre_we(self):
        self.sim_manager.pre_we()

    def test_post_we(self):
        self.sim_manager.post_we()


class TestMABSimManager(TestSimManager):
    def setUp(self):
        parser = argparse.ArgumentParser()
        westpa.rc.add_args(parser)

        here = os.path.dirname(__file__)
        os.environ['WEST_SIM_ROOT'] = os.path.join(here, 'refs', 'odld')

        config_file_name = os.path.join(here, 'refs', 'odld', 'west_mab.cfg')
        args = parser.parse_args(['-r={}'.format(config_file_name)])
        westpa.rc.process_args(args)
        self.sim_manager = MABSimManager()
        self.test_dir = tempfile.mkdtemp()
        self.hdf5 = os.path.join(self.test_dir, "west.h5")
        self.basis_states = [BasisState(label="label", probability=1.0)]
        self.segments = [self.segment(0.0, 1.5, weight=0.125) for _i in range(4)] + [
            self.segment(1.5, 0.5, weight=0.125) for _i in range(4)
        ]
        self.sim_manager.we_driver.new_iteration()
        self.sim_manager.we_driver.assign(self.segments)
        self.sim_manager.we_driver.construct_next()
        self.sim_manager.segments = {segment.seg_id: segment for segment in self.segments}
        self.sim_manager.incomplete_segments = self.sim_manager.segments
        self.sim_manager.current_iter_istates = self.sim_manager.segments
        self.sim_manager.completed_segments = self.sim_manager.segments
        self.sim_manager.report_bin_statistics = MagicMock(return_value=True)

        data = self.sim_manager.we_driver.rc.get_data_manager()
        data.we_h5filename = self.hdf5
        data.prepare_backing()
        data.create_ibstate_group([])
        data.create_initial_states(1)
        data.save_target_states([])
        data.update_segments = MagicMock(return_value=None)

        n_iter = 0
        it_name = data.iter_group_name(n_iter)
        for group in ["seg_index", "parents", "ibstates", "pcoord"]:
            data.we_h5file.create_group(it_name + "/" + group)
        data.get_new_weight_data = MagicMock(return_value=None)
        data.get_segments = MagicMock(return_value=self.segments)
        self.sim_manager.we_driver.rc.get_data_manager = MagicMock(return_value=data)
        self.sim_manager.n_iter = n_iter


class TestBinlessSimManager(TestSimManager):
    def setUp(self):
        parser = argparse.ArgumentParser()
        westpa.rc.add_args(parser)

        here = os.path.dirname(__file__)
        os.environ['WEST_SIM_ROOT'] = os.path.join(here, 'refs', 'odld')

        config_file_name = os.path.join(here, 'refs', 'odld', 'west_binless.cfg')
        args = parser.parse_args(['-r={}'.format(config_file_name)])
        westpa.rc.process_args(args)
        self.sim_manager = BinlessSimManager()
        self.test_dir = tempfile.mkdtemp()
        self.hdf5 = os.path.join(self.test_dir, "west.h5")
        self.basis_states = [BasisState(label="label", probability=1.0)]
        self.segments = [self.segment(0.0, 1.5, weight=0.125) for _i in range(4)] + [
            self.segment(1.5, 0.5, weight=0.125) for _i in range(4)
        ]
        self.sim_manager.we_driver.new_iteration()
        self.sim_manager.we_driver.assign(self.segments)
        self.sim_manager.we_driver.construct_next()
        self.sim_manager.segments = {segment.seg_id: segment for segment in self.segments}
        self.sim_manager.incomplete_segments = self.sim_manager.segments
        self.sim_manager.current_iter_istates = self.sim_manager.segments
        self.sim_manager.completed_segments = self.sim_manager.segments
        self.sim_manager.report_bin_statistics = MagicMock(return_value=True)

        data = self.sim_manager.we_driver.rc.get_data_manager()
        data.we_h5filename = self.hdf5
        data.prepare_backing()
        data.create_ibstate_group([])
        data.create_initial_states(1)
        data.save_target_states([])
        data.update_segments = MagicMock(return_value=None)

        n_iter = 0
        it_name = data.iter_group_name(n_iter)
        for group in ["seg_index", "parents", "ibstates", "pcoord"]:
            data.we_h5file.create_group(it_name + "/" + group)
        data.get_new_weight_data = MagicMock(return_value=None)
        data.get_segments = MagicMock(return_value=self.segments)
        self.sim_manager.we_driver.rc.get_data_manager = MagicMock(return_value=data)
        self.sim_manager.n_iter = n_iter

    @pytest.mark.skip('Not configured')
    def test_run(self):
        self.sim_manager.run()
