#!/usr/bin/env python3

import os
import sys
import argparse

import matplotlib as mpl
import numpy as np
import swiftsimio
from matplotlib import pyplot as plt
from scipy.optimize import curve_fit as cf

# Parameters users should/may tweak
snapshot_base = "output"  # snapshot basename
plot_physical_quantities = True

nPhotonGroups = 11

mpl.rcParams["text.usetex"] = True


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-z", "--redshift", help="Redshift domain to plot advection for", default="high"
    )

    args = parser.parse_args()
    return args


def get_snapshot_list(snapshot_basename="output"):
    """
    Find the snapshots that are to be plotted 
    and return their names as list
    """
    snaplist = []

    dirlist = os.listdir()
    for f in dirlist:
        if f.startswith(snapshot_basename) and f.endswith("hdf5"):
            snaplist.append(f)

    snaplist = sorted(snaplist)
    if len(snaplist) == 0:
        print(f"No snapshots with base {snapshot_basename} found!")
        sys.exit(1)
    return snaplist


def plot_param_over_time(
    snapshot_list, param="energy density", redshift_domain="high_redshift"
):
    print(f"Now plotting {param} over time")

    # Arrays to keep track of plot_param and scale factor
    plot_param = [[], []]
    scale_factor = []
    analytic_exponent = [0, 0]

    # Functions to convert between scale factor and redshift
    a2z = lambda a: 1 / a - 1
    z2a = lambda z: 1 / (z + 1)

    for file in snapshot_list:
        data = swiftsimio.load(file)
        meta = data.metadata

        # Read comoving quantities
        energy = data.gas.photon_energies.group1
        for i in range(nPhotonGroups - 1):
            energy += getattr(data.gas.photon_energies, f"group{i+2}")
        mass = data.gas.masses
        rho = data.gas.densities
        vol = mass / rho
        energy_density = energy / vol

        if plot_physical_quantities:
            physical_energy_density = energy_density.to_physical()
            physical_mass = mass.to_physical()
            physical_vol = vol.to_physical()
            physical_energy = physical_energy_density * physical_vol
            if param == "energy density":
                plot_param[1].append(
                    1
                    * np.sum(physical_energy_density)
                    / physical_energy_density.shape[0]
                )
                analytic_exponent[1] = -4.0
            elif param == "total energy":
                plot_param[1].append(1 * np.sum(physical_energy))
                analytic_exponent[1] = -1.0

        if param == "energy density":
            plot_param[0].append(1 * np.sum(energy_density) / energy_density.shape[0])
            analytic_exponent[0] = -1.0
        elif param == "total energy":
            plot_param[0].append(1 * np.sum(energy))
            analytic_exponent[0] = -1.0

        scale_factor.append(meta.scale_factor)

    fig = plt.figure(figsize=(5.05 * (1 + plot_physical_quantities), 5.4), dpi=200)

    x = np.linspace(min(scale_factor), max(scale_factor), 1000)

    if param == "energy density":
        titles = ["Comoving energy density", "Physical energy density"]
        ylabel = "Average energy density"
        figname = f"output_energy_density_over_time-{redshift_domain}.png"
    elif param == "total energy":
        titles = ["Comoving total energy", "Physical total energy"]
        ylabel = "Total energy"
        figname = f"output_total_energy_over_time-{redshift_domain}.png"

    for i in range(1 + plot_physical_quantities):
        ax = fig.add_subplot(1, (1 + plot_physical_quantities), (1 + i))
        ax.scatter(scale_factor, plot_param[i], label="Simulation")

        # Analytic scale-factor relation
        analytic = x ** analytic_exponent[i]

        # Scale solution to correct offset
        analytic = analytic / analytic[0] * plot_param[i][0]
        ax.plot(
            x,
            analytic,
            c="r",
            label=f"Analytic solution $\propto a^{{{analytic_exponent[i]}}}$",
        )

        ax.legend()
        ax.set_title(titles[i] + " sum of all groups")

        ax.set_xlabel("Scale factor")
        secax = ax.secondary_xaxis("top", functions=(a2z, z2a))
        secax.set_xlabel("Redshift")

        ax.yaxis.get_offset_text().set_position((-0.05, 1))

        if analytic_exponent[i] == 0.0:
            ax.set_ylim(plot_param[i][0] * 0.95, plot_param[i][0] * 1.05)
        if i == 0:
            units = plot_param[i][0].units.latex_representation()
            ax.set_ylabel(f"{ylabel} [${units}$]")

    plt.tight_layout()
    plt.savefig(figname)
    plt.close()


if __name__ in ("__main__"):
    # Get command line args
    args = parse_args()
    domain = args.redshift.lower()
    if domain in ("low", "l", "low_redshift", "low redshift", "low-redshift"):
        redshift_domain = "low_redshift"
    elif domain in (
        "medium",
        "m",
        "medium_redshift",
        "medium redshift",
        "medium-redshift",
    ):
        redshift_domain = "medium_redshift"
    elif domain in ("high", "h", "high_redshift", "high redshift", "high-redshift"):
        redshift_domain = "high_redshift"
    else:
        print("Redshift domain not recognised!")
        sys.exit(1)

    snaplist = get_snapshot_list(snapshot_base + f"_{redshift_domain}")
    if len(snaplist) < 1:
        print("No snapshots found!")
        exit(1)

    for param in ["energy density", "total energy"]:
        plot_param_over_time(snaplist, param, redshift_domain)