"""
Defines a class storing data in memory. 
   
.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de> 
"""

from __future__ import annotations

from contextlib import contextmanager
from typing import List, Optional, Sequence

import numpy as np

from ..fields import FieldCollection
from ..fields.base import FieldBase
from .base import InfoDict, StorageBase, WriteModeType


class MemoryStorage(StorageBase):
    """store discretized fields in memory"""

    def __init__(
        self,
        times: Optional[Sequence[float]] = None,
        data: Optional[List[np.ndarray]] = None,
        *,
        info: Optional[InfoDict] = None,
        field_obj: Optional[FieldBase] = None,
        write_mode: WriteModeType = "truncate_once",
    ):
        """
        Args:
            times (:class:`~numpy.ndarray`):
                Sequence of times for which data is known
            data (list of :class:`~numpy.ndarray`):
                The field data at the given times
            field_obj (:class:`~pde.fields.base.FieldBase`):
                An instance of the field class store data for a single time point.
            info (dict):
                Supplies extra information that is stored in the storage
            write_mode (str):
                Determines how new data is added to already existing data. Possible
                values are: 'append' (data is always appended), 'truncate' (data is
                cleared every time this storage is used for writing), or 'truncate_once'
                (data is cleared for the first writing, but appended subsequently).
                Alternatively, specifying 'readonly' will disable writing completely.
        """
        super().__init__(info=info, write_mode=write_mode)
        self.times: List[float] = [] if times is None else list(times)
        if field_obj is not None:
            self._field = field_obj.copy()
            self._grid = field_obj.grid
            self._data_shape = field_obj.data.shape

        self.data: List[np.ndarray] = [] if data is None else data
        if self._data_shape is None and len(self.data) > 0:
            self._data_shape = self.data[0].shape

        # check consistency
        if len(self.times) != len(self.data):
            raise ValueError(
                "Length of the supplied `times` and `fields` are inconsistent "
                f"({len(self.times)} != {len(self.data)})"
            )

    @classmethod
    def from_fields(
        cls,
        times: Optional[Sequence[float]] = None,
        fields: Optional[Sequence[FieldBase]] = None,
        info: Optional[InfoDict] = None,
        write_mode: WriteModeType = "truncate_once",
    ) -> MemoryStorage:
        """create MemoryStorage from a list of fields

        Args:
            times (:class:`~numpy.ndarray`):
                Sequence of times for which data is known
            fields (list of :class:`~pde.fields.FieldBase`):
                The fields at all given time points
            info (dict):
                Supplies extra information that is stored in the storage
            write_mode (str):
                Determines how new data is added to already existing data. Possible
                values are: 'append' (data is always appended), 'truncate' (data is
                cleared every time this storage is used for writing), or 'truncate_once'
                (data is cleared for the first writing, but appended subsequently).
                Alternatively, specifying 'readonly' will disable writing completely.
        """
        if fields is None:
            field_obj = None
            data = None
        else:
            field_obj = fields[0]
            data = [fields[0].data]
            for field in fields[1:]:
                if field_obj.grid != field.grid:
                    raise ValueError("Grids of the fields are incompatible")
                data.append(field.data)

        return cls(
            times, data=data, field_obj=field_obj, info=info, write_mode=write_mode
        )

    @classmethod
    def from_collection(
        cls,
        storages: Sequence[StorageBase],
        label: Optional[str] = None,
        *,
        rtol: float = 1.0e-5,
        atol: float = 1.0e-8,
    ) -> MemoryStorage:
        """combine multiple memory storages into one

        This method can be used to combine multiple time series of different fields into
        a single representation. This requires that all time series contain data at the
        same time points.

        Args:
            storages (list):
                A collection of instances of :class:`~pde.storage.base.StorageBase`
                whose data will be concatenated into a single MemoryStorage
            label (str, optional):
                The label of the instances of :class:`~pde.fields.FieldCollection` that
                represent the concatenated data
            rtol (float):
                Relative tolerance used when checking times for merging
            atol (float):
                Absolute tolerance used when checking times for merging

        Returns:
            :class:`~pde.storage.memory.MemoryStorage`: Storage containing all the data.
        """
        if len(storages) == 0:
            return cls()

        # initialize the combined data
        times = storages[0].times
        data = [[field] for field in storages[0]]

        # append data from further storages
        for storage in storages[1:]:
            if not np.allclose(times, storage.times, rtol=rtol, atol=atol):
                raise ValueError("Storages have incompatible times")
            for i, field in enumerate(storage):
                data[i].append(field)

        # convert data format to FieldCollections
        fields = [FieldCollection(d, label=label) for d in data]  # type: ignore

        return cls.from_fields(times, fields=fields)

    def clear(self, clear_data_shape: bool = False) -> None:
        """truncate the storage by removing all stored data.

        Args:
            clear_data_shape (bool):
                Flag determining whether the data shape is also deleted.
        """
        self.times = []
        self.data = []
        super().clear(clear_data_shape=clear_data_shape)

    def start_writing(self, field: FieldBase, info: Optional[InfoDict] = None) -> None:
        """initialize the storage for writing data

        Args:
            field (:class:`~pde.fields.FieldBase`):
                An instance of the field class store data for a single time point.
            info (dict):
                Supplies extra information that is stored in the storage
        """
        super().start_writing(field, info=info)

        # update info after opening file because otherwise information can be
        # overwritten by data that is already present in the file
        if info is not None:
            self.info.update(info)

        # handle the different write modes
        if self.write_mode == "truncate_once":
            self.clear()
            self.write_mode = "append"  # do not truncate in subsequent calls

        elif self.write_mode == "truncate":
            self.clear()

        elif self.write_mode == "readonly":
            raise RuntimeError("Cannot write in read-only mode")

        elif self.write_mode != "append":
            raise ValueError(
                f"Unknown write mode `{self.write_mode}`. Possible values are "
                "`truncate_once`, `truncate`, and `append`"
            )

    def _append_data(self, data: np.ndarray, time: float) -> None:
        """append a new data set

        Args:
            data (:class:`~numpy.ndarray`): The actual data
            time (float, optional): The time point associated with the data
        """
        assert data.shape == self.data_shape, f"Data must have shape {self.data_shape}"
        self.data.append(np.array(data))  # store copy of the data
        self.times.append(time)


@contextmanager
def get_memory_storage(field: FieldBase, info: Optional[InfoDict] = None):
    """a context manager that can be used to create a MemoryStorage

    Example:
        This can be used to quickly store data::

            with get_memory_storage(field_class) as storage:
                storage.append(numpy_array0, 0)
                storage.append(numpy_array1, 1)

            # use storage thereafter

    Args:
        field (:class:`~pde.fields.FieldBase`):
            An instance of the field class store data for a single time point.
        info (dict):
            Supplies extra information that is stored in the storage

    Yields:
        :class:`MemoryStorage`
    """
    storage = MemoryStorage()
    storage.start_writing(field, info)
    yield storage
    storage.end_writing()
