#!/usr/bin/env python

# © 2022. Triad National Security, LLC. All rights reserved.  This
# program was produced under U.S. Government contract
# 89233218CNA000001 for Los Alamos National Laboratory (LANL), which
# is operated by Triad National Security, LLC for the U.S.  Department
# of Energy/National Nuclear Security Administration. All rights in
# the program are reserved by Triad National Security, LLC, and the
# U.S. Department of Energy/National Nuclear Security
# Administration. The Government is granted for itself and others
# acting on its behalf a nonexclusive, paid-up, irrevocable worldwide
# license in this material to reproduce, prepare derivative works,
# distribute copies to the public, perform publicly and display
# publicly, and to permit others to do so.

import argparse
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from mpl_toolkits.axes_grid1 import make_axes_locatable
import sys
import shutil
import os
from subprocess import call, DEVNULL
import glob

# from parthenon_tools import phdf
import time
from enum import Enum
from phoebus_constants import cgs, scalefree
import phoebus_utils
from phoedf import phoedf


def plot_frame_from_phoedf(
    dfile,
    ifname,
    savefig,
    geomfile=None,
    rlim=40,
    coords="cartesian",
    draw_bh=True,
    custom_name=None,
):

    cmap_uniform = "viridis"
    cmap_diverging = "RdYlBu"

    rad_active = dfile.Params["radiation/active"]

    a = dfile.Params["geometry/a"]
    hslope = dfile.Params["geometry/h"]
    reh = 1.0 + np.sqrt(1.0 - a ** 2)

    if geomfile is None:
        geomfile = dfile

    nblocks = geomfile.NumBlocks
    nx = geomfile.MeshBlockSize[0]
    ny = geomfile.MeshBlockSize[1]
    nz = geomfile.MeshBlockSize[2]
    ndim = int(nx > 1) + int(ny > 1) + int(nz > 1)
    if ndim == 2:
        xlim = 0
    else:
        xlim = -rlim

    time = dfile.Time

    # Get pcolormesh grid for each block
    ke = nz + 1 if nz > 1 else 1
    x1block = np.zeros([nblocks, nx + 1, ny + 1, ke])
    x2block = np.zeros([nblocks, nx + 1, ny + 1, ke])
    x3block = np.zeros([nblocks, nx + 1, ny + 1, ke])
    for b in range(nblocks):
        dx1 = (geomfile.BlockBounds[b][1] - geomfile.BlockBounds[b][0]) / nx
        dx2 = (geomfile.BlockBounds[b][3] - geomfile.BlockBounds[b][2]) / ny
        dx3 = (geomfile.BlockBounds[b][5] - geomfile.BlockBounds[b][4]) / nz
        for i in range(nx + 1):
            for j in range(ny + 1):
                for k in range(ke):
                    x1block[b, i, j, k] = geomfile.BlockBounds[b][0] + i * dx1
                    x2block[b, i, j, k] = geomfile.BlockBounds[b][2] + j * dx2
                    x3block[b, i, j, k] = geomfile.BlockBounds[b][4] + k * dx3
    rblock = np.exp(x1block)
    thblock = np.pi * x2block + ((1.0 - hslope) / 2.0) * np.sin(2.0 * np.pi * x2block)
    phiblock = x3block

    if ndim == 2:
        xblock = rblock * np.sin(thblock)
        yblock = rblock * np.cos(thblock)
        zblock = None
    elif ndim == 3:
        xblock = rblock * np.sin(thblock) * np.cos(phiblock)
        yblock = rblock * np.cos(thblock)
        zblock = rblock * np.cos(phiblock)  # Fake cylindrical for midplane plots

    if coords == "cartesian":
        xplot = xblock
        yplot = yblock
        zplot = zblock
    elif coords == "mks":
        xplot = x1block
        yplot = x2block
        zplot = x3block

    if rad_active:
        if ndim == 2:
            fig, axes = plt.subplots(2, 4, figsize=(8, 6))
        elif ndim == 3:
            fig, axes = plt.subplots(2, 4, figsize=(10, 6))
    else:
        if ndim == 2:
            fig, axes = plt.subplots(2, 2, figsize=(4, 6))
        elif ndim == 3:
            fig, axes = plt.subplots(2, 2, figsize=(6, 6))

    ax = axes[0, 0]
    ldensity = np.log10(dfile.GetRho())

    density = dfile.Get("p.density", flatten=False)
    J = dfile.Get("r.p.J", flatten=False)
    # print(f'density min: {density.min()} max: {density.max()}')
    # if J is not None:
    #  print(f'J min:       {J.min()} max: {J.max()}')

    if ndim == 2:
        k = 0
    else:
        k = 0

    for b in range(nblocks):
        im = ax.pcolormesh(
            xplot[b, :, :, k],
            yplot[b, :, :, k],
            ldensity[b, k, :, :].transpose(),
            vmin=-5,
            vmax=0,
            cmap=cmap_uniform,
        )
        if ndim == 3:
            im = ax.pcolormesh(
                xplot[b, :, :, k + nz // 2],
                yplot[b, :, :, k + nz // 2],
                ldensity[b, k + nz // 2, :, :].transpose(),
                vmin=-5,
                vmax=0,
                cmap=cmap_uniform,
            )
    div = make_axes_locatable(ax)
    cax = div.append_axes("right", size="5%", pad=0.05)
    fig.colorbar(im, cax=cax, orientation="vertical")
    ax.set_title(r"$\log_{10}~\rho$")

    ax = axes[0, 1]
    Gamma = dfile.GetGamma()
    for b in range(nblocks):
        im = ax.pcolormesh(
            xplot[b, :, :, k],
            yplot[b, :, :, k],
            Gamma[b, k, :, :].transpose(),
            vmin=1,
            vmax=5,
            cmap=cmap_uniform,
        )
        if ndim == 3:
            im = ax.pcolormesh(
                xplot[b, :, :, k + nz // 2],
                yplot[b, :, :, k + nz // 2],
                Gamma[b, k + nz // 2, :, :].transpose(),
                vmin=1,
                vmax=5,
                cmap=cmap_uniform,
            )

    div = make_axes_locatable(ax)
    cax = div.append_axes("right", size="5%", pad=0.05)
    fig.colorbar(im, cax=cax, orientation="vertical")
    ax.set_title(r"$\Gamma$")

    if rad_active:

        ax = axes[0, 2]
        Pg = dfile.GetPg()
        Pm = np.clip(dfile.GetPm(), 1.0e-20, 1.0e20)
        lbeta = np.log10(Pg / Pm)
        for b in range(nblocks):
            im = ax.pcolormesh(
                xplot[b, :, :, k],
                yplot[b, :, :, k],
                lbeta[b, k, :, :].transpose(),
                vmin=-2,
                vmax=2,
                cmap=cmap_diverging,
            )
            if ndim == 3:
                im = ax.pcolormesh(
                    xplot[b, :, :, k + nz // 2],
                    yplot[b, :, :, k + nz // 2],
                    lbeta[b, k + nz // 2, :, :].transpose(),
                    vmin=-2,
                    vmax=2,
                    cmap=cmap_diverging,
                )
        div = make_axes_locatable(ax)
        cax = div.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(im, cax=cax, orientation="vertical")
        ax.set_title(r"$\log_{10}~\beta_{\rm M}$")

        ax = axes[0, 3]
        lTg = np.log10(dfile.GetTg() * dfile.TemperatureCodeToCGS)
        for b in range(nblocks):
            im = ax.pcolormesh(
                xplot[b, :, :, k],
                yplot[b, :, :, k],
                lTg[b, k, :, :].transpose(),
                vmin=5,
                vmax=10,
                cmap=cmap_diverging,
            )
            if ndim == 3:
                im = ax.pcolormesh(
                    xplot[b, :, :, k + nz // 2],
                    yplot[b, :, :, k + nz // 2],
                    lTg[b, k + nz // 2, :, :].transpose(),
                    vmin=5,
                    vmax=10,
                    cmap=cmap_diverging,
                )
        div = make_axes_locatable(ax)
        cax = div.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(im, cax=cax, orientation="vertical")
        ax.set_title(r"$\log_{10}~T_{\rm g}~({\rm K})$")

        ax = axes[1, 0]
        lJ = np.log10(np.clip(dfile.Get("r.p.J", flatten=False), 1.0e-20, None))[
            :, 0, :, :, :
        ]
        for b in range(nblocks):
            im = ax.pcolormesh(
                xplot[b, :, :, k],
                yplot[b, :, :, k],
                lJ[b, k, :, :].transpose(),
                vmin=-5,
                vmax=0,
                cmap=cmap_uniform,
            )
            if ndim == 3:
                im = ax.pcolormesh(
                    xplot[b, :, :, k + nz // 2],
                    yplot[b, :, :, k + nz // 2],
                    lJ[b, k + nz // 2, :, :].transpose(),
                    vmin=-5,
                    vmax=0,
                    cmap=cmap_uniform,
                )
        div = make_axes_locatable(ax)
        cax = div.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(im, cax=cax, orientation="vertical")
        ax.set_title(r"$\log_{10}~J$")

        ax = axes[1, 1]
        lxi = np.log10(dfile.GetXi())[:, 0, :, :, :]
        for b in range(nblocks):
            im = ax.pcolormesh(
                xplot[b, :, :, k],
                yplot[b, :, :, k],
                lxi[b, k, :, :].transpose(),
                vmin=-3,
                vmax=0,
                cmap=cmap_uniform,
            )
            if ndim == 3:
                im = ax.pcolormesh(
                    xplot[b, :, :, k + nz // 2],
                    yplot[b, :, :, k + nz // 2],
                    lxi[b, k + nz // 2, :, :].transpose(),
                    vmin=-3,
                    vmax=0,
                    cmap=cmap_uniform,
                )
        div = make_axes_locatable(ax)
        cax = div.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(im, cax=cax, orientation="vertical")
        ax.set_title(r"$\log_{10}~\xi$")

        ax = axes[1, 2]
        Pg = dfile.GetPg()
        Pr = dfile.GetPr()
        lbetar = np.log10(Pg / Pr)
        for b in range(nblocks):
            im = ax.pcolormesh(
                xplot[b, :, :, k],
                yplot[b, :, :, k],
                lbetar[b, k, :, :].transpose(),
                vmin=-2,
                vmax=2,
                cmap=cmap_diverging,
            )
            if ndim == 3:
                im = ax.pcolormesh(
                    xplot[b, :, :, k + nz // 2],
                    yplot[b, :, :, k + nz // 2],
                    lbetar[b, k + nz // 2, :, :].transpose(),
                    vmin=-2,
                    vmax=2,
                    cmap=cmap_diverging,
                )
        div = make_axes_locatable(ax)
        cax = div.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(im, cax=cax, orientation="vertical")
        ax.set_title(r"$\log_{10}~\beta_{\rm R}$")

        ax = axes[1, 3]
        ltau = np.log10(dfile.GetTau())
        for b in range(nblocks):
            im = ax.pcolormesh(
                xplot[b, :, :, k],
                yplot[b, :, :, k],
                ltau[b, k, :, :].transpose(),
                vmin=-3,
                vmax=3,
                cmap=cmap_diverging,
            )
            if ndim == 3:
                im = ax.pcolormesh(
                    xplot[b, :, :, k + nz // 2],
                    yplot[b, :, :, k + nz // 2],
                    ltau[b, k + nz // 2, :, :].transpose(),
                    vmin=-3,
                    vmax=3,
                    cmap=cmap_diverging,
                )
        div = make_axes_locatable(ax)
        cax = div.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(im, cax=cax, orientation="vertical")
        ax.set_title(r"$\log_{10}~\tau_{\rm zone}$")

    else:
        ax = axes[1, 0]
        lenergy = np.log10(np.clip(dfile.Get("p.energy", flatten=False), 1.0e-20, None))
        for b in range(nblocks):
            im = ax.pcolormesh(
                xplot[b, :, :, k],
                yplot[b, :, :, k],
                lenergy[b, k, :, :].transpose(),
                vmin=-5,
                vmax=0,
                cmap=cmap_uniform,
            )
            if ndim == 3:
                im = ax.pcolormesh(
                    xplot[b, :, :, k + nz // 2],
                    yplot[b, :, :, k + nz // 2],
                    lenergy[b, k + nz // 2, :, :].transpose(),
                    vmin=-5,
                    vmax=0,
                    cmap=cmap_uniform,
                )
        div = make_axes_locatable(ax)
        cax = div.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(im, cax=cax, orientation="vertical")
        ax.set_title(r"$\log_{10}~u$")

        ax = axes[1, 1]
        Pg = dfile.GetPg()
        Pm = np.clip(dfile.GetPm(), 1.0e-20, 1.0e20)
        lbeta = np.log10(Pg / Pm)
        for b in range(nblocks):
            im = ax.pcolormesh(
                xplot[b, :, :, k],
                yplot[b, :, :, k],
                lbeta[b, k, :, :].transpose(),
                vmin=-2,
                vmax=2,
                cmap=cmap_diverging,
            )
            if ndim == 3:
                im = ax.pcolormesh(
                    xplot[b, :, :, k + nz // 2],
                    yplot[b, :, :, k + nz // 2],
                    lbeta[b, k + nz // 2, :, :].transpose(),
                    vmin=-2,
                    vmax=2,
                    cmap=cmap_diverging,
                )
        div = make_axes_locatable(ax)
        cax = div.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(im, cax=cax, orientation="vertical")
        ax.set_title(r"$\log_{10}~\beta_{\rm M}$")

    for nrow, tmp_ax in enumerate(axes):
        for ncol, ax in enumerate(tmp_ax):
            if coords == "cartesian":
                ax.set_xlim([xlim, rlim])
                ax.set_ylim([-rlim, rlim])
                ax.set_aspect("equal")
            if ncol > 0:
                ax.yaxis.set_ticklabels([])
            if draw_bh:
                bh = plt.Circle((0, 0), reh, color="k")
                ax.add_patch(bh)

    plt.suptitle("Time = %g M" % dfile.Time)

    fig.tight_layout()

    if custom_name is not None:
        savename = custom_name
    else:
        savename = str(ifname).rjust(5, "0") + ".png"
    plt.savefig(savename, dpi=300, bbox_inches="tight")
    plt.close(fig)


def plot_frame(
    ifname,
    fname,
    savefig,
    geomfile=None,
    rlim=40,
    coords="cartesian",
    draw_bh=True,
    custom_name=None,
):
    print(fname)

    dfile = phoedf(fname)

    plot_frame_from_phoedf(
        dfile, ifname, savefig, geomfile, rlim, coords, draw_bh, custom_name
    )


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Plot neutrino thermalization")
    parser.add_argument(
        "-gf",
        "--geometry_file",
        dest="geometry_file",
        default=None,
        help="Optional dump file containing the time-independent (no AMR) geometry data",
    )
    parser.add_argument(
        "--savefig", type=bool, default=False, help="Whether to save figure"
    )
    parser.add_argument(
        "--numin", type=float, default=1.0e-4, help="Minimum frequency (Hz)"
    )
    parser.add_argument(
        "--numax", type=float, default=1.0e2, help="Maximum frequency (Hz)"
    )
    parser.add_argument(
        "--nnu", type=int, default=100, help="Number of frequency support points"
    )
    parser.add_argument(
        "--coords", type=str, default="cartesian", help="Coordinates to plot in"
    )
    parser.add_argument(
        "files", type=str, nargs="+", help="Files to take a snapshot of"
    )
    args = parser.parse_args()

    matplotlib.use("Agg")

    for i, fname in enumerate(args.files):
        plot_frame(i, fname, args.savefig, coords=args.coords)
