# coding: utf-8

import numpy as np
import h5py
from astropy import units as u


def load_fields_arepo(path):
    """
    Parameters
    ----------
    path : _type_
        _description_
    output : _type_
        _description_

    Returns
    -------
    dict
        A dataset of cells in the following format:
        gas:
            position (Ngas, 3) [kpc], centered
            volume   (Ngas)    [pc^3]
            velocity (Ngas, 3) [km/s]
            mass     (Ngas)    [Msun]
        stars:
            position (Nstar, 3) [kpc], centered
            velocity (Nstar, 3) [km/s]
            mass     (Nstar)    [Msun]
            birth_time      (Nstar)    [Myr]
        dm:
            position (Ngas, 3) [kpc], centered
            velocity (Ngas, 3) [km/s]
            mass     (Ngas)    [Msun]
        maps:
            extent      (xmin, xmax, ymin, ymax) Coordinates of the edges of the map, centered
            gas_coldens (Nx, Ny) [Msun/pc^2], map of column density
    """

    snap = h5py.File(path)
    info = snap["Header"].attrs
    gas = snap["PartType0"]
    stars = snap["PartType4"]
    time = info["Time"]
    box_size = info["BoxSize"]
    UnitLength_in_kpc = info["UnitLength_in_cm"] * u.cm.to(u.kpc)
    UnitVelocity_in_km_per_s = info["UnitVelocity_in_cm_per_s"] * u.cm.to(u.km)
    UnitTime_in_Myr = (
        info["UnitLength_in_cm"] / info["UnitVelocity_in_cm_per_s"]
    ) * u.s.to(u.Myr)
    UnitMass_in_Msun = info["UnitMass_in_g"] * u.g.to(u.Msun)

    # Create dataset
    data = {
        "header": {"time": time * UnitTime_in_Myr, 
                   "fluids": ["gas", "stars"],
                   "box_size": box_size * UnitLength_in_kpc },
        "gas": {
            "position": (
                np.asarray(gas["CenterOfMass"]) - np.array([0.5, 0.5, 0.5]) * box_size
            )
            * UnitLength_in_kpc,
            "volume": (np.asarray(gas["Masses"]) / np.asarray(gas["Density"]))
            * UnitLength_in_kpc**3,
            "velocity": np.asarray(gas["Velocities"]) * UnitVelocity_in_km_per_s,
            "mass": np.asarray(gas["Masses"]) * UnitMass_in_Msun,
        },
        "stars": {
            "position": (
                np.asarray(stars["Coordinates"]) - np.array([0.5, 0.5, 0.5]) * box_size
            )
            * UnitLength_in_kpc,
            "velocity": np.asarray(stars["Velocities"]) * UnitVelocity_in_km_per_s,
            "mass": np.asarray(stars["Masses"]) * UnitMass_in_Msun,
            "birth_time": np.asarray(stars["StellarFormationTime"]) * UnitTime_in_Myr,
        },
    }
    snap.close()
    return data