"""
This module provides a wrapper around a Cactus simulation output, i.e. a directory of HDF5 files.
"""

import os
from collections.abc import MutableMapping

from . import datafile


class _RefinedSlice(object):
    # simulation time on this slice
    time = -1
    # simulation iteration on this slice
    it   = -1

    # a list of arrays of coordinate points along each direction [x, y, z]
    c = None

    # arrays of coordinate points along each direction
    x = None
    y = None
    z = None

    rl = None

    # private
    _sd = None

    def __init__(self, sd, rl, time = -1, iteration = -1):
        self._sd = sd
        self.rl  = rl

        if time < 0 and iteration < 0:
            raise TypeError('Neither time nor iteration provided')

        trial_df = sd._df

        if time >= 0:
            self.it = trial_df._iter_from_time(time)
        else:
            self.it = iteration

        try:
            s = trial_df.slice(it = self.it, rl = rl.n)
            self.time = s.time
        except KeyError:
            if time < 0:
                raise IndexError('No such iteration: ' + str(iteration))
            else:
                raise IndexError('No such time: ' + str(time))

        ds = trial_df.slice(self.it, rl = rl.n)
        self.c = ds.layout.coords
        self.x, self.y, self.z = self.c

    def __getitem__(self, key):
        if key in self._sd.df:
            return self._sd.df[key].slice(self.it, rl = self.rl.n)
        raise IndexError


class _RefinementLevel(object):
    # refinement level number
    n  = -1

    # private
    _sd = None

    def __init__(self, sd, n):
        self._sd = sd
        self.n   = n

    def slice(self, iteration = -1, time = -1):
        return _RefinedSlice(self._sd, self, time, iteration)

    @property
    def itertimes(self):
        return self._sd._df.itertimes[self.n]

class _DataDir(MutableMapping):
    _files = None
    _paths = None

    def __init__(self):
        self._files = {}
        self._paths = {}

    def __setitem__(self, key, val):
        if key in self._files:
            raise ValueError

        self._files[key] = None
        self._paths[key] = val

    def __delitem__(self, key):
        if self._files[key]:
            self._files[key].close()
        del self._files[key]
        del self._paths[key]

    def __getitem__(self, key):
        if not key in self._paths:
            raise KeyError(key)
        elif self._files[key] is None:
            self._files[key] = datafile.DataFile(self._paths[key])

        return self._files[key]

    def __len__(self):
        return len(self._paths)

    def __iter__(self):
        for it in self._paths:
            yield it

    def close(self):
        for f in self._files.values():
            if f is not None:
                f.close()
        self._files.clear()
        self._paths.clear()

class SimulationData(object):
    "The full path to the simulation directory"
    dirname = None

    "A dictionary of { basename : DataFile } items representing the data files. The basename is the file name without the .h5 extension"
    df = None

    # per-refinement level parameters
    rl = None

    _df = None

    def __init__(self, dirname):
        self.dirname = os.path.abspath(dirname)
        self.df = _DataDir()

        # open all the hdf5 files in the dir
        for f in os.listdir(self.dirname):
            if (not f.endswith('.h5') or f.startswith('checkpoint') or f.endswith('.d.h5')):
                continue

            self.df[f[:-3]] = '%s/%s' % (self.dirname, f)
        if len(self.df) == 0:
            raise ValueError('No HDF5 data files in the directory.')

        # pick a representative datafile and get the grid properties from it
        if 'H' in self.df:
            df = self.df['H']
        else:
            df = self.df.values()[0]

        self._df = df

        # get the refinement levels, iterations and times
        self.rl = []
        for rl in df.rl:
            self.rl.append(_RefinementLevel(self, len(self.rl)))

    def close(self):
        self.df.close()