<commit_before>"""File-backed data store
"""
import glob
import os
from abc import abstractmethod
from logging import getLogger

import numpy as np  # type: ignore
import pandas  # type: ignore
import pyarrow as pa  # type: ignore
from smif.data_layer.abstract_data_store import DataStore
from smif.data_layer.data_array import DataArray
from smif.exception import SmifDataMismatchError, SmifDataNotFoundError


class FileDataStore(DataStore):
    """Abstract file data store
    """
    def __init__(self, base_folder):
        super().__init__()
        self.logger = getLogger(__name__)
        # extension for DataArray/list-of-dict data - override in implementations
        self.ext = ''
        # extension for bare numpy.ndarray data - override in implementations
        self.coef_ext = ''

        self.base_folder = str(base_folder)
        self.data_folder = str(os.path.join(self.base_folder, 'data'))
        self.data_folders = {}
        self.results_folder = str(os.path.join(self.base_folder, 'results'))
        data_folders = [
            'coefficients',
            'strategies',
            'initial_conditions',
            'interventions',
            'narratives',
            'scenarios',
            'strategies',
            'parameters'
        ]
        for folder in data_folders:
            dirname = os.path.join(self.data_folder, folder)
            # ensure each directory exists
            if not os.path.exists(dirname):
                msg = "Expected data folder at '{}' but it does not exist"
                abs_path = os.path.abspath(dirname)
                raise SmifDataNotFoundError(msg.format(abs_path))
            self.data_folders[folder] = dirname

    # region Abstract methods
    @abstractmethod
    def _read_data_array(self, path, spec, timestep=None):
        """Read DataArray from file
        """

    @abstractmethod
    def _write_data_array(self, path, data_array, timestep=None):
        """Write DataArray to file
        """

    @abstractmethod
    def _read_list_of_dicts(self, path):
        """Read file to list[dict]
        """

    @abstractmethod
    def _write_list_of_dicts(self, path, data):
        """Write list[dict] to file
        """

    @abstractmethod
    def _read_ndarray(self, path):
        """Read numpy.ndarray
        """

    @abstractmethod
    def _write_ndarray(self, path, data, header=None):
        """Write numpy.ndarray
        """
    # endregion

    # region Data Array
    def read_scenario_variant_data(self, key, spec, timestep=None):
        path = os.path.join(self.data_folders['scenarios'], key)
        data = self._read_data_array(path, spec, timestep)
        data.validate_as_full()
        return data

    def write_scenario_variant_data(self, key, data, timestep=None):
        path = os.path.join(self.data_folders['scenarios'], key)
        self._write_data_array(path, data, timestep)

    def read_narrative_variant_data(self, key, spec, timestep=None):
        path = os.path.join(self.data_folders['narratives'], key)
        return self._read_data_array(path, spec, timestep)

    def write_narrative_variant_data(self, key, data, timestep=None):
        path = os.path.join(self.data_folders['narratives'], key)
        path = self._set_file_extension(path)
        self._write_data_array(path, data, timestep)

    def read_model_parameter_default(self, key, spec):
        self.logger.debug("Trying to read model parameter default from key {}".format(key))
        path = os.path.join(self.data_folders['parameters'], key)
        data = self._read_data_array(path, spec)
        data.validate_as_full()
        return data

    def write_model_parameter_default(self, key, data):
        path = os.path.join(self.data_folders['parameters'], key)
        path = self._set_file_extension(path)
        self._write_data_array(path, data)
    # endregion

    # region Interventions
    def read_interventions(self, keys):
        all_interventions = []
        for key in keys:
            path = os.path.join(self.data_folders['interventions'], key)
            interventions = self._read_list_of_dicts(path)
            all_interventions.extend(interventions)

        seen = set()
        dups = set()

        for intervention in all_interventions:
            try:
                name = intervention['name']
            except KeyError:
                msg = "Could not find `name` key in {} for {}"
                raise KeyError(msg.format(intervention, keys))
            if name in seen:
                dups.add(name)
            else:
                seen.add(name)

        if dups:
            name = dups.pop()
            msg = "An entry for intervention {} already exists. Also found duplicates for {}"
            raise ValueError(msg.format(name, dups))

        return {
            intervention['name']: _nest_keys(intervention)
            for intervention in all_interventions
        }

    def write_interventions(self, key, interventions):
        # convert dict[str, dict] to list[dict]
        data = [
            _unnest_keys(intervention)
            for intervention in interventions.values()
        ]
        path = os.path.join(self.data_folders['interventions'], key)
        self._write_list_of_dicts(path, data)

    def read_strategy_interventions(self, strategy):
        path = os.path.join(self.data_folders['strategies'], strategy['filename'])
        return self._read_list_of_dicts(path)

    def write_strategy_interventions(self, strategy, data):
        path = os.path.join(self.data_folders['strategies'], strategy['filename'])
        path = self._set_file_extension(path)
        return self._write_list_of_dicts(path, data)

    def read_initial_conditions(self, keys):
        conditions = []
        for key in keys:
            path = os.path.join(self.data_folder, 'initial_conditions', key)
            data = self._read_list_of_dicts(path)
            conditions.extend(data)
        return conditions

    def write_initial_conditions(self, key, initial_conditions):
        path = os.path.join(self.data_folders['initial_conditions'], key)
        path = self._set_file_extension(path)
        self._write_list_of_dicts(path, initial_conditions)
    # endregion

    # region State
    def read_state(self, modelrun_name, timestep, decision_iteration=None):
        path = self._get_state_path(modelrun_name, timestep, decision_iteration)
        try:
            state = self._read_list_of_dicts(path)
        except FileNotFoundError:<commit_msg>Remove calls to _set_file_extension<commit_after>"""File-backed data store
"""
import glob
import os
from abc import abstractmethod
from logging import getLogger

import numpy as np  # type: ignore
import pandas  # type: ignore
import pyarrow as pa  # type: ignore
from smif.data_layer.abstract_data_store import DataStore
from smif.data_layer.data_array import DataArray
from smif.exception import SmifDataMismatchError, SmifDataNotFoundError


class FileDataStore(DataStore):
    """Abstract file data store
    """
    def __init__(self, base_folder):
        super().__init__()
        self.logger = getLogger(__name__)
        # extension for DataArray/list-of-dict data - override in implementations
        self.ext = ''
        # extension for bare numpy.ndarray data - override in implementations
        self.coef_ext = ''

        self.base_folder = str(base_folder)
        self.data_folder = str(os.path.join(self.base_folder, 'data'))
        self.data_folders = {}
        self.results_folder = str(os.path.join(self.base_folder, 'results'))
        data_folders = [
            'coefficients',
            'strategies',
            'initial_conditions',
            'interventions',
            'narratives',
            'scenarios',
            'strategies',
            'parameters'
        ]
        for folder in data_folders:
            dirname = os.path.join(self.data_folder, folder)
            # ensure each directory exists
            if not os.path.exists(dirname):
                msg = "Expected data folder at '{}' but it does not exist"
                abs_path = os.path.abspath(dirname)
                raise SmifDataNotFoundError(msg.format(abs_path))
            self.data_folders[folder] = dirname

    # region Abstract methods
    @abstractmethod
    def _read_data_array(self, path, spec, timestep=None):
        """Read DataArray from file
        """

    @abstractmethod
    def _write_data_array(self, path, data_array, timestep=None):
        """Write DataArray to file
        """

    @abstractmethod
    def _read_list_of_dicts(self, path):
        """Read file to list[dict]
        """

    @abstractmethod
    def _write_list_of_dicts(self, path, data):
        """Write list[dict] to file
        """

    @abstractmethod
    def _read_ndarray(self, path):
        """Read numpy.ndarray
        """

    @abstractmethod
    def _write_ndarray(self, path, data, header=None):
        """Write numpy.ndarray
        """
    # endregion

    # region Data Array
    def read_scenario_variant_data(self, key, spec, timestep=None):
        path = os.path.join(self.data_folders['scenarios'], key)
        data = self._read_data_array(path, spec, timestep)
        data.validate_as_full()
        return data

    def write_scenario_variant_data(self, key, data, timestep=None):
        path = os.path.join(self.data_folders['scenarios'], key)
        self._write_data_array(path, data, timestep)

    def read_narrative_variant_data(self, key, spec, timestep=None):
        path = os.path.join(self.data_folders['narratives'], key)
        return self._read_data_array(path, spec, timestep)

    def write_narrative_variant_data(self, key, data, timestep=None):
        path = os.path.join(self.data_folders['narratives'], key)
        self._write_data_array(path, data, timestep)

    def read_model_parameter_default(self, key, spec):
        self.logger.debug("Trying to read model parameter default from key {}".format(key))
        path = os.path.join(self.data_folders['parameters'], key)
        data = self._read_data_array(path, spec)
        data.validate_as_full()
        return data

    def write_model_parameter_default(self, key, data):
        path = os.path.join(self.data_folders['parameters'], key)
        self._write_data_array(path, data)
    # endregion

    # region Interventions
    def read_interventions(self, keys):
        all_interventions = []
        for key in keys:
            path = os.path.join(self.data_folders['interventions'], key)
            interventions = self._read_list_of_dicts(path)
            all_interventions.extend(interventions)

        seen = set()
        dups = set()

        for intervention in all_interventions:
            try:
                name = intervention['name']
            except KeyError:
                msg = "Could not find `name` key in {} for {}"
                raise KeyError(msg.format(intervention, keys))
            if name in seen:
                dups.add(name)
            else:
                seen.add(name)

        if dups:
            name = dups.pop()
            msg = "An entry for intervention {} already exists. Also found duplicates for {}"
            raise ValueError(msg.format(name, dups))

        return {
            intervention['name']: _nest_keys(intervention)
            for intervention in all_interventions
        }

    def write_interventions(self, key, interventions):
        # convert dict[str, dict] to list[dict]
        data = [
            _unnest_keys(intervention)
            for intervention in interventions.values()
        ]
        path = os.path.join(self.data_folders['interventions'], key)
        self._write_list_of_dicts(path, data)

    def read_strategy_interventions(self, strategy):
        path = os.path.join(self.data_folders['strategies'], strategy['filename'])
        return self._read_list_of_dicts(path)

    def write_strategy_interventions(self, strategy, data):
        path = os.path.join(self.data_folders['strategies'], strategy['filename'])
        return self._write_list_of_dicts(path, data)

    def read_initial_conditions(self, keys):
        conditions = []
        for key in keys:
            path = os.path.join(self.data_folder, 'initial_conditions', key)
            data = self._read_list_of_dicts(path)
            conditions.extend(data)
        return conditions

    def write_initial_conditions(self, key, initial_conditions):
        path = os.path.join(self.data_folders['initial_conditions'], key)
        self._write_list_of_dicts(path, initial_conditions)
    # endregion

    # region State
    def read_state(self, modelrun_name, timestep, decision_iteration=None):
        path = self._get_state_path(modelrun_name, timestep, decision_iteration)
        try:
            state = self._read_list_of_dicts(path)
        except FileNotFoundError:
{'fraction_non_alphanumeric': 0.0984977543750968, 'fraction_numerical': 0.0, 'mean_word_length': 2.5637417218543046, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4430593', 'n_tokens_mistral': 3805, 'n_tokens_neox': 3712, 'n_words': 836}