# create_h5.py
"""
Created on Fri Apr 22 11:14:28 2016

@author: fagent

Create structure for hdf5 storage of 2D averages time series, motivated
to save emf tensors for use with mean-field pencil code module.
To generalise hdf5 for all data structures
"""


def create_aver_sph(
    filename, dataset, fields, nzyxt, zyxt, hdf5dir="data", dgroup="emftensor"
):
    """
    call signature:

    create_aver_sph(filename, dataset, fields, nzyxt, zyxt,
                    hdf5dir='data', dgroup='emftensor')

    Keyword arguments:

      *filename*:
        Output filename.

      *dataset*:
        HDF5 dataset.

      *fields*:
        Names of the fields to be written.

      *nzyxt*:
        Array containing [nz, ny, nx, nt].

      *zyxt*:
        Array containing the arrays z, y, x and t.

      *hdf5dir*:
        Output data directory.

      *dgroup*:
        Data group name.
    """

    import numpy as np
    import os
    import h5py

    # Unpack the tuples.
    nz, ny, nx, nt = nzyxt
    z, y, x, t = zyxt

    # Prepare hdf5 directory.
    if not os.path.exists(hdf5dir):
        os.makedirs(hdf5dir)
    print(hdf5dir, os.path.exists(hdf5dir), "constructing hdf5 file")
    if os.path.exists(filename):
        print(filename + " already exists, skipping file creation.")
    with h5py.File(filename, "a") as hf:
        # Existing grid will not be rewritten by subsequent datasets.
        if "grid" not in hf.keys():
            hf.create_group("grid")
        if "t" not in hf["grid"].keys():
            hf.create_dataset("grid/t", (t.size,), data=np.asfarray(t))
        if "x" not in hf["grid"].keys():
            hf.create_dataset("grid/x", (nx,), data=np.asfarray(x))
        if "y" not in hf["grid"].keys():
            hf.create_dataset("grid/y", (ny,), data=np.asfarray(y))
        if "z" not in hf["grid"].keys():
            hf.create_dataset("grid/z", (nz,), data=np.asfarray(z))
        if dgroup not in hf.keys():
            hf.create_group(dgroup)
        for field, comp in fields:
            shape = comp + (nz, ny, nx, nt)
            if field not in hf[dgroup].keys():
                hf.create_group("{0}/{1}".format(dgroup, field))
            if dataset not in hf["{0}/{1}".format(dgroup, field)].keys():
                hf.create_dataset(
                    "{0}/{1}/{2}".format(dgroup, field, dataset), shape, "f"
                )
            if shape != hf["{0}/{1}/{2}".format(dgroup, field, dataset)].shape:
                del hf["{0}/{1}/{2}".format(dgroup, field, dataset)]
                hf.create_dataset(
                    "{0}/{1}/{2}".format(dgroup, field, dataset), shape, "f"
                )


fvars = [
    ("utensor", (3,)),
    (
        "alpha",
        (
            3,
            3,
        ),
    ),
    (
        "beta",
        (
            3,
            3,
        ),
    ),
    ("gamma", (3,)),
    ("delta", (3,)),
    (
        "kappa",
        (
            3,
            3,
            3,
        ),
    ),
    (
        "acoef",
        (
            3,
            3,
        ),
    ),
    (
        "bcoef",
        (
            3,
            3,
            3,
        ),
    ),
]