import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors
import h5py


def main() -> None:

    fp = h5py.File(f"ab_initio_results.h5", "r")
    
    def read(name: str) -> np.ndarray:
        return np.array(fp[name])
          
    E = read("E")  # carrier energies in eV
    omega = read("omega")  # photon energies in eV
    lbda = read("lambda")  # corresponding wavelengths in nm
    df_photon = read("df_photon")  # occupation change / photon in nm^3
    df_per_power = read("df_per_power")  # occupation change / absorbed power
    emissivity = read("emissivity")  # emissivity in eV^-2 nm^-3
    self_emissivity = read("self_emissivity")  # self-emissivity contribution
    emission = read("emission")  # net emission per photon energy in eV^-1
    emission_lbda = read("emission_lambda")  # same per wavelength in nm^-1

    # Pump wavelengths and corresponding photon energies to plot:
    lbda_list = np.array([488., 502., 517., 532., 600., 678., 785.])  # in nm
    omega_list = 1239.84193 / lbda_list  # in eV
    lbda_labels = [
        f"{lbda_i:.0f} nm ({omega_i:.2f} eV)"
        for lbda_i, omega_i in zip(lbda_list, omega_list)
    ]
    lbda_indices = [np.abs(omega - omega_i).argmin() for omega_i in omega_list]

    # Carrier energies to plot for emissivity:
    E_list = np.arange(1.5, 2.6, 0.2)
    E_list = np.sort(np.concatenate((E_list, -E_list)))
    E_indices = [np.abs(E - Ei).argmin() for Ei in E_list]
    E_labels = [f"{Ei:.1f} eV" for Ei in E_list]

    # Plot all carrier-resolved properties:
    fig, axes = plt.subplots(4, 3, figsize=(15, 17), width_ratios=(8, 1, 6))
    common_kwargs = dict(
        x=E, xlabel=r"$\varepsilon - \varepsilon_F$ (eV)",
        y=omega, ylabel=r"$\hbar\omega$ (eV)",
    )
    lbda_slice_kwargs = dict(
        slice_indices=lbda_indices, slice_names=lbda_labels, slice_axis='y'
    )
    E_slice_kwargs = dict(
        slice_indices=E_indices, slice_names=E_labels, slice_axis='x'
    )
    plot_quantity(
        quantity=df_photon,
        qlabel=r"Excited $\Delta f$/(photon/vol) (nm$^3$)",
        ax_img=axes[0, 0], ax_cbar=axes[0, 1], ax_slices=axes[0, 2],
        cmap="RdBu", linthresh=0.01, vmin=-0.4, vmax=0.4,
        **common_kwargs, **lbda_slice_kwargs,
    )
    plot_quantity(
        quantity=df_per_power,
        qlabel=r"Steady $\Delta f/p_{\mathrm{abs}}$ ($\mu$m$^3$/W)",
        ax_img=axes[1, 0], ax_cbar=axes[1, 1], ax_slices=axes[1, 2],
        cmap="RdBu", linthresh=1e-7, vmin=-0.03, vmax=0.03,
        **common_kwargs, **lbda_slice_kwargs,
    )
    plot_quantity(
        quantity=emissivity,
        qlabel=r"Emissivity $\Gamma(\varepsilon, \omega)$ (eV$^{-2}$nm$^{-3}$)",
        ax_img=axes[2, 0], ax_cbar=axes[2, 1], ax_slices=axes[2, 2],
        cmap="RdBu", linthresh=1e-10, vmin=-3e-7, vmax=3e-7,
        **common_kwargs, **E_slice_kwargs,
        legend_kwargs=dict(ncol=4, columnspacing=1, loc="upper right"),
    )
    plot_quantity(
        quantity=self_emissivity,
        qlabel=r"Self-emissivity $P(\varepsilon, \omega)$",
        ax_img=axes[3, 0], ax_cbar=axes[3, 1], ax_slices=axes[3, 2],
        cmap="RdBu", linthresh=1e-9, vmin=-4e-7, vmax=5e-7,
        **common_kwargs, **E_slice_kwargs,
        legend_kwargs=dict(ncol=4, columnspacing=1, loc="lower right"),
    )
    plt.savefig("carriers.pdf", bbox_inches="tight")
    
    # Plot final emission spectrum result:
    fig, axes = plt.subplots(2, 3, figsize=(12, 11), width_ratios=(6, 1, 6))
    common_kwargs = dict(
        x=omega, xlabel=r"Emitted $\hbar\omega$ (eV)",
        y=omega, ylabel=r"Incident $\hbar\omega$ (eV)",
    )
    plot_quantity(
        quantity=emission,
        qlabel=r"Emission $\Gamma(\omega)$/photon (eV$^{-1}$)",
        ax_img=axes[0, 0], ax_cbar=axes[0, 1], ax_slices=axes[0, 2],
        cmap="PuRd", linthresh=1e-8, vmin=0, vmax=1e-4,
        **common_kwargs, **lbda_slice_kwargs,
        legend_kwargs=dict(loc="upper right"),
    )
    plot_quantity(
        quantity=emission_lbda,
        qlabel=r"Emission $\Gamma(\lambda)$/photon (nm$^{-1}$)",
        ax_img=axes[1, 0], ax_cbar=axes[1, 1], ax_slices=axes[1, 2],
        cmap="PuRd", linthresh=1e-11, vmin=0, vmax=3e-8,
        x=lbda, xlabel=r"Emitted $\lambda$ (nm)",
        y=lbda, ylabel=r"Incident $\lambda$ (nm)",
        slice_indices=[np.abs(lbda - lbda_i).argmin() for lbda_i in lbda_list],
        slice_names=lbda_labels, slice_axis='y',
        legend_kwargs=dict(loc="upper right"),
    )
    plt.savefig("emission.pdf", bbox_inches="tight")
    plt.show()


def plot_quantity(
    *, x, y, quantity, xlabel, ylabel, qlabel,
    ax_img, cmap, linthresh, vmin, vmax, ax_cbar,
    ax_slices, slice_indices, slice_names, slice_axis,
    legend_kwargs=dict(loc="upper left"),
):
    """
    Plot quantity(x, y) as an image on ax_img, and with slices on ax_slices.
    All axes are labeled appropriately using xlabel, ylabel and qlabel.
    The image plot uses colormap cmap with SymLog scaling specified by
    linthresh, vmin and vmax, and puts the colorbar on ax_cbar.
    The slices plot selects slice_indices from x or y depending on slice_axis,
    labels them on the legend using slice_names and scales the quantity axis
    using the same SymLog parameters.
    """
    # Image plot:
    plt.sca(ax_img)
    plt.axvline(0, color='k', ls='dotted', lw=1)
    plt.imshow(
        quantity.T, cmap=cmap, origin="lower",
        extent=(x.min(), x.max(), y.min(), y.max()),
        norm=colors.SymLogNorm(linthresh=linthresh, vmin=vmin, vmax=vmax),
    )
    plt.xlabel(xlabel, labelpad=0)
    plt.ylabel(ylabel, labelpad=0)
    
    # Color bar:
    bbox = ax_cbar.get_position()
    bbox.x1 = bbox.x0 + 0.3*(bbox.x1 - bbox.x0)  # shrink for label
    ax_cbar.set_position(bbox)
    plt.colorbar(cax=ax_cbar, orientation="vertical")
    ax_cbar.yaxis.set_tick_params(labelright=False)
    
    # Slices:
    plt.sca(ax_slices)
    plt.yscale("symlog", linthresh=linthresh)
    plt.ylim(vmin, vmax)
    if slice_axis == "x":
        horiz = y
        horiz_label = ylabel
        vert = quantity
    else:
        horiz = x
        horiz_label = xlabel
        vert = quantity.T
    for slice_index, slice_name in zip(slice_indices, slice_names):
        plt.plot(horiz, vert[slice_index], label=slice_name)
    plt.xlim(horiz.min(), horiz.max())
    plt.xlabel(horiz_label, labelpad=0)
    plt.ylabel(qlabel, labelpad=0)
    plt.axhline(0, color='k', ls='dotted', lw=1)
    plt.axvline(0, color='k', ls='dotted', lw=1)
    plt.legend(labelspacing=0, **legend_kwargs)


if __name__ == "__main__":
    main()
