"""Load Dyablo simulation from a hdf5-formatted dataset

Requires yt_experiments to be installed:
$ pip install git+https://github.com/cphyc/yt_experiments@octree-converter

Example:
   >>> ds = load("./test_gravity_spheres_3D_iter0008987.h5")
   ... p = yt.SlicePlot(ds, "x", ("gas","density"))
   ... p.set_unit("density", "mp/cm**3")
   ... p.save("/tmp/")
"""

import yt
from yt_experiments.octree.converter import OctTree
from pathlib import Path

import h5py
import numpy as np


def readFile(filename: str | Path):
    with h5py.File(filename) as f:
        connectivity = f["connectivity"][:].astype(np.int64)
        coords = f["coordinates"][:].astype(np.float64)

        meta = dict(f["scalar_data"].attrs)
        data = {}

        for key in f:
            if key in ("connectivity", "coordinates", "scalar_data"):
                continue
            tmp = f[key][:].astype(np.float64).reshape(-1, 1)

            data["connect1", key] = tmp

    return meta, connectivity, coords, data


def load(filename: str | Path):
    meta, connectivity, coords, data = readFile(filename)

    # Get domain extent
    left_edge = coords.min(axis=0)
    right_edge = coords.max(axis=0)
    domain_width = right_edge - left_edge

    # Get cell positions
    cell_pos = (coords[connectivity[:, 6]] + coords[connectivity[:, 0]]) / 2
    cell_dx = coords[connectivity[:, 6]] - coords[connectivity[:, 0]]

    # Make sure we have cubic cells
    assert np.allclose(cell_dx, cell_dx[:, 0:1])

    cell_lvl = np.log2(domain_width[0] / cell_dx[:, 0]).astype(int)

    yt.mylog.info("Creating octree")
    oct = OctTree.from_list(cell_pos, cell_lvl, check=False)

    yt.mylog.info("Creating refmask")
    ref_mask, leaf_order = oct.get_refmask()

    unit_time = meta.get("tstar", 1)
    unit_length = meta.get("vstar", 1) * unit_time
    unit_mass = meta.get("rhostar", 1) * unit_length**3

    # Somehow yt ignores unit_system when loading octrees
    # and assumes the units are in cgs
    unit_length *= 100
    unit_mass *= 1000

    ds = yt.load_octree(
        ref_mask,
        {("gas", key[1]): val[leaf_order] for key, val in data.items()},
        bbox=np.array([left_edge, right_edge]).T,
        sim_time=meta["time"],
        length_unit=unit_length,
        time_unit=unit_time,
        mass_unit=unit_mass,
        unit_system="cgs",
        num_zones=1,
    )
    ds.stream_handler.name = filename
    return ds