"""Writing the optional field export (HDF5) described in the challenge README."""
from __future__ import annotations

from pathlib import Path

import h5py
import numpy as np


def field_grid(challenge: dict):
    """The recommended export grid for a challenge, from challenge.json."""
    fg = challenge["fields"]
    p = challenge["parameters"]
    nt, nx = fg["shape"]
    x = np.linspace(0.0, p["L"], int(nx))
    t = np.arange(int(nt)) / float(fg["frame_rate_hz"])
    return x, t


def write_fields(path: Path, x: np.ndarray, t: np.ndarray, u: np.ndarray, challenge: dict, solution_id: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with h5py.File(path, "w") as f:
        f.attrs["challenge_id"] = challenge["challenge_id"]
        f.attrs["solution_id"] = solution_id
        f.attrs["quantity"] = "transverse displacement"
        f.attrs["units"] = "m"
        f.create_dataset("x", data=np.asarray(x, dtype=np.float64))
        f.create_dataset("t", data=np.asarray(t, dtype=np.float64))
        f.create_dataset("u", data=np.asarray(u, dtype=np.float32), compression="gzip", compression_opts=4, chunks=(min(64, len(t)), len(x)))