"""
    class for loading the data from the HDF5 output file from S^3 and assembling the data matrix
"""
import logging
import torch as pt

from h5py import File
from os.path import join
from typing import Union, List, Dict

from .dataloader import Dataloader

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

CONST = "constant"
GRID = "grid"
DATA = "data"
FACES = "faces"
CENTERS = "centers"
VERTICES = "vertices"


class SCUBEDataloader(Dataloader):
    """
    class for loading the data from the HDF5 output file from S^3 and assembling the data matrix
    """
    def __init__(self, load_path: str, file_name: str, dtype: pt.dtype = pt.float32):
        """
        implements a DataLoader class for conveniently loading the HDF5 file generated by S^3

        :param load_path: path to the HDF5 file
        :type load_path: str
        :param file_name: name of the HDF5 file
        :type file_name: str
        :param dtype: the type of the tensor, default is set to float32 (single precision)
        :type dtype: pt.dtype
        """
        # create a file object
        self._load_path = load_path
        self._file_name = file_name
        self._dtype = dtype

        # load some values we may need, we always have to close the file after reading, otherwise we can't access it
        # later for writing (or sth. else)
        with File(join(self._load_path, self._file_name), "r") as f:
            self._n_cells = f.get(f"{GRID}/{CENTERS}")[()].shape[0]
            self._n_dimensions = f.get(f"{GRID}/{CENTERS}")[()].shape[1]
            self._size_initial_cell = f.get(f"{CONST}/size_initial_cell")[()]

        # placeholders for properties which are only loaded if requested
        self._write_times = None
        self._weights = None  # cell areas (2D) / volumes (3D)
        self._levels = None
        self._field_names = None

        # properties of the grid
        self._vertices = None
        self._faces = None
        self._nodes = None

    @property
    def write_times(self) -> List[str]:
        """
        load all available time steps present in the HDF5 file, the time steps at which a specific field is present are
        stored in the 'field_names' property.

        :return: available write times
        :rtype list
        """
        if self._write_times is None:
            with File(join(self._load_path, self._file_name), "r") as f:
                if DATA in f.keys():
                    self._write_times = list(f.get(f"{DATA}").keys())

        return self._write_times

    @property
    def weights(self) -> pt.Tensor:
        """
        load and compute the cell areas (2d) / volumes (3D) of the grid

        :return: cell areas / volumes
        :rtype: pt.Tensor
        """
        if self._weights is None:
            self._compute_cell_area()

        return self._weights

    @property
    def vertices(self) -> pt.Tensor:
        """
        load the cell centers of the grid

        :return: cell centers of the grid
        :rtype: pt.Tensor
        """
        if self._vertices is None:
            with File(join(self._load_path, self._file_name), "r") as f:
                self._vertices = pt.from_numpy(f.get(f"{GRID}/{CENTERS}")[()])

        return self._vertices

    @property
    def nodes(self) -> pt.Tensor:
        """
        load the cell nodes of the grid

        :return: cell nodes of the grid
        :rtype: pt.Tensor
        """
        if self._nodes is None:
            with File(join(self._load_path, self._file_name), "r") as f:
                self._nodes = pt.from_numpy(f.get(f"{GRID}/{VERTICES}")[()])

        return self._nodes

    @property
    def faces(self) -> pt.Tensor:
        """
        load the cell faces of the grid

        :return: idx cell faces of the grid
        :rtype: pt.Tensor
        """
        if self._faces is None:
            with File(join(self._load_path, self._file_name), "r") as f:
                self._faces = pt.from_numpy(f.get(f"{GRID}/{FACES}")[()])

        return self._faces

    @property
    def field_names(self) -> dict:
        """
        create a dictionary containing the available fields as a list for each available time step, the time steps are
        the keys while the available fields are the values for each key

        :return: dict with available fields for each available time step
        :rtype: Dict[list]
        """
        if self._field_names is None:
            # we can't use CENTERS, because CENTERS only stands for the cell centers of the grid; the location of the
            # field is marked as center or vertices
            with File(join(self._load_path, self._file_name), "r") as f:
                self._field_names = {k: [f.split("_")[0] for f in f[f"{DATA}/{k}"].keys() if f.endswith("center")]
                                     for k in f[DATA].keys()}

        return self._field_names

    @property
    def levels(self) -> pt.Tensor:
        """
        load the cell levels of the grid

        :return: cell levels of the grid
        :rtype: pt.Tensor
        """
        if self._levels is None:
            with File(join(self._load_path, self._file_name), "r") as f:
                self._levels = pt.from_numpy(f.get(f"{CONST}/levels")[()]).squeeze()

        return self._levels

    @property
    def load_path(self) -> str:
        """
        get the current load path

        :return: current load path
        :rtype: str
        """
        return self._load_path

    @load_path.setter
    def load_path(self, value: str) -> None:
        """
        Set a new load path, this automatically resets the properties to avoid inconsistencies

        :return: None
        """
        self._load_path = value
        self._reset()

    @property
    def file_name(self) -> str:
        """
        get the current file name

        :return: current file name
        :rtype: str
        """
        return self._file_name

    @file_name.setter
    def file_name(self, value: str) -> None:
        """
        Set a new file name, this automatically resets the properties to avoid inconsistencies

        :return: None
        """
        self._file_name = value
        self._reset()

    def _reset(self) -> None:
        """
        Reset the properties of the class to avoid inconsistencies

        :return: None
        """
        with File(join(self._load_path, self._file_name), "r") as f:
            self._n_cells = f.get(f"{GRID}/{CENTERS}")[()].shape[0]
            self._n_dimensions = f.get(f"{GRID}/{CENTERS}")[()].shape[1]
            self._size_initial_cell = f.get(f"{CONST}/size_initial_cell")[()]
        self._write_times = None
        self._weights = None
        self._levels = None
        self._field_names = None
        self._vertices = None

    def _compute_cell_area(self) -> None:
        """
        Compute the cell areas (2d) / volumes (3D) of the grid

        :return: None
        """
        self._weights = (1 / pow(2, self._n_dimensions) * pow(self._size_initial_cell / pow(2, self.levels),
                                                              self._n_dimensions)).squeeze()

    def load_snapshot(self, field_name: Union[List[str], str],
                      write_times: Union[str, List[str]] = None) -> Union[List[pt.Tensor], pt.Tensor]:
        """
        load the snapshots for the given field(s) and time steps

        :param field_name: field name for which the data matrix should be created
        :type field_name: str
        :param write_times: time steps for which the data matrix should be created, if 'None' then all available time
                            steps will be used
        :type write_times: Union[str, List[str]]
        :return: either data matrix containing the snapshots as: [N_cells, N_dimensions, N_snapshots] (vector field) and
                 [N_cells, N_snapshots] (scalar field) or list with data matrices (if multiple fields should be loaded)
        :rtype: Union[List[pt.Tensor], pt.Tensor]
        """
        if write_times is None:
            write_times = self.write_times

        if type(write_times) is str:
            write_times = [write_times]
        if type(field_name) is str:
            field_name = [field_name]

        # create fiel object and empty list for storing the data matrices
        _file = File(join(self._load_path, self._file_name), "r")
        _dm = []

        for f in field_name:
            # get the shape of the field (scalar / vector field)
            shape = _file.get(f"{DATA}/{write_times[0]}/{f}_center")[()].shape

            # allocate the data matrix accordingly
            if len(shape) == 1:
                _data_matrix = pt.zeros((self._n_cells, len(write_times)), dtype=self._dtype)
            else:
                _data_matrix = pt.zeros((shape[0], shape[1], len(write_times)), dtype=self._dtype)

            # assemble the data matrix for the given time steps
            for i, k in enumerate(write_times):
                if len(shape) == 1:
                    _data_matrix[:, i] = pt.from_numpy(_file.get(f"{DATA}/{k}/{f}_center")[()])
                else:
                    _data_matrix[:, :, i] = pt.from_numpy(_file.get(f"{DATA}/{k}/{f}_center")[()])

            _dm.append(_data_matrix)

        # close the file
        _file.close()

        return _dm[0] if len(_dm) == 1 else _dm


if __name__ == "__main__":
    pass
