import pytest
from unittest import TestCase

import numpy as np

from westpa.core.binning import RectilinearBinMapper
from westpa.core.segment import Segment
from westpa.core.states import TargetState, InitialState
from westpa.core.systems import WESTSystem
from westpa.core.we_driver import WEDriver

EPS = np.finfo(np.float64).eps


class TestWEDriver(TestCase):
    def setUp(self):
        system = WESTSystem()
        system.bin_mapper = RectilinearBinMapper([[0.0, 1.0, 2.0]])
        system.bin_target_counts = np.array([4, 4])
        system.pcoord_len = 2
        self.we_driver = WEDriver(system=system)
        self.system = system
        self._seg_id = 0

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

    def tearDown(self):
        self.we_driver.clear()
        del self.we_driver
        del self.system
        del self._seg_id

    def test_assign(self):
        segments = [self.segment(0.0, 1.5, weight=0.5), self.segment(1.5, 0.5, weight=0.5)]
        self.we_driver.new_iteration()
        n_recycled = self.we_driver.assign(segments)
        assert n_recycled == 0
        assert len(self.we_driver.initial_binning[0]) == 1
        assert len(self.we_driver.initial_binning[1]) == 1
        assert len(self.we_driver.final_binning[0]) == 1
        assert len(self.we_driver.final_binning[1]) == 1
        assert (self.we_driver.flux_matrix == np.array([[0.0, 0.5], [0.5, 0.0]])).all()

    def test_passthrough(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)
        ]
        segs_by_id = {segment.seg_id: segment for segment in segments}
        self.we_driver.new_iteration()
        self.we_driver.assign(segments)
        self.we_driver.construct_next()
        out_segs = set()
        for bin in self.we_driver.next_iter_binning:
            out_segs.update(bin)
        assert out_segs == set(self.we_driver.next_iter_segments)
        assert abs(sum(seg.weight for seg in self.we_driver.next_iter_segments) - 1.0) < 8 * EPS

        for segment in self.we_driver.next_iter_segments:
            # has n_iter been advanced?
            assert segment.n_iter == 2

            # is weight correct?
            assert segment.weight == 0.125  # rigorous floating point comparison okay here because no math was done

            # is parent set correctly?
            assert segment.parent_id is not None

            # was pcoord set correctly?
            assert (segs_by_id[segment.parent_id].pcoord[-1] == segment.pcoord[0]).all()

            # was status set correctly?
            assert segment.status == Segment.SEG_STATUS_PREPARED

        # were parent endpoint types set correctly
        for segment in segments:
            assert segment.endpoint_type == Segment.SEG_ENDPOINT_CONTINUES

    def test_split_by_weight(self):
        segments = [self.segment(1.5, 0.5, weight=0.25), self.segment(0.0, 1.5, weight=0.75)]
        self.we_driver.new_iteration()
        self.we_driver.assign(segments)
        self.we_driver.construct_next()
        assert len(self.we_driver.next_iter_binning[0]) == 4
        assert len(self.we_driver.next_iter_binning[1]) == 4
        assert abs(sum(seg.weight for seg in self.we_driver.next_iter_binning[0]) - 0.25) < 4 * EPS
        assert abs(sum(seg.weight for seg in self.we_driver.next_iter_binning[1]) - 0.75) < 4 * EPS
        assert np.allclose([seg.weight for seg in self.we_driver.next_iter_binning[0]], [0.25 / 4.0 for _i in range(4)])
        assert np.allclose([seg.weight for seg in self.we_driver.next_iter_binning[1]], [0.75 / 4.0 for _i in range(4)])

        for ibin in range(2):
            for segment in self.we_driver.next_iter_binning[ibin]:
                print(segment)
                assert segment.n_iter == 2
                assert segment.parent_id is not None
                assert segment.parent_id == segments[ibin].seg_id
                assert segment.status == Segment.SEG_STATUS_PREPARED

    # this test will fail up to alpha of the time
    @pytest.mark.flaky(reruns=5, only_rerun="AssertionError")
    def test_merge_by_weight(self):
        selected_counts = {0: 0, 1: 0}
        alpha = 0.01
        nrounds = 1000
        from scipy.stats import binom

        # lower and upper bounds of 95% CI for selecting the segment with weight 1/3
        lb = binom.ppf(alpha / 2.0, nrounds, 1.0 / 3.0)
        ub = binom.ppf(1.0 - alpha / 2.0, nrounds, 1.0 / 3.0)

        system = WESTSystem()
        system.bin_mapper = RectilinearBinMapper([[0.0, 1.0]])
        system.bin_target_counts = np.array([1])
        system.pcoord_len = 2
        self.we_driver = WEDriver(system=system)
        self.system = system
        self._seg_id = 0

        segments = [
            Segment(n_iter=1, seg_id=0, pcoord=np.array([[0], [0.25]], dtype=np.float32), weight=1.0 / 3.0),
            Segment(n_iter=1, seg_id=1, pcoord=np.array([[0], [0.75]], dtype=np.float32), weight=2.0 / 3.0),
        ]

        for _iround in range(nrounds):
            for segment in segments:
                segment.endpoint_type = Segment.SEG_ENDPOINT_UNSET

            self.we_driver.new_iteration()
            self.we_driver.assign(segments)
            self.we_driver.construct_next()

            assert len(self.we_driver.next_iter_binning[0]) == 1
            newseg = self.we_driver.next_iter_binning[0].pop()

            assert segments[newseg.parent_id].endpoint_type == Segment.SEG_ENDPOINT_CONTINUES
            assert segments[~newseg.parent_id].endpoint_type == Segment.SEG_ENDPOINT_MERGED

            selected_counts[newseg.parent_id] += 1

        print(selected_counts)
        assert (
            lb <= selected_counts[0] <= ub
        ), 'Incorrect proportion of histories selected.' 'this is expected about {:%} of the time; retry test.'.format(alpha)

    def test_split_with_adjust(self):
        # this is a split followed by merge
        self.system.bin_target_counts = np.array([5, 5])
        segments = [
            self.segment(1.5, 0.5, weight=0.125),
            self.segment(1.5, 0.5, weight=0.125),
            self.segment(0.0, 1.5, weight=0.375),
            self.segment(0.0, 1.5, weight=0.375),
        ]
        self.we_driver.new_iteration()
        self.we_driver.assign(segments)
        self.we_driver.construct_next()
        assert len(self.we_driver.next_iter_binning[0]) == 5
        assert len(self.we_driver.next_iter_binning[1]) == 5
        assert abs(sum(seg.weight for seg in self.we_driver.next_iter_binning[0]) - 0.25) < 5 * EPS
        assert abs(sum(seg.weight for seg in self.we_driver.next_iter_binning[1]) - 0.75) < 5 * EPS

        for ibin in range(2):
            for segment in self.we_driver.next_iter_binning[ibin]:
                print(segment)
                assert segment.n_iter == 2
                assert segment.parent_id is not None
                assert segment.status == Segment.SEG_STATUS_PREPARED

    def test_split_with_adjust_istates(self):
        # this is a split followed by merge, for segments which are initial states
        self.system.bin_target_counts = np.array([5, 5])
        # segments = [
        # self.segment(1.5, 0.5, weight=0.125),
        # self.segment(1.5, 0.5, weight=0.125),
        # self.segment(0.0, 1.5, weight=0.375),
        # self.segment(0.0, 1.5, weight=0.375),
        # ]
        self.we_driver.new_iteration()
        self.we_driver._prep_we()
        self.we_driver.used_initial_states[-1] = None
        self.we_driver.used_initial_states[-2] = None

        for ibin, bin in enumerate(self.we_driver.next_iter_binning):
            pc = np.array([[0.5 + ibin], [0.0]])
            for iseg in range(6):
                segment = Segment(n_iter=1, seg_id=None, weight=1.0 / 12.0, parent_id=-(ibin + 1), pcoord=pc)
                bin.add(segment)

        # for ibin in range(len(self.we_driver.next_iter_binning)):
        # This will raise KeyError if initial state tracking is done improperly
        #     self.we_driver._adjust_count(ibin)

        for ibin, bin in enumerate(self.we_driver.next_iter_binning):
            target_count = self.we_driver.bin_target_counts[ibin]
            subgroups = self.we_driver.subgroup_function(self.we_driver, ibin, **self.we_driver.subgroup_function_kwargs)
            self.we_driver._adjust_count(bin, subgroups, target_count)

        assert len(self.we_driver.next_iter_binning[0]) == 5
        assert len(self.we_driver.next_iter_binning[1]) == 5

    def test_recycle(self):
        segments = [self.segment(0.0, 1.5, weight=0.5), self.segment(0.0, 0.5, weight=0.5)]
        tstate = TargetState('recycle', [1.5], 0)
        istate = InitialState(0, 0, 0, pcoord=[0.0])

        self.we_driver.new_iteration(initial_states=[istate], target_states=[tstate])
        n_needed = self.we_driver.assign(segments)
        assert n_needed == 0
        self.we_driver.construct_next()

        n_recycled = len(list(self.we_driver.recycling_segments))
        assert n_recycled == 1

        assert len(self.we_driver.next_iter_binning[0]) == 4
        assert len(self.we_driver.next_iter_binning[1]) == 0
        assert abs(sum(seg.weight for seg in self.we_driver.next_iter_binning[0]) - 1.0) < 4 * EPS
        assert np.allclose([seg.weight for seg in self.we_driver.next_iter_binning[0]], [0.25 for _i in range(4)])
        assert segments[0].endpoint_type == Segment.SEG_ENDPOINT_RECYCLED

    def test_multiple_merge(self):
        # This weight and count combination is known to trigger a split to 51
        # followed by a count adjustment to 50 (thanks to Josh Adelman)
        segment = self.segment(0.0, 0.5, weight=0.9999999999970001)

        # Initial state ID 0
        segment.parent_id = -1
        segment.wtg_parent_ids = set([-1])
        assert segment.initpoint_type == segment.SEG_INITPOINT_NEWTRAJ

        self.system.bin_target_counts = np.array([50, 50])
        self.we_driver.new_iteration()
        self.we_driver.assign([segment])
        self.we_driver.construct_next()

        assert len(self.we_driver.next_iter_binning[0]) == 50

    def check_populate_initial(self, prob, target_counts):
        istate = InitialState(0, 0, 0, pcoord=[0.0])
        self.system.bin_target_counts = np.array([target_counts, target_counts])

        self.we_driver.populate_initial([istate], [prob], system=self.system)
        assert len(self.we_driver.next_iter_binning[0]) == target_counts

    @pytest.mark.skip
    def test_populate_initial(self):
        for prob in [0.1, 1.0 / 3.0, 0.9999999999970001, 1.0]:
            for tcount in range(30, 60):
                yield self.check_populate_initial, prob, tcount

    # TODO: add test for seeding the flux matrix based on recycling
    # TODO: add test for split after merge in adjust count
