"""
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 functools import cached_property

import numpy as np

from math_utils.array_utils import array_reflect

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 _AxisData:
    """
    Wrapper providing convenient access to quantities along a given axis
    as functions of time. Data is exported as (t, x) 2D arrays.


    """

    sd        = None
    "parent simulation data"

    axis      = None
    "axis along which the data is loaded, one of: x/y/z"
    axis_idx  = None
    "axis index: 0/1/2 for x/y/z"

    rl        = None
    "refinement level number"

    it        = None
    "1D array of simulation iterations on which data is available"

    t         = None
    "1D array of coordinate times, same shape as it"

    x         = None
    "1D array of spatial coordinates on the axis"

    T         = None
    "2D array of times corresponding to returned data"

    X         = None
    "2D array of spatial coordinates corresponding to returned data"

    _reflect    = None
    _zero_shift = None
    _data_sl    = None

    def __init__(self, sd, rl, *,
                 axis = 'x', maxtime = None, reflect = False,
                 zero_shift = None):
        self.sd = sd
        self.rl = sd.rl[rl].n

        if axis == 'x':
            self.axis_idx = 0
        elif axis == 'y':
            self.axis_idx = 1
        elif axis == 'z':
            self.axis_idx = 2
        else:
            raise ValueError('axis must be one of: x y z')
        self.axis      = axis

        self._reflect    = reflect
        self._zero_shift = zero_shift

        df_name = os.path.splitext(os.path.basename(sd._df.path))[0].split('.')[0]
        df      = sd.df['%s.%s' % (df_name, axis)]

        itertimes  = df.itertimes[rl]
        if maxtime is not None:
            itertimes = list(filter(lambda it: it[1] <= maxtime, itertimes))
        if len(itertimes) < 1:
            raise ValueError('No data for maxtime: %s' % maxtime)

        self.it, self.t = map(np.array, zip(*itertimes))

        s      = df.slice(it = self.it[0], rl = self.rl)
        x      = s.layout.coords[0]

        ghosts        = np.nonzero(np.abs(x) < 1e-10)[0][0]
        self._data_sl = slice(ghosts, -ghosts)

        x = x[self._data_sl]
        if reflect:
            x = array_reflect(x, parity = -1.0)

        self.x = x

        self.T, self.X = np.meshgrid(self.t, self.x, indexing = 'ij')

    def _load_data(self, gf, parity):
        df = self.sd.df['%s.%s' % (gf, self.axis)]

        ret = []
        for it in self.it:
            d = df.slice(it, rl = self.rl).data[self._data_sl]
            if self._reflect:
                d = array_reflect(d, parity = parity)

            ret.append(d)
        return np.array(ret)

    # scalars
    @cached_property
    def alpha(self):    return self._load_data('alpha', 1.0)
    @cached_property
    def W(self):        return self._load_data('W', 1.0)
    @cached_property
    def phi(self):      return self._load_data('phi', 1.0)
    @cached_property
    def trK(self):      return self._load_data('trK', 1.0)
    @cached_property
    def zeta(self):     return self._load_data('zeta', 1.0)
    @cached_property
    def Kretsch(self):  return self._load_data('Kretsch', 1.0)
    @cached_property
    def H(self):        return self._load_data('H', 1.0)

    # vectors
    @cached_property
    def betax(self):
        if self._zero_shift:
            return np.zeros_like(self.T)
        return self._load_data('beta1', -1.0 if self.axis == 'x' else 1.0)
    @cached_property
    def betay(self):
        if self._zero_shift:
            return np.zeros_like(self.T)
        return self._load_data('beta2', -1.0 if self.axis == 'y' else 1.0)
    @cached_property
    def betaz(self):
        if self._zero_shift:
            return np.zeros_like(self.T)
        return self._load_data('beta3', -1.0 if self.axis == 'z' else 1.0)

    @cached_property
    def conf_fact(self):      return 1.0 / (self.phi * self.phi)

    # 2-tensors
    @cached_property
    def gtxx(self):   return self._load_data('gt11', 1.0)
    @cached_property
    def gtyy(self):   return self._load_data('gt22', 1.0)
    @cached_property
    def gtzz(self):   return self._load_data('gt33', 1.0)
    @cached_property
    def Atxx(self):   return self._load_data('At11', 1.0)
    @cached_property
    def Atyy(self):   return self._load_data('At22', 1.0)
    @cached_property
    def Atzz(self):   return self._load_data('At33', 1.0)

    @cached_property
    def gtxy(self):
        return self._load_data('gt12', 1.0 if self.axis == 'z' else -1.0)
    @cached_property
    def gtxz(self):
        return self._load_data('gt13', 1.0 if self.axis == 'y' else -1.0)
    @cached_property
    def gtyz(self):
        return self._load_data('gt23', 1.0 if self.axis == 'x' else -1.0)
    @cached_property
    def Atxy(self):
        return self._load_data('At12', 1.0 if self.axis == 'z' else -1.0)
    @cached_property
    def Atxz(self):
        return self._load_data('At13', 1.0 if self.axis == 'y' else -1.0)
    @cached_property
    def Atyz(self):
        return self._load_data('At23', 1.0 if self.axis == 'x' else -1.0)

    @cached_property
    def gxx(self):   return self.gtxx * self.conf_fact
    @cached_property
    def gyy(self):   return self.gtyy * self.conf_fact
    @cached_property
    def gzz(self):   return self.gtzz * self.conf_fact
    @cached_property
    def Kxx(self):   return self.Atxx * self.conf_fact + self.gxx * self.trK / 3.
    @cached_property
    def Kyy(self):   return self.Atyy * self.conf_fact + self.gyy * self.trK / 3.
    @cached_property
    def Kzz(self):   return self.Atzz * self.conf_fact + self.gzz * self.trK / 3.

    @cached_property
    def gxy(self):   return self.gtxy * self.conf_fact
    @cached_property
    def gxz(self):   return self.gtxz * self.conf_fact
    @cached_property
    def gyz(self):   return self.gtyz * self.conf_fact
    @cached_property
    def Kxy(self):   return self.Atxy * self.conf_fact + self.gxy * self.trK / 3.
    @cached_property
    def Kxz(self):   return self.Atxz * self.conf_fact + self.gxz * self.trK / 3.
    @cached_property
    def Kyz(self):   return self.Atyz * self.conf_fact + self.gyz * self.trK / 3.

    @cached_property
    def gamma(self):
        """
        spatial metric
        """
        return np.array([[self.gxx, self.gxy, self.gxz],
                         [self.gxz, self.gyy, self.gyz],
                         [self.gxz, self.gyz, self.gzz]])

    @cached_property
    def K(self):
        """
        spatial extrinsic curvature
        """
        ret = np.array([[self.Axx, self.Axy, self.Axz],
                        [self.Axz, self.Ayy, self.Ayz],
                        [self.Axz, self.Ayz, self.Azz]])

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 = next(iter(self.df.values()))

        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 axis_data(self, rl, **kwargs):
        return _AxisData(self, rl, **kwargs)

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