"""
Makes a gas density projection plot. Uses the swiftsimio library.
"""

import matplotlib.pyplot as plt
import numpy as np

import swiftsimio as sw
from swiftsimio.visualisation.projection import project_gas

import unyt
from unyt import mh, cm
from tqdm import tqdm
from matplotlib.colors import LogNorm
from matplotlib.animation import FuncAnimation

# %%
# Constants; these could be put in the parameter file but are rarely changed.
image_resolution = 1024

# Plotting controls
cmap = "inferno"

# %%


def get_gas_mu(data: sw.SWIFTDataset) -> np.array:
    """
    Return the mean molecular weight of the gas.
    """
    # Get the cooling model
    cooling = data.metadata.subgrid_scheme["Cooling Model"]

    # Use a variable for better readability
    gas = data.gas

    # Get rho
    rho = gas.densities.to(unyt.g / unyt.cm ** 3)  # self.Rho(units='g/cm3')

    # hydrogen mass in gram
    mh.convert_to_cgs()
    mH_in_g = mh

    if cooling == b"Grackle3":  # grackle3
        nHI = gas.hi * rho / (mH_in_g)
        nHII = gas.hii * rho / (mH_in_g)
        nHeI = gas.he_i * rho / (4 * mH_in_g)
        nHeII = gas.he_ii * rho / (4 * mH_in_g)
        nHeIII = gas.he_iii * rho / (4 * mH_in_g)
        nH2I = gas.h2_i * rho / (2 * mH_in_g)
        nH2II = gas.h2_ii * rho / (2 * mH_in_g)
        nHDI = gas.hdi * rho / (3 * mH_in_g)

        nel = nHII + nHeII + 2 * nHeIII + nH2II
        mu = (
            (nHI + nHII) + (nHeI + nHeII + nHeIII) * 4 + (nH2I + nH2II) * 2 + nHDI * 3
        ) / (nHI + nHII + nHeI + nHeII + nHeIII + nH2I + nH2II + nHDI + nel)
        return mu

    elif cooling == b"Grackle2":  # grackle2
        nHI = gas.hi * rho / (mH_in_g)
        nHII = gas.hii * rho / (mH_in_g)
        nHeI = gas.he_i * rho / (4 * mH_in_g)
        nHeII = gas.he_ii * rho / (4 * mH_in_g)
        nHeIII = gas.he_iii * rho / (4 * mH_in_g)
        nH2I = gas.h2_i * rho / (2 * mH_in_g)
        nH2II = gas.h2_ii * rho / (2 * mH_in_g)

        nel = nHII + nHeII + 2 * nHeIII + nH2II
        mu = ((nHI + nHII) + (nHeI + nHeII + nHeIII) * 4 + (nH2I + nH2II) * 2) / (
            nHI + nHII + nHeI + nHeII + nHeIII + nH2I + nH2II + nel
        )
        return mu

    elif cooling == b"Grackle1":  # grackle1
        nHI = gas.hi * rho / (mH_in_g)
        nHII = gas.hii * rho / (mH_in_g)
        nHeI = gas.he_i * rho / (4 * mH_in_g)
        nHeII = gas.he_ii * rho / (4 * mH_in_g)
        nHeIII = gas.he_iii * rho / (4 * mH_in_g)
        nel = nHII + nHeII + 2 * nHeIII
        mu = ((nHI + nHII) + (nHeI + nHeII + nHeIII) * 4) / (
            nHI + nHII + nHeI + nHeII + nHeIII + nel
        )
        return mu

    else:  # Grackle0
        # print("... found info for grackle mode=0")

        from unyt.physical_constants import kboltz_cgs as k_B_cgs

        gamma = data.metadata.gas_gamma[0]

        # Get internal energy
        u = data.gas.internal_energies
        u = u.to_physical()
        u = u.to(unyt.erg / unyt.g)

        # Get hydrigen fraction
        H_frac = float(
            data.metadata.parameters["GrackleCooling:HydrogenFractionByMass"]
        )

        # Compute T/mu
        T_over_mu = (gamma - 1.0) * u.value * mH_in_g.value / k_B_cgs.value
        T_trans = 1.1e4
        mu_trans = 4.0 / (8.0 - 5.0 * (1.0 - H_frac))

        # Determine if we are ionized or not
        mu = np.ones(np.size(u))
        mask_ionized = T_over_mu > (T_trans + 1) / mu_trans
        mask_neutral = T_over_mu < (T_trans + 1) / mu_trans

        # Give the right mu
        mu[mask_ionized] = 4.0 / (8.0 - 5.0 * (1.0 - H_frac))
        mu[mask_neutral] = 4.0 / (1.0 + 3.0 * H_frac)

        return mu


def get_gas_temperatures(data: sw.SWIFTDataset) -> np.array:
    """
        Compute the temperature of the gas.
    """
    from unyt.physical_constants import kboltz_cgs as k_B

    # from unyt.physical_constants import mh

    # Convert to cgs
    mh.convert_to_cgs()

    # Get the cooling model
    cooling = data.metadata.subgrid_scheme["Cooling Model"]

    # Get internal energy and convert to physical units in cgs
    u = data.gas.internal_energies
    u = u.to_physical()
    u = u.to(unyt.erg / unyt.g)

    # Get gamm and compute mu
    gamma = data.metadata.gas_gamma[0]
    mu = get_gas_mu(data)

    # FInally compute the the temperature
    if cooling == b"Grackle3" or cooling == b"Grackle2" or cooling == b"Grackle1":
        T = mu * (gamma - 1.0) * u * mh / k_B
    else:
        a = (gamma - 1.0) * (mu * mh) / k_B * u
        T = np.where((u.value > 0), a.value, 0) * unyt.kelvin
    return T


def get_sink_and_stars_positions(filename):
    data = sw.load(filename)

    sink_pos = data.sinks.coordinates
    star_pos = data.stars.coordinates

    return sink_pos, star_pos


def make_projection(filename, image_resolution):
    """
    Compute a mass projection with swiftsimio.
    """
    data = sw.load(filename)

    # Compute projected density
    projected_mass = project_gas(
        data,
        resolution=image_resolution,
        project="masses",
        parallel=True,
        periodic=True,
    )

    boxsize = data.metadata.boxsize
    x_edges = np.linspace(0 * unyt.kpc, boxsize[0], image_resolution)
    y_edges = np.linspace(0 * unyt.kpc, boxsize[1], image_resolution)

    # Convert to 1/cm**2
    projected_mass = projected_mass.to(mh / (cm ** 2))

    return projected_mass.T, x_edges, y_edges


def setup_axes():
    """
    Creates the figure and axis object.
    """
    fig, ax = plt.subplots(1, figsize=(6, 5), dpi=300)

    ax.set_xlabel("x [kpc]")
    ax.set_ylabel("y [kpc]")

    return fig, ax


def make_single_image(filename, image_resolution):
    """
    Makes a single image and saves it to mass_projection_{snapshot_number}.png.

    Filename should be given _without_ hdf5 extension.
    """
    file = "{:s}.hdf5".format(filename)
    fig, ax = setup_axes()
    projected_mass, x_edges, y_edges = make_projection(file, image_resolution)

    mappable = ax.pcolormesh(
        x_edges, y_edges, projected_mass, cmap=cmap, norm=LogNorm()
    )
    fig.colorbar(mappable, label="Surface density [cm$^{-2}]$", pad=0)

    sink_pos, star_pos = get_sink_and_stars_positions(file)

    if star_pos.size != 0:
        ax.scatter(star_pos[:, 0], star_pos[:, 1], c="limegreen", zorder=1, marker="*")

    if sink_pos.size != 0:
        ax.scatter(sink_pos[:, 0], sink_pos[:, 1], c="blue", zorder=2, marker=".")

    ax.text(
        0.7,
        0.95,
        "$N_{\mathrm{star}}" + " = {}$".format(len(star_pos)),
        transform=ax.transAxes,
        fontsize=7,
        bbox=dict(facecolor="white", alpha=0.8),
    )
    ax.text(
        0.1,
        0.95,
        "$N_{\mathrm{sink}}" + " = {}$".format(len(sink_pos)),
        transform=ax.transAxes,
        fontsize=7,
        bbox=dict(facecolor="white", alpha=0.8),
    )

    fig.tight_layout()

    image = "mass_projection_{:s}.png".format(filename[-4:])
    fig.savefig(image)

    return


def make_movie(args, image_resolution):
    """
    Makes a movie and saves it to mass_projection_movie.mp4.
    """

    fig, ax = setup_axes()

    def grab_metadata(n):
        filename = "{:s}_{:04d}.hdf5".format(args["stub"], n)
        data = sw.load(filename)

        return data.metadata

    def grab_data(n):
        filename = "{:s}_{:04d}.hdf5".format(args["stub"], n)

        H, _, _ = make_projection(filename, image_resolution)

        # Need to ravel because pcolormesh's set_array takes a 1D array. Might
        # as well do it here, because 1d arrays are easier to max() than 2d.
        return H.ravel()

    def grab_sink_star_pos(n):
        filename = "{:s}_{:04d}.hdf5".format(args["stub"], n)

        sink_pos, star_pos = get_sink_and_stars_positions(filename)

        return sink_pos, star_pos

    histograms = [
        grab_data(n)
        for n in tqdm(
            range(args["initial"], args["final"] + 1),
            desc="Computing gas mass projection",
        )
    ]

    sink_star = [
        grab_sink_star_pos(n)
        for n in tqdm(
            range(args["initial"], args["final"] + 1),
            desc="Getting sink and stars positions",
        )
    ]

    sink_pos = [s[0] for s in sink_star]
    star_pos = [s[1] for s in sink_star]

    metadata = [
        grab_metadata(n)
        for n in tqdm(
            range(args["initial"], args["final"] + 1), desc="Grabbing metadata"
        )
    ]

    # Need to get a reasonable norm so that we don't overshoot.
    max_particles = max([x.max() for x in histograms])
    min_particles = max([x.min() for x in histograms])
    norm = LogNorm(vmin=min_particles, vmax=max_particles)

    # First, let's make the initial frame (we need this for our d, T values that we
    # got rid of in grab_data.
    hist, x_edges, y_edges = make_projection(
        "{:s}_{:04d}.hdf5".format(args["stub"], args["initial"]), image_resolution
    )

    mappable = ax.pcolormesh(x_edges, y_edges, hist, cmap=cmap, norm=norm)
    fig.colorbar(mappable, label="Surface density [cm$^{-2}]$", pad=0)

    fig.tight_layout()

    # Once we've rearranged the figure with tight_layout(), we can start laying
    # Down the metadata text.

    def format_metadata(metadata: sw.SWIFTMetadata):
        t = metadata.t
        t.convert_to_units(unyt.Myr)

        x = "$a$: {:2.2f}\n$z$: {:2.2f}\n$t$: {:2.2f}".format(metadata.a, metadata.z, t)

        return x

    text = ax.text(
        0.025,
        0.975,
        format_metadata(metadata[0]),
        ha="left",
        va="top",
        transform=ax.transAxes,
        color="white",
    )

    ax.text(
        0.975,
        0.975,
        metadata[0].code["Git Revision"].decode("utf-8"),
        ha="right",
        va="top",
        transform=ax.transAxes,
        color="white",
    )

    def animate(data):
        mappable.set_array(histograms[data])
        text.set_text(format_metadata(metadata[data]))

        if star_pos[data].size != 0:
            ax.scatter(
                star_pos[data][:, 0],
                star_pos[data][:, 1],
                c="limegreen",
                zorder=1,
                marker="*",
            )

        if sink_pos[data].size != 0:
            ax.scatter(
                sink_pos[data][:, 0],
                sink_pos[data][:, 1],
                c="blue",
                zorder=2,
                marker=".",
            )

        return mappable

    animation = FuncAnimation(
        fig, animate, range(len(histograms)), fargs=[], interval=1000 / 10
    )

    animation.save("mass_projection_movie.mp4")

    return


# %%
if __name__ == "__main__":
    import argparse as ap

    parser = ap.ArgumentParser(
        description="""
             Plotting script for making a mass projection plot.
             Takes the filename handle, start, and (optionally) stop
             snapshots. If stop is not given, png plot is produced for
             that snapshot. If given, a movie is made.
             """
    )

    parser.add_argument(
        "-i",
        "--initial",
        help="""Initial snapshot number. Default: 0.""",
        default=0,
        required=False,
        type=int,
    )

    parser.add_argument(
        "-f",
        "--final",
        help="""Final snapshot number. Default: 0.""",
        default=0,
        required=False,
        type=int,
    )

    parser.add_argument(
        "-s",
        "--stub",
        help="""Root of the snapshots filenames (e.g. snapshot). This is
                the first part of the filename for the snapshots,
                not including the final underscore. Required.""",
        type=str,
        required=True,
    )

    args = vars(parser.parse_args())

    if args["final"] <= args["initial"]:
        # Run in single image mode.
        filename = "{:s}_{:04d}".format(args["stub"], args["initial"])

        make_single_image(filename, image_resolution=image_resolution)

    else:
        # Movie mode!
        make_movie(args, image_resolution=image_resolution)