#!/usr/bin/env python3

import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
from kuibit import simdir as sd
from kuibit import grid_data_utils as gdu
from kuibit import visualize_matplotlib as viz
import numpy as np
import sys

viz.setup_matplotlib(
    {
        "mathtext.fontset": "cm",
        "mathtext.rm": "serif",
        "font.size": 30,
        "axes.labelsize": 30,
        "xtick.labelsize": 26,
        "ytick.labelsize": 26,
        "xtick.major.size": 10,
        "ytick.major.size": 10,
        "xtick.minor.size": 4,
        "ytick.minor.size": 4,
    }
)

def plot_2d_snapshot(path_to_folder):
    print("Opening SimDir")
    with (sd.SimDir(path_to_folder) as sim):

        # Get x-z slice of the scalar field (real part) grid function data
        gf_Theta_xz = sim.gf.xz.get("Theta_gf")
        gf_Phi_xz = sim.gf.xz.get("Phi_gf")
        times = gf_Theta_xz.available_times

        # Set uniform grid to resample the hierarchical grid data on
        gridsize = 6
        grid_2d = gdu.UniformGrid([1000, 1000], x0=[-gridsize, -gridsize], x1=[gridsize, gridsize])

        time = 100

        # Axion (Theta) field
        it = gf_Theta_xz.iteration_at_time(time)
        Theta_xz = gf_Theta_xz.read_on_grid(it, grid_2d)

        fig = plt.figure(figsize=(12, 8))
        gs = fig.add_gridspec(1, 2, wspace=0)
        (ax1, ax2) = gs.subplots(sharey=True)

        img = ax1.imshow(
            Theta_xz.data.T,
            vmin=-0.002, vmax=0.002,
            extent=[-gridsize, gridsize, -gridsize, gridsize],
            cmap="vanimo_r",
            origin="lower"
        )
        cb = plt.colorbar(img, ax=ax1, extend="both", label=r'$\Theta$',
                         fraction=0.046, pad=0.04, shrink=0.8, location='top')
        cb.formatter.set_powerlimits((0, 0))

        # Dilaton (Phi) field
        it = gf_Phi_xz.iteration_at_time(time)
        Phi_xz = gf_Phi_xz.read_on_grid(it, grid_2d)

        img2 = ax2.imshow(
            Phi_xz.data.T,
            vmin=0.0005, vmax=0.005,
            extent=[-gridsize, gridsize, -gridsize, gridsize],
            cmap="inferno",
            origin="lower"
        )
        cb = plt.colorbar(img2, ax=ax2, extend="both", label=r'$\Phi$',
                         fraction=0.046, shrink=0.8, pad=0.04, location='top')
        cb.formatter.set_powerlimits((0, 0))

        # Horizon
        horizons = sim.horizons
        hor = horizons[(0, 1)]
        centroid_x_timeseries = hor.ah.centroid_x
        centroid_z_timeseries = hor.ah.centroid_z
        min_x_timeseries = hor.ah.min_x
        max_x_timeseries = hor.ah.max_x
        min_z_timeseries = hor.ah.min_z
        max_z_timeseries = hor.ah.max_z

        horizon_iteration = np.argmax(np.isclose(centroid_x_timeseries.t, np.asarray(time), atol=1e-3))
        centroid_x = centroid_x_timeseries.values[horizon_iteration]
        centroid_z = centroid_z_timeseries.values[horizon_iteration]
        min_x = min_x_timeseries.values[horizon_iteration]
        max_x = max_x_timeseries.values[horizon_iteration]
        min_z = min_z_timeseries.values[horizon_iteration]
        max_z = max_z_timeseries.values[horizon_iteration]
        xdiameter = max_x - min_x
        zdiameter = max_z - min_z

        circle_1 = Ellipse((centroid_x, centroid_z), xdiameter, zdiameter, color='black', fill=True)
        circle_2 = Ellipse((centroid_x, centroid_z), xdiameter, zdiameter, color='black', fill=True)
        ax1.add_patch(circle_1)
        ax2.add_patch(circle_2)

        # Add title and labels
        ax1.set_xlabel(r"$x/M$")
        ax2.set_xlabel(r"$x/M$")
        ax1.set_ylabel(r"$z/M$")
        t_maj = [-gridsize, -gridsize / 2, 0, gridsize / 2, gridsize]

        ax1.set_xticks(t_maj)
        xticks = ax1.xaxis.get_major_ticks()
        xticks[-1].label1.set_visible(False)
        ax2.set_xticks(t_maj)
        ax1.set_yticks(t_maj)

        plt.tight_layout()
        save_path = f'xz_t{time}.png'
        plt.savefig(
           save_path,
           facecolor="w",
           dpi=100
        )
        plt.close()

####################################################################
# Example to use plotting script
####################################################################

path_to_folder = None  # path to top level directory of the simulation

if path_to_folder is None:
  print("Please set path_to_folder variable in source code to point to the main directory of the simulation", file=sys.stderr)
  sys.exit(1)

plot_2d_snapshot(path_to_folder)