from typing import Dict, List
from unittest.mock import Mock

from pytest import fixture, raises
from smif.data_layer.store import Store
from smif.decision.decision import DecisionManager, RuleBased
from smif.exception import SmifDataNotFoundError


@fixture(scope='function')
def plan():
    planned_interventions = [
        {'name': 'small_pumping_station_oxford', 'build_year': 2010},
        {'name': 'small_pumping_station_abingdon', 'build_year': 2015},
        {'name': 'large_pumping_station_oxford', 'build_year': 2020}
    ]

    return planned_interventions


@fixture(scope='function')
def get_strategies():
    strategies = [{'strategy': 'pre-specified-planning',
                   'description': 'build_nuclear',
                   'interventions': [
                       {'name': 'nuclear_large', 'build_year': 2012},
                       {'name': 'carrington_retire', 'build_year': 2011}]
                   }]

    return strategies


@fixture(scope='function')
def get_register():
    lifetime = {'technical_lifetime': {'value': 99}}
    register = {'nuclear_large': lifetime,
                'carrington_retire': lifetime,
                'small_pumping_station_oxford': lifetime,
                'small_pumping_station_abingdon': lifetime,
                'large_pumping_station_oxford': lifetime
                }
    return register


class TestRuleBasedProperties:

    def test_timesteps(self):

        timesteps = [2010, 2015, 2020]
        dm = RuleBased(timesteps, Mock())
        assert dm.current_timestep == 2010
        assert dm.next_timestep == 2015

    def test_timesteps_end(self):
        timesteps = [2010, 2015, 2020]
        dm = RuleBased(timesteps, Mock())
        dm.current_timestep = 2020
        assert dm.next_timestep is None

    def test_timesteps_begin(self):
        timesteps = [2010, 2015, 2020]
        dm = RuleBased(timesteps, Mock())
        dm.current_timestep = 2010
        assert dm.previous_timestep is None

    def test_timesteps_first_last(self):

        timesteps = [2010, 2015, 2020]
        dm = RuleBased(timesteps, Mock())
        assert dm.first_timestep == 2010
        assert dm.last_timestep == 2020

    def test_interventions(self):

        all_interventions = {'test_intervention': {'name': 'test_intervention'}}

        timesteps = [2010, 2015, 2020]
        dm = RuleBased(timesteps, all_interventions)
        assert dm.available_interventions([]) == ['test_intervention']

    def test_interventions_planned(self):

        all_interventions = {'test_intervention': {'name': 'test_intervention'},
                             'planned_intervention': {'name': 'planned_intervention'}}

        timesteps = [2010, 2015, 2020]
        dm = RuleBased(timesteps, all_interventions)
        actual = dm.available_interventions([{'name': 'planned_intervention'}])
        expected = ['test_intervention']
        assert actual == expected

    def test_get_intervention(self):

        interventions = {'a': {'name': 'a'},
                         'b': {'name': 'b'},
                         'c': {'name': 'c'}}

        timesteps = [2010, 2015, 2020]
        dm = RuleBased(timesteps, interventions)
        assert dm.get_intervention('a') == interventions['a']
        with raises(SmifDataNotFoundError) as ex:
            dm.get_intervention('z')
        msg = "Intervention 'z' is not found in the list of available interventions"
        assert msg in str(ex.value)


class TestRuleBasedIterationTimestepAccounting:
    """Test that the iteration and timestep accounting methods properly follow
    the path through the decision iterations

    2010 - 0, 1
    2015 - 2, 3
    """

    @fixture(scope='function')
    def dm(self):
        timesteps = [2010, 2015, 2020]
        dm = RuleBased(timesteps, Mock())
        return dm

    def test_first_iteration_base_year(self, dm):

        dm.current_timestep = 2010
        dm.current_iteration = 1
        dm._max_iteration_by_timestep[2010] = 1
        assert dm.get_previous_iteration_timestep() is None

    def test_second_iteration_base_year(self, dm):

        dm.current_timestep = 2010
        dm.current_iteration = 2
        dm._max_iteration_by_timestep[2010] = 2
        assert dm.get_previous_iteration_timestep() == (2010, 1)

    def test_second_iteration_next_year(self, dm):

        dm.current_timestep = 2015
        dm.current_iteration = 3
        dm._max_iteration_by_timestep[2010] = 2
        dm._max_iteration_by_timestep[2015] = 3
        assert dm.get_previous_iteration_timestep() == (2010, 2)

    def test_third_iteration_next_year(self, dm):

        dm.current_timestep = 2015
        dm.current_iteration = 4
        dm._max_iteration_by_timestep[2010] = 2
        dm._max_iteration_by_timestep[2015] = 4
        assert dm.get_previous_iteration_timestep() == (2015, 3)


class TestRuleBased:

    def test_initialisation(self):

        timesteps = [2010, 2015, 2020]
        dm = RuleBased(timesteps, Mock())
        assert dm.timesteps == timesteps
        assert dm.satisfied is False
        assert dm.current_timestep == 2010
        assert dm.current_iteration == 0

    def test_generator(self):

        timesteps = [2010, 2015, 2020]
        dm = RuleBased(timesteps, Mock())

        actual = next(dm)
        assert actual == {
            'decision_iterations': [1],
            'timesteps': [2010]
        }

        dm.satisfied = True
        actual = next(dm)
        assert actual == {
            'decision_iterations': [2],
            'timesteps': [2015],
            'decision_links': {
                2: 1
            }
        }
        assert dm.satisfied is False

        actual = next(dm)
        assert actual == {
            'decision_iterations': [3],
            'timesteps': [2015],
            'decision_links': {
                3: 1
            }
        }
        assert dm.satisfied is False

        dm.satisfied = True
        dm.current_timestep = 2020
        actual = next(dm)
        assert actual is None


class TestDecisionManager():

    @fixture(scope='function')
    def decision_manager(self, empty_store) -> DecisionManager:
        empty_store.write_model_run({'name': 'test', 'sos_model': 'test_sos_model'})
        empty_store.write_sos_model({'name': 'test_sos_model', 'sector_models': []})
        empty_store.write_strategies('test', [])
        sos_model = Mock()
        sos_model.name = 'test_sos_model'
        sos_model.sector_models = []

        df = DecisionManager(empty_store, [2010, 2015], 'test', sos_model)
        return df

    def test_decision_manager_init(self,  decision_manager: DecisionManager):
        df = decision_manager
        dm = df.decision_loop()
        bundle = next(dm)
        assert bundle == {
            'decision_iterations': [0],
            'timesteps': [2010, 2015]
        }
        with raises(StopIteration):
            next(dm)

    def test_available_interventions(self, decision_manager: DecisionManager):
        df = decision_manager
        df._register = {'a': {'name': 'a'},
                        'b': {'name': 'b'},
                        'c': {'name': 'c'}}

        assert df.available_interventions == df._register

        df.planned_interventions = {(2010, 'a'), (2010, 'b')}

        expected = {'c': {'name': 'c'}}

        assert df.available_interventions == expected

    def test_get_intervention(self,  decision_manager: DecisionManager):
        df = decision_manager
        df._register = {'a': {'name': 'a'},
                        'b': {'name': 'b'},
                        'c': {'name': 'c'}}

        assert df.get_intervention('a') == {'name': 'a'}

        with raises(SmifDataNotFoundError):
            df.get_intervention('z')

    def test_buildable(self, decision_manager):

        decision_manager._timesteps = [2010, 2015]
        assert decision_manager.buildable(2010, 2010) is True
        assert decision_manager.buildable(2011, 2010) is True

    def test_historical_intervention_buildable(self, decision_manager):
        decision_manager._timesteps = [2020, 2030]
        assert decision_manager.buildable(1980, 2020) is True
        assert decision_manager.buildable(1990, 2020) is True

    def test_buildable_raises(self, decision_manager):

        with raises(ValueError):
            decision_manager.buildable(2015, 2014)

    def test_within_lifetime(self, decision_manager):

        assert decision_manager.within_lifetime(2010, 2010, 1)

    def test_within_lifetime_does_not_check_start(self, decision_manager):
        """Note that the ``within_lifetime`` method does not check
        that the build year is compatible with timestep
        """
        assert decision_manager.within_lifetime(2011, 2010, 1)

    def test_negative_lifetime_raises(self, decision_manager):
        with raises(ValueError):
            decision_manager.within_lifetime(2010, 2010, -1)


class TestDecisionManagerDecisions:

    @fixture(scope='function')
    def decision_manager(self, empty_store) -> DecisionManager:
        empty_store.write_model_run({'name': 'test', 'sos_model': 'test_sos_model'})
        empty_store.write_sos_model({'name': 'test_sos_model', 'sector_models': []})
        empty_store.write_strategies('test', [])
        sos_model = Mock()
        sos_model.name = 'test_sos_model'
        sos_model.sector_models = []

        interventions = {'test': {'technical_lifetime': {'value': 99}},
                         'planned': {'technical_lifetime': {'value': 99}},
                         'decided': {'technical_lifetime': {'value': 99}}
                         }

        df = DecisionManager(empty_store, [2010, 2015], 'test', sos_model)
        df._register = interventions
        return df

    def test_get_decisions(self, decision_manager: DecisionManager):
        dm = decision_manager

        mock_handle = Mock()
        dm._decision_module = Mock()
        dm._decision_module.get_decision = Mock(
            return_value=[{'name': 'test', 'build_year': 2010}])

        actual = dm._get_decisions(dm._decision_module, mock_handle)
        expected = [(2010, 'test')]
        assert actual == expected

    def test_get_and_save_decisions_dm(self, decision_manager: DecisionManager):
        """Test that the ``get_and_save_decisions`` method updates pre-decision
        state with a new decision and writes it to store
        """
        dm = decision_manager

        dm._decision_module = Mock()
        dm._decision_module.get_decision = Mock(
            return_value=[{'name': 'decided', 'build_year': 2010}])
        dm._decision_module.get_previous_state = Mock(return_value=[])

        dm.get_and_save_decisions(0, 2010)

        actual = dm._store  # type: Store
        expected = [{'name': 'decided', 'build_year': 2010}]  # type: List[Dict]
        assert actual.read_state('test', 2010, decision_iteration=0) == expected

    def test_get_and_save_decisions_prespec(self,
                                            decision_manager: DecisionManager):
        """Test that the ``get_and_save_decisions`` method updates pre-decision
        state with a pre-specified planning and writes it to store
        """
        dm = decision_manager

        dm.planned_interventions = [(2010, 'planned')]

        dm.get_and_save_decisions(0, 2010)

        actual = dm._store  # type: Store
        expected = [{'name': 'planned', 'build_year': 2010}]  # type: List[Dict]
        assert actual.read_state('test', 2010, decision_iteration=0) == expected

    def test_pre_spec_and_decision_module(self,
                                          decision_manager: DecisionManager):
        dm = decision_manager

        dm._decision_module = Mock()
        dm._decision_module.get_decision = Mock(
            return_value=[{'name': 'decided', 'build_year': 2010}])
        dm._decision_module.get_previous_state = Mock(return_value=[])

        dm.planned_interventions = [(2010, 'planned')]

        dm.get_and_save_decisions(0, 2010)

        actual = dm._store.read_state('test', 2010, decision_iteration=0)  # type: List[Dict]

        expected = set([('decided', 2010), ('planned', 2010)])
        assert set([(x['name'], x['build_year']) for x in actual]) == expected
