"""
Makes a rho-T plot. Uses the swiftsimio library.
"""

import matplotlib.pyplot as plt
import numpy as np

from swiftsimio import SWIFTDataset, SWIFTMetadata, SWIFTUnits

from unyt import mh, cm, Gyr
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.
density_bounds = [1e-8, 1e4]  # in nh/cm^3
temperature_bounds = [1e2, 1e8]  # in K
bins = 128

# Plotting controls
cmap = "viridis"


def get_data(filename):
    """
    Grabs the data (T in Kelvin and density in mh / cm^3).
    """

    data = SWIFTDataset(filename)

    data.gas.densities.convert_to_units(mh / (cm ** 3))
    data.gas.temperatures.convert_to_cgs()

    return data.gas.densities, data.gas.temperatures


def make_hist(filename, density_bounds, temperature_bounds, bins):
    """
    Makes the histogram for filename with bounds as lower, higher
    for the bins and "bins" the number of bins along each dimension.

    Also returns the edges for pcolormesh to use.
    """

    density_bins = np.logspace(
        np.log10(density_bounds[0]), np.log10(density_bounds[1]), bins
    )
    temperature_bins = np.logspace(
        np.log10(temperature_bounds[0]), np.log10(temperature_bounds[1]), bins
    )

    H, density_edges, temperature_edges = np.histogram2d(
        *get_data(filename), bins=[density_bins, temperature_bins]
    )

    return H.T, density_edges, temperature_edges


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

    ax.set_xlabel("Density [$n_H$ cm$^{-3}$]")
    ax.set_ylabel("Temperature [K]")

    ax.loglog()

    return fig, ax


def make_single_image(filename, density_bounds, temperature_bounds, bins):
    """
    Makes a single image and saves it to rhoTPlot_{filename}.png.
    
    Filename should be given _without_ hdf5 extension.
    """

    fig, ax = setup_axes()
    hist, d, T = make_hist(
        "{:s}.hdf5".format(filename), density_bounds, temperature_bounds, bins
    )

    mappable = ax.pcolormesh(d, T, hist, cmap=cmap, norm=LogNorm())
    fig.colorbar(mappable, label="Number of particles", pad=0)

    fig.tight_layout()

    fig.savefig("rhoTPlot_{:s}.png".format(filename))

    return


def make_movie(args, density_bounds, temperature_bounds, bins):
    """
    Makes a movie and saves it to rhoTPlot_{stub}.mp4.
    """

    fig, ax = setup_axes()

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

        return data

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

        H, _, _ = make_hist(filename, density_bounds, temperature_bounds, bins)

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

    histograms = [
        grab_data(n)
        for n in tqdm(
            range(args["initial"], args["final"] + 1), desc="Histogramming data"
        )
    ]

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

    units = SWIFTUnits("{:s}_{:04d}.hdf5".format(args["stub"], args["initial"]))

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

    norm = LogNorm(vmin=1, 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, d, T = make_hist(
        "{:s}_{:04d}.hdf5".format(args["stub"], args["initial"]),
        density_bounds,
        temperature_bounds,
        bins,
    )

    mappable = ax.pcolormesh(d, T, hist, cmap=cmap, norm=norm)
    fig.colorbar(mappable, label="Number of particles", pad=0)

    fig.tight_layout()

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

    def format_metadata(metadata: SWIFTMetadata):
        t = metadata.t * units.units["Unit time in cgs (U_t)"]
        t.convert_to_units(Gyr)

        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,
    )

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

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

        return mappable

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

    animation.save("rhoTPlot_{:s}.mp4".format(args["stub"]))

    return


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

    parser = ap.ArgumentParser(
        description="""
             Plotting script for making a rho-T 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="""Stub for the filename (e.g. santabarbara). 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,
            density_bounds=density_bounds,
            temperature_bounds=temperature_bounds,
            bins=bins,
        )

    else:
        # Movie mode!
        make_movie(
            args,
            density_bounds=density_bounds,
            temperature_bounds=temperature_bounds,
            bins=bins,
        )