"""Test ModelData
"""
# pylint: disable=redefined-outer-name
from copy import copy, deepcopy
from unittest.mock import Mock

import numpy as np
from pytest import fixture, raises
from smif.data_layer import DataHandle
from smif.data_layer.data_array import DataArray
from smif.data_layer.data_handle import ResultsHandle
from smif.exception import (
    SmifDataError,
    SmifDataMismatchError,
    SmifDataNotFoundError,
    SmifTimestepResolutionError,
)
from smif.metadata import Spec
from smif.model import SectorModel, SosModel


@fixture(scope="function")
def mock_store(sample_dimensions, annual, get_sector_model, empty_store):
    """Store with minimal setup"""
    store = empty_store

    store.write_unit_definitions(["people", "million people = 1000000 * people"])

    for dim in sample_dimensions:
        store.write_dimension(dim)

    # energy demand model
    store.write_model(get_sector_model)
    for param in get_sector_model["parameters"]:
        spec = Spec.from_dict(param)
        data = DataArray(spec, np.array(42))
        store.write_model_parameter_default("energy_demand", spec.name, data)

    store.write_initial_conditions("energy_demand", [])
    data = {
        "water_asset_a": {
            "name": "water_asset_a",
            "build_year": 2010,
            "capacity": 50,
            "location": None,
            "sector": "",
        },
        "water_asset_b": {
            "name": "water_asset_b",
            "build_year": 2015,
            "capacity": 150,
            "location": None,
            "sector": "",
        },
        "water_asset_c": {
            "name": "water_asset_c",
            "capacity": 100,
            "build_year": 2015,
            "location": None,
            "sector": "",
        },
    }
    store.write_interventions("energy_demand", data)

    # source scenario
    pop_spec = {
        "name": "population_count",
        "dtype": "int",
        "dims": ["lad", "annual"],
        "coords": {"lad": sample_dimensions[0]["elements"], "annual": annual},
        "absolute_range": [0, int(1e12)],
        "expected_range": [0, 100000],
        "unit": "people",
    }
    store.write_scenario(
        {
            "name": "population",
            "description": "Projected annual fertility rates",
            "provides": [pop_spec],
            "variants": [
                {
                    "name": "low",
                    "description": "Population (Low)",
                    "data": {
                        "population_count": "population_low_count",
                    },
                }
            ],
        }
    )
    pop_spec_t = deepcopy(pop_spec)
    pop_spec_t["dims"] = ["timestep"] + pop_spec["dims"]
    pop_spec_t["coords"]["timestep"] = [{"name": 2015}]
    pop_da = DataArray(Spec.from_dict(pop_spec_t), np.array([[[1.0], [2.0]]]))
    store.write_scenario_variant_data("population", "low", pop_da)

    # source model
    # - test output matches population input to energy demand
    # - test_alt_dims output matches test input to test_convertor
    store.write_model(
        {
            "name": "test_source",
            "description": "",
            "inputs": [],
            "outputs": [
                {
                    "name": "test",
                    "dtype": "int",
                    "dims": ["lad", "annual"],
                    "absolute_range": [0, int(1e12)],
                    "expected_range": [0, 100000],
                    "unit": "people",
                },
                {
                    "name": "test_alt_dims",
                    "dtype": "int",
                    "dims": ["lad", "annual"],
                    "absolute_range": [0, int(1e12)],
                    "expected_range": [0, 100000],
                    "unit": "million people",
                },
            ],
            "parameters": [],
            "interventions": [],
            "initial_conditions": [],
        }
    )

    # convertor model
    # - test input matches test output from test_source
    # - test output matches population input to energy demand
    store.write_model(
        {
            "name": "test_convertor",
            "description": "",
            "inputs": [
                {
                    "name": "test",
                    "dtype": "int",
                    "dims": ["lad", "annual"],
                    "absolute_range": [0, int(1e12)],
                    "expected_range": [0, 100000],
                    "unit": "million people",
                }
            ],
            "outputs": [
                {
                    "name": "test",
                    "dtype": "int",
                    "dims": ["lad", "annual"],
                    "absolute_range": [0, int(1e12)],
                    "expected_range": [0, 100000],
                    "unit": "people",
                }
            ],
            "parameters": [],
            "interventions": [],
            "initial_conditions": [],
        }
    )

    store.write_sos_model(
        {
            "name": "energy_demand_only",
            "sector_models": ["energy_demand"],
            "model_dependencies": [],
            "scenario_dependencies": [
                {
                    "source": "population",
                    "source_output": "population_count",
                    "sink_input": "population",
                    "sink": "energy_demand",
                }
            ],
            "narratives": [],
        }
    )

    store.write_sos_model(
        {
            "name": "test_sos_model",
            "sector_models": ["energy_demand", "test_source"],
            "scenario_dependencies": [],
            "model_dependencies": [
                {
                    "source": "test_source",
                    "source_output": "test",
                    "sink_input": "population",
                    "sink": "energy_demand",
                }
            ],
            "narratives": [
                {
                    "name": "test_narrative",
                    "description": "a narrative config",
                    "provides": {"energy_demand": ["smart_meter_savings"]},
                    "variants": [
                        {
                            "name": "high_tech_dsm",
                            "description": "High takeup",
                            "data": {"smart_meter_savings": "filename.csv"},
                        }
                    ],
                }
            ],
        }
    )
    parameter_spec = Spec.from_dict(get_sector_model["parameters"][0])
    da = DataArray(parameter_spec, np.array(99))
    store.write_narrative_variant_data(
        "test_sos_model", "test_narrative", "high_tech_dsm", da
    )

    store.write_sos_model(
        {
            "name": "test_converting_sos_model",
            "sector_models": ["energy_demand", "test_source", "test_convertor"],
            "scenario_dependencies": [],
            "model_dependencies": [
                {
                    "source": "test_source",
                    "source_output": "test_alt_dims",
                    "sink_input": "test",
                    "sink": "test_convertor",
                },
                {
                    "source": "test_convertor",
                    "source_output": "test",
                    "sink_input": "population",
                    "sink": "energy_demand",
                },
            ],
            "narratives": [],
        }
    )

    store.write_model_run(
        {"name": 1, "narratives": {}, "sos_model": "test_sos_model", "scenarios": {}}
    )
    store.write_model_run(
        {
            "name": 2,
            "narratives": {},
            "sos_model": "test_converting_sos_model",
            "scenarios": {},
        }
    )
    store.write_model_run(
        {
            "name": 3,
            "narratives": {},
            "sos_model": "energy_demand_only",
            "scenarios": {"population": "low"},
        }
    )

    return store


class EmptySectorModel(SectorModel):
    """Sector Model implementation"""

    def simulate(self, data):
        """no-op"""
        pass


@fixture
def mock_model(mock_store):
    return EmptySectorModel.from_dict(mock_store.read_model("energy_demand"))


@fixture
def mock_model_with_conversion(mock_store):
    return EmptySectorModel.from_dict(mock_store.read_model("test_convertor"))


@fixture
def mock_sos_model(mock_store, mock_model):
    test_source_model = EmptySectorModel.from_dict(mock_store.read_model("test_source"))
    return SosModel.from_dict(
        mock_store.read_sos_model("test_sos_model"), [mock_model, test_source_model]
    )


class TestDataHandle:
    """Test data handle"""

    def test_create(self, mock_store, mock_model):
        """should be created with a Store"""
        DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)

    def test_get_data_from_scenario(self, mock_store, mock_model):
        """should allow read access to input data from scenarios"""
        data_handle = DataHandle(mock_store, 3, 2015, [2015, 2020], mock_model)
        data = np.array([[1.0], [2.0]])

        input_spec = mock_model.inputs["population"]
        input_da = DataArray(input_spec, data)

        actual = data_handle.get_data("population")
        assert actual.name == "population"
        np.testing.assert_equal(actual, input_da)

    def test_get_data_from_model_output(self, mock_store, mock_model):
        """should allow read access to input data from model results"""
        data_handle = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)
        data = np.array([[1.0], [2.0]])

        input_spec = mock_model.inputs["population"]
        input_da = DataArray(input_spec, data)
        output_spec = copy(input_spec)
        output_spec.name = "test"  # use source spec name
        output_da = DataArray(output_spec, data)

        mock_store.write_results(
            output_da, 1, "test_source", 2015, None  # write source model results
        )
        actual = data_handle.get_data("population")
        assert output_da.name == "test"
        assert actual.name == "population"
        np.testing.assert_equal(actual, input_da)

    def test_get_data_with_conversion(self, mock_store, mock_model_with_conversion):
        """should convert liters to milliliters (1 -> 0.001)"""
        modelrun_name = 2
        data_handle = DataHandle(
            mock_store, modelrun_name, 2015, [2015, 2020], mock_model_with_conversion
        )
        data = np.array([[0.001], [0.002]])

        spec = copy(mock_model_with_conversion.inputs["test"])
        spec.name = "test_alt_dims"  # use source spec name
        da = DataArray(spec, data)

        mock_store.write_results(
            da,
            modelrun_name,
            "test_source",  # write results as though from source
            2015,
            None,
        )
        actual = data_handle.get_data("test")

        expected = DataArray(spec, data)

        np.testing.assert_equal(actual, expected)

    def test_get_base_timestep_data(self, mock_store, mock_model):
        """should allow read access to input data from base timestep"""
        data_handle = DataHandle(mock_store, 1, 2025, [2015, 2020, 2025], mock_model)
        data = np.array([[1.0], [2.0]])

        spec = copy(mock_model.inputs["population"])
        spec.name = "test"  # use source spec name
        da = DataArray(spec, data)

        mock_store.write_results(
            da,
            1,
            "test_source",  # write source model results
            2015,  # base timetep
            None,
        )

        actual = data_handle.get_base_timestep_data("population")
        expected = DataArray(spec, data)
        np.testing.assert_equal(actual, expected)

    def test_get_base_timestep_data_before_model_run(self, mock_store, mock_model):
        """Prior to a model run, there is no current timestep

        This should allow read access to input data from base timestep
        """
        data_handle = DataHandle(mock_store, 1, None, [2015, 2020, 2025], mock_model)
        data = np.array([[1.0], [2.0]])

        spec = copy(mock_model.inputs["population"])
        spec.name = "test"  # use source spec name
        da = DataArray(spec, data)

        mock_store.write_results(
            da,
            1,
            "test_source",  # write source model results
            2015,  # base timetep
            None,
        )

        actual = data_handle.get_base_timestep_data("population")
        expected = DataArray(spec, data)
        np.testing.assert_equal(actual, expected)

    def test_get_previous_timestep_data(self, mock_store, mock_model):
        """should allow read access to input data from previous timestep"""
        data_handle = DataHandle(mock_store, 1, 2025, [2015, 2020, 2025], mock_model)
        data = np.random.rand(*mock_model.inputs["population"].shape)
        spec = copy(mock_model.inputs["population"])
        spec.name = "test"  # use source spec name
        da = DataArray(spec, data)

        mock_store.write_results(
            da,
            1,
            "test_source",  # write source model results
            2020,  # previous timetep
            None,
        )
        actual = data_handle.get_previous_timestep_data("population")

        expected = DataArray(spec, data)

        assert actual == expected

    def test_get_data_with_square_brackets(self, mock_store, mock_model):
        """should allow dict-like read access to input data"""
        data_handle = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)
        data = np.random.rand(*mock_model.inputs["population"].shape)

        spec = copy(mock_model.inputs["population"])
        spec.name = "test"  # use source spec name
        da = DataArray(spec, data)

        mock_store.write_results(
            da,
            1,
            "test_source",  # write source model results
            2015,  # current timetep
            None,
        )
        actual = data_handle["population"]
        assert actual == da

    def test_set_data(self, mock_store, mock_model_with_conversion):
        """should allow write access to output data"""
        data = np.array([[1.0], [4.0]])
        data_handle = DataHandle(
            mock_store, 1, 2015, [2015, 2020], mock_model_with_conversion
        )

        spec = mock_model_with_conversion.outputs["test"]
        da = DataArray(spec, data)

        data_handle.set_results("test", data)
        actual = mock_store.read_results(
            1,
            "test_convertor",  # read results from model
            mock_model_with_conversion.outputs["test"],
            2015,
            None,
        )
        np.testing.assert_equal(actual.as_ndarray(), data)
        assert actual == da

    def test_set_data_wrong_shape(self, mock_store, mock_model_with_conversion):
        """should allow write access to output data"""
        data = np.array([[1.0, 1.0]])  # regions is 1, intervals is 1 not 2

        spec = mock_model_with_conversion.outputs["test"]
        with raises(SmifDataMismatchError) as ex:
            DataArray(spec, data)

        msg = "Data shape (1, 2) does not match spec " "(2, 1)"
        assert msg in str(ex.value)

    def test_set_data_with_square_brackets(self, mock_store, mock_model):
        """should allow dict-like write access to output data"""
        data = np.random.rand(2, 8)
        data_handle = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)

        spec = mock_model.outputs["gas_demand"]
        da = DataArray(spec, data)

        data_handle["gas_demand"] = data
        actual = mock_store.read_results(
            1,
            "energy_demand",  # read results from model
            mock_model.outputs["gas_demand"],
            2015,
        )
        np.testing.assert_equal(actual.as_ndarray(), data)
        assert actual == da

    def test_set_data_with_square_brackets_raises(self, mock_store, mock_model):
        """should allow dict-like write access to output data"""
        data = np.random.rand(2, 8)
        data_handle = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)

        spec = mock_model.outputs["gas_demand"]
        da = DataArray(spec, data)

        with raises(TypeError) as ex:
            data_handle["gas_demand"] = da

        assert "Pass in a numpy array" in str(ex.value)


class TestDataHandleState:
    """Test handling of initial conditions, decision interventions and intervention state."""

    def test_get_state(self, mock_store, mock_model):
        """should get decision module state for given timestep/decision_iteration

        A call to ``get_state`` method on the data handle calls the read_state
        method of the store with arguments for model run name, current
        timestep and decision iteration.
        """
        mock_store.read_state = Mock(
            return_value=[{"name": "test", "build_year": 2010}]
        )
        mock_store.write_interventions(
            "energy_demand",
            [{"name": "test", "capital_cost": {"value": 2500, "unit": "£/GW"}}],
        )
        data_handle = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)
        expected = [{"name": "test", "build_year": 2010}]
        actual = data_handle.get_state()
        mock_store.read_state.assert_called_with(1, 2015, None)
        assert actual == expected

    def test_get_interventions_for_sector_model(self, mock_store, mock_model):
        """

        A call to the ``get_current_interventions`` method of the data_handle
        returns a dict of interventions from the current state filtered by the
        current model

        """
        state = [
            {"name": "water_asset_a", "build_year": 2010},
            {"name": "water_asset_b", "build_year": 2015},
        ]
        mock_store.write_state(state, 1, 2015, None)
        # mock_store._strategies[1] = []

        data_handle = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)

        actual = data_handle.get_current_interventions()
        expected = {
            "water_asset_a": {
                "name": "water_asset_a",
                "build_year": 2010,
                "capacity": 50,
                "location": None,
                "sector": "",
            },
            "water_asset_b": {
                "name": "water_asset_b",
                "build_year": 2015,
                "capacity": 150,
                "location": None,
                "sector": "",
            },
        }
        assert actual == expected

    def test_interventions_sector_model_ignore_unrecog(self, mock_store, mock_model):
        """Ignore unrecognised interventions

        Interventions that are listed in state, but not included in the
        intervention list are ignored

        """
        state = [
            {"name": "water_asset_a", "build_year": 2010},
            {"name": "energy_asset_unexpected", "build_year": 2015},
        ]

        mock_store.write_state(state, 1, 2015, None)

        data_handle = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)
        actual = data_handle.get_current_interventions()
        expected = {
            "water_asset_a": {
                "name": "water_asset_a",
                "build_year": 2010,
                "capacity": 50,
                "location": None,
                "sector": "",
            }
        }
        assert actual == expected

    def test_pass_none_to_timestep_raises(self, mock_store, mock_model):
        data_handle = DataHandle(mock_store, 1, None, [2015, 2020], mock_model)
        with raises(ValueError):
            data_handle.get_state()


class TestDataHandleTimesteps:
    """Test timestep helper properties"""

    def test_current_timestep(self, mock_model, mock_store):
        """should return current timestep"""
        data_handle = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)
        assert data_handle.current_timestep == 2015

    def test_base_timestep(self, mock_model, mock_store):
        """should return first timestep in list"""
        data_handle = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)
        assert data_handle.base_timestep == 2015

    def test_previous_timestep(self, mock_model, mock_store):
        """should return previous timestep from list"""
        data_handle = DataHandle(mock_store, 1, 2020, [2015, 2020], mock_model)
        assert data_handle.previous_timestep == 2015

    def test_previous_timestep_error(self, mock_model, mock_store):
        """should raise error if there's no previous timestep in the list"""
        data_handle = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)
        with raises(SmifTimestepResolutionError) as ex:
            data_handle.previous_timestep
        assert "no previous timestep" in str(ex.value)


class TestDataHandleGetResults:
    def test_get_results_sectormodel(self, mock_store, mock_model):
        """Get results from a sector model"""
        store = mock_store
        spec = mock_model.outputs["gas_demand"]

        data = np.random.rand(2, 8)
        da = DataArray(spec, data)

        store.write_results(da, 1, "energy_demand", 2010)

        dh = DataHandle(mock_store, 1, 2010, [2010], mock_model)
        actual = dh.get_results("gas_demand")
        assert actual == DataArray(spec, data)

    def test_get_results_no_output_sector(self, mock_store, mock_model):

        with raises(KeyError):
            dh = DataHandle(mock_store, 1, 2010, [2010], mock_model)
            dh.get_results("no_such_output")


class TestDataHandleGetParameters:
    def test_load_parameter_defaults(self, mock_store, mock_model):

        dh = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)

        actual = dh.get_parameter("smart_meter_savings")
        spec = Spec.from_dict(
            {
                "name": "smart_meter_savings",
                "description": "Difference in floor area per person"
                "in end year compared to base year",
                "absolute_range": [0, float("inf")],
                "expected_range": [0.5, 2],
                "unit": "%",
                "dtype": "float",
            }
        )
        expected = DataArray(spec, np.array(42, dtype=float))

        assert actual == expected

    def test_load_parameters_override(self, mock_store, mock_model):

        mock_store.update_model_run(
            1,
            {
                "name": 1,
                "narratives": {"test_narrative": ["high_tech_dsm"]},
                "sos_model": "test_sos_model",
                "scenarios": {},
            },
        )
        dh = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)

        actual = dh.get_parameter("smart_meter_savings")
        spec = Spec.from_dict(
            {
                "name": "smart_meter_savings",
                "description": "Difference in floor area per person"
                "in end year compared to base year",
                "absolute_range": [0, float("inf")],
                "expected_range": [0.5, 2],
                "unit": "%",
                "dtype": "float",
            }
        )
        expected = DataArray(spec, np.array(99))

        assert actual == expected

    def test_load_parameters_override_ordered(self, mock_store, mock_model):
        """Parameters in a narrative variants listed later override parameters
        contained in earlier variants
        """
        sos_model = mock_store.read_sos_model("test_sos_model")
        sos_model["narratives"] = [
            {
                "name": "test_narrative",
                "description": "a narrative config",
                "sos_model": "test_sos_model",
                "provides": {"energy_demand": ["smart_meter_savings"]},
                "variants": [
                    {
                        "name": "first_variant",
                        "description": "This variant should be overridden",
                        "data": {"smart_meter_savings": "filename.csv"},
                    },
                    {
                        "name": "second_variant",
                        "description": "This variant should override the first",
                        "data": {"smart_meter_savings": "filename.csv"},
                    },
                ],
            }
        ]
        mock_store.update_sos_model("test_sos_model", sos_model)

        mock_store.update_model_run(
            1,
            {
                "name": 1,
                "narratives": {"test_narrative": ["first_variant", "second_variant"]},
                "sos_model": "test_sos_model",
                "scenarios": {},
            },
        )

        parameter_spec = Spec(name="smart_meter_savings", dtype="float", unit="%")
        first_variant = DataArray(parameter_spec, np.array(1))
        mock_store.write_narrative_variant_data(
            "test_sos_model", "test_narrative", "first_variant", first_variant
        )

        second_variant = DataArray(parameter_spec, np.array(2))
        mock_store.write_narrative_variant_data(
            "test_sos_model", "test_narrative", "second_variant", second_variant
        )

        dh = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)

        actual = dh.get_parameter("smart_meter_savings")

        assert actual == second_variant

    def test_load_parameters_partial_override(self, mock_store, mock_model):
        # add parameter to model
        param_spec = Spec(
            name="multi_savings",
            description="The savings from various technologies",
            abs_range=(0, 100),
            exp_range=(3, 10),
            dims=["technology_type"],
            coords={"technology_type": ["water_meter", "electricity_meter"]},
            unit="%",
            dtype="float",
        )
        mock_model.add_parameter(param_spec)
        mock_store.update_model(mock_model.name, mock_model.as_dict())

        # default values
        param_defaults = DataArray(param_spec, np.array([3, 3]))
        mock_store.write_model_parameter_default(
            mock_model.name, param_spec.name, param_defaults
        )

        # add narrative to sosmodel
        sos_model = mock_store.read_sos_model("test_sos_model")
        sos_model["narratives"] = [
            {
                "name": "test_narrative",
                "description": "a narrative config",
                "sos_model": "test_sos_model",
                "provides": {"energy_demand": ["multi_savings"]},
                "variants": [
                    {
                        "name": "low_multi_save",
                        "description": "Low values",
                        "data": {"multi_savings": "multi_savings.csv"},
                    }
                ],
            }
        ]
        mock_store.update_sos_model("test_sos_model", sos_model)

        # narrative values - one NaN i.e. not overridden
        param_narrative = DataArray(param_spec, np.array([np.nan, 99]))
        mock_store.write_narrative_variant_data(
            "test_sos_model", "test_narrative", "low_multi_save", param_narrative
        )

        # expect combination
        expected = DataArray(param_spec, np.array([3, 99]))

        mock_store.update_model_run(
            1,
            {
                "name": 1,
                "narratives": {"test_narrative": ["low_multi_save"]},
                "sos_model": "test_sos_model",
                "scenarios": {},
            },
        )

        dh = DataHandle(mock_store, 1, 2015, [2015, 2020], mock_model)
        actual = dh.get_parameter("multi_savings")

        assert actual == expected


class TestDataHandleCoefficients:
    """Tests the interface for reading and writing coefficients"""

    def test_read_coefficient_raises(self, mock_store, mock_model):
        """Reading coefficients that aren't cached should raise SmifDataNotFoundError"""
        store = mock_store
        dh = DataHandle(store, 1, 2010, [2010], mock_model)

        with raises(SmifDataNotFoundError):
            dh.read_coefficients("from_dim", "to_dim")

    def test_write_coefficients(self, mock_store, mock_model):
        """Writing coefficients should succeed"""
        store = mock_store
        data = np.zeros(4)

        dh = DataHandle(store, 1, 2010, [2010], mock_model)
        dh.write_coefficients("from_dim", "to_dim", data)

    def test_read_coefficient_(self, mock_store, mock_model):
        """Reading written coefficients should return the cached data"""
        store = mock_store
        dh = DataHandle(store, 1, 2010, [2010], mock_model)
        data = np.zeros(4)
        dh.write_coefficients("from_dim", "to_dim", data)

        actual = dh.read_coefficients("from_dim", "to_dim")
        np.testing.assert_equal(actual, np.array([0, 0, 0, 0]))

    def test_read_unit_definitions(self, mock_store, mock_model):
        """Reading list of unit definition strings"""
        store = mock_store
        dh = DataHandle(store, 1, 2010, [2010], mock_model)
        actual = dh.read_unit_definitions()
        assert actual == ["people", "million people = 1000000 * people"]


class TestResultsHandle:
    """Get results from any model"""

    def test_get_state(self, mock_store, mock_model, mock_sos_model):
        """should get decision module state for given timestep/decision_iteration

        A call to ``get_state`` method on the data handle calls the read_state
        method of the store with arguments for model run name, current
        timestep and decision iteration.
        """
        mock_store.read_state = Mock(
            return_value=[{"name": "test", "build_year": 2010}]
        )
        mock_store.write_interventions(
            "energy_demand",
            [{"name": "test", "capital_cost": {"value": 2500, "unit": "£/GW"}}],
        )
        data_handle = ResultsHandle(
            mock_store, 1, mock_sos_model, 2015, timesteps=[2015, 2020]
        )
        expected = [{"name": "test", "build_year": 2010}]
        actual = data_handle.get_state(2015, None)
        mock_store.read_state.assert_called_with(1, 2015, None)
        assert actual == expected

    def test_get_results_sos_model(self, mock_store, mock_model, mock_sos_model):
        """Get results from a sector model within a sos model"""
        store = mock_store
        spec = mock_model.outputs["gas_demand"]

        da = DataArray(spec, np.random.rand(2, 8))

        store.write_results(da, "test_modelrun", "energy_demand", 2010, None)

        dh = ResultsHandle(mock_store, "test_modelrun", mock_sos_model, 2010)
        actual = dh.get_results("energy_demand", "gas_demand", 2010, None)
        spec = mock_model.outputs["gas_demand"]
        assert actual == da

    def test_get_results_no_output_sos(self, mock_store, mock_sos_model):
        with raises(KeyError):
            dh = ResultsHandle(mock_store, "test_modelrun", mock_sos_model, 2010)
            dh.get_results("energy_demand", "no_such_output", 2010, None)

    def test_get_results_wrong_name_sos(self, mock_store, mock_sos_model):
        with raises(KeyError):
            dh = ResultsHandle(mock_store, "test_modelrun", mock_sos_model, 2010)
            dh.get_results("no_such_model", "test_output", 2010, None)

    def test_get_results_not_exists(self, mock_store, mock_sos_model, mock_model):
        store = mock_store

        spec = mock_model.outputs["gas_demand"]
        da = DataArray(spec, np.random.rand(2, 8))

        store.write_results(da, "test_modelrun", "energy_demand", 2010, None)
        dh = ResultsHandle(store, "test_modelrun", mock_sos_model, 2100)
        with raises(SmifDataError):
            dh.get_results("energy_demand", "gas_demand", 2099, None)
