# coding: utf-8

import numpy as np
from snapshotprocessor import U
from tables import NoSuchNodeError


def load_fields(pp):
    """
    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
    """

    # Load  arrays
    pp.load_cells(keys=["pos", "vel", "dx", "rho"])

    try:
        pp.load_parts(keys=["pos", "vel", "mass", "epoch"])
    except (KeyError, NoSuchNodeError):
        pp.load_parts(keys=["pos", "vel", "mass"])

    cells = pp.cells
    parts = pp.parts

    if "epoch" not in parts:
        parts["epoch"] = np.zeros(len(pp.parts["mass"]))

    # Compute extra fields and convert units
    for dset in (cells, parts):
        dset["position"] = dset["pos"] - np.array([0.5, 0.5, 0.5])

    cells["mass"] = (
        cells["rho"]
        * cells["dx"] ** 3
        * (pp.info["unit_density"] * pp.info["unit_length"] ** 3).express(U.Msun)
    )

    cells["volume"] = cells["dx"] ** 3 * (pp.info["unit_length"] ** 3).express(
        U.kpc**3
    )

    # Separate DM from stars
    mass_dm = np.max(parts["mass"])
    mask_dm = parts["mass"] == mass_dm
    mask_star = parts["mass"] < mass_dm

    # Create dataset
    data = {
        "header": {"time": pp.time * pp.info["unit_time"].express(U.Myr)},
        "gas": {
            "position": cells["position"] * pp.info["unit_length"].express(U.kpc),
            "volume": cells["volume"],
            "velocity": cells["vel"] * pp.info["unit_velocity"].express(U.km_s),
            "mass": cells["mass"],
        },
        "stars": {
            "position": parts["position"][mask_star]
            * pp.info["unit_length"].express(U.kpc),
            "velocity": parts["vel"][mask_star]
            * pp.info["unit_velocity"].express(U.km_s),
            "mass": parts["mass"][mask_star] * pp.info["unit_mass"].express(U.Msun),
            "birth_time": parts["epoch"][mask_star]
            * pp.info["unit_time"].express(U.Myr),
        },
        "dm": {
            "position": parts["position"][mask_dm]
            * pp.info["unit_length"].express(U.kpc),
            "velocity": parts["vel"][mask_dm]
            * pp.info["unit_velocity"].express(U.km_s),
            "mass": parts["mass"][mask_dm] * pp.info["unit_mass"].express(U.Msun),
        },
    }
    return data


if __name__ == "__main__":
    from snapshotprocessor import SnapshotProcessor

    pp = SnapshotProcessor(
        "/home/nbrucy/ecogal/galturb/F20H_7_1pc_frig", num=80, params="params_gal.yml"
    )
    data = load_fields(pp)