"""
Script to test functionality from namelist creation to run and postprocessing.

It is recommended to run this in IPython.
"""

import numpy as np
import matplotlib.pyplot as plt
from omfit_classes import omfit_eqdsk
import pickle as pkl
import scipy
import sys
import os
import time
from copy import deepcopy

plt.ion()

# Make sure that package home is added to sys.path
sys.path.append("../")
import aurora

# number of repetitions to accurately time runs
num = 10

# read in default Aurora namelist
namelist = aurora.default_nml.load_default_namelist()
kp = namelist["kin_profs"]

# Use gfile and statefile in local directory:
examples_dir = os.path.dirname(os.path.abspath(__file__))
geqdsk = omfit_eqdsk.OMFITgeqdsk(examples_dir + "/example.gfile")

# save kinetic profiles on a rhop (sqrt of norm. pol. flux) grid
# parameterization f=(f_center-f_edge)*(1-rhop**alpha1)**alpha2 + f_edge
kp = namelist["kin_profs"]
T_core = 5e3  # eV
T_edge = 100  # eV
T_alpha1 = 2.0
T_alpha2 = 1.5
n_core = 1e14  # cm^-3
n_edge = 0.4e14  # cm^-3
n_alpha1 = 2
n_alpha2 = 0.5

rhop = kp["Te"]["rhop"] = kp["ne"]["rhop"] = np.linspace(0, 1, 100)
ne_cm3 = (n_core - n_edge) * (1 - rhop ** n_alpha1) ** n_alpha2 + n_edge
kp["ne"]["vals"] = ne_cm3[None, :]
Te_eV = (T_core - T_edge) * (1 - rhop ** T_alpha1) ** T_alpha2 + T_edge
kp["Te"]["vals"] = Te_eV[None, :]

# set impurity species and sources rate
imp = namelist["imp"] = "C"
namelist["source_type"] = "const"
namelist["source_rate"] = 1e21  # particles/s

# Change radial resolution from default:
# namelist['dr_0']=0.3
# namelist['dr_1']=0.05

# Change time resolution from default:
# namelist['timing']['dt_increase'] = np.array([1.01, 1.])
# namelist['timing']['dt_start'] = np.array([1e-5, 0.001])
# namelist['timing']['steps_per_cycle'] = np.array([1,1])
# namelist['timing']['times'] = np.array([0.,0.2])

# Now get aurora setup (let setup process create inputs needed to produce radial and temporal grids)
asim = aurora.core.aurora_sim(namelist, geqdsk=geqdsk)

# Visualize radial grid
rvol_grid, pro_grid, qpr_grid, prox_param = aurora.grids_utils.create_radial_grid(
    asim.namelist, plot=True
)

# Visualize time resolution
time_grid, save_grid = aurora.grids_utils.create_time_grid(
    asim.namelist["timing"], plot=True
)

# set time-independent transport coefficients (flat D=1 m^2/s, V=-2 cm/s)
D_z = 1e4 * np.ones(len(asim.rvol_grid))  # cm^2/s
V_z = -2e2 * asim.rhop_grid ** 10  # cm/s, increasing towards the edge

# plot transport coefficients
fig, ax = plt.subplots(2, 1, sharex=True, figsize=(8, 8))
ax[0].plot(asim.rhop_grid, D_z / 1e4)
ax[1].plot(asim.rhop_grid, V_z / 1e2)
ax[1].set_xlabel(r"$\rho_p$")
ax[0].set_ylabel(r"$D$ [$m^2/s$]")
ax[1].set_ylabel(r"$v$ [$m/s$]")
plt.subplots_adjust(wspace=0, hspace=0)


####### Finite-differences method #########
start = time.time()
for n in np.arange(num):
    out = asim.run_aurora(D_z, V_z, alg_opt=0)
print("Average FD time per run: ", (time.time() - start) / num)
nz = out[
    0
]  # extract only charge state densities in the plasma from output -- (time,nZ,space)

# plot charge state distributions over radius and time
aurora.plot_tools.slider_plot(
    asim.rvol_grid,
    asim.time_out,
    nz.transpose(1, 0, 2),
    xlabel=r"$r_V$ [cm]",
    ylabel="time [s]",
    zlabel=r"$n_z$ [$cm^{-3}$]",
    labels=[str(i) for i in np.arange(0, nz.shape[1])],
    plot_sum=True,
    x_line=asim.rvol_lcfs,
)

# Check particle conservation
out_fd, axs = asim.check_conservation(plot=True)


####### Finite volumes (Linder) method #########
start = time.time()
for n in np.arange(num):
    out_2 = asim.run_aurora(D_z, V_z, alg_opt=1, evolneut=False)
print("Average FV time per run: ", (time.time() - start) / num)
nz_2 = out_2[0]

# plot charge state distributions over radius and time
aurora.plot_tools.slider_plot(
    asim.rvol_grid,
    asim.time_out,
    nz_2.transpose(1, 0, 2),
    xlabel=r"$r_V$ [cm]",
    ylabel="time [s]",
    zlabel=r"$n_z$ [$cm^{-3}$]",
    labels=[str(i) for i in np.arange(0, nz_2.shape[1])],
    plot_sum=True,
    x_line=asim.rvol_lcfs,
)

# Check particle conservation
out_fv, axs = asim.check_conservation(plot=True, axs=axs)


######### Finite volumes (Linder) evolving neutrals -- under development! #########
start = time.time()
for n in np.arange(num):
    out_3 = asim.run_aurora(D_z, V_z, alg_opt=1, evolneut=True)
print("Average FVN time per run: ", (time.time() - start) / num)
nz_3 = out_3[0]

# plot charge state distributions over radius and time
aurora.plot_tools.slider_plot(
    asim.rvol_grid,
    asim.time_out,
    nz_3.transpose(1, 0, 2),
    xlabel=r"$r_V$ [cm]",
    ylabel="time [s]",
    zlabel=r"$n_z$ [$cm^{-3}$]",
    labels=[str(i) for i in np.arange(0, nz_3.shape[1])],
    plot_sum=True,
    x_line=asim.rvol_lcfs,
)

# Check particle conservation
out_fvn, axs = asim.check_conservation(plot=True, axs=axs)


###########################################################################
#
#           Compare all algorithmic choices at last time slice
#
############################################################################
labels = [fr"{imp}$^{{{str(i)}}}$" for i in np.arange(0, nz_3.shape[1])]
colors = plt.cm.rainbow(np.linspace(0, 1, nz.shape[1]))
fig = plt.figure()
fig.set_size_inches(10, 7, forward=True)
a_plot = plt.subplot2grid((10, 10), (0, 0), rowspan=10, colspan=8, fig=fig)
a_legend = plt.subplot2grid((10, 10), (0, 8), rowspan=10, colspan=2, fig=fig)
for ii, cc in zip(np.arange(nz.shape[1]), colors):
    a_plot.plot(asim.rhop_grid, nz[:, ii, -1].T, c=cc, ls="-")
    a_plot.plot(asim.rhop_grid, nz_2[:, ii, -1].T, c=cc, ls="--")

    #########
    factor = np.max(nz_2) / np.max(
        nz_3
    )  # factor needed to match neutral evolution to basic FV case
    a_plot.plot(asim.rhop_grid, factor * nz_3[:, ii, -1].T, c=cc, ls=":")
    ########
    a_legend.plot([], [], c=cc, label=labels[ii], ls="-")
a_legend.plot([], [], c="k", ls="-", lw=2, label="Original")
a_legend.plot([], [], c="k", ls="--", lw=2, label="Linder w/o n")
a_legend.plot([], [], c="k", ls=":", lw=2, label="Linder w/ n")

a_legend.legend(loc="best").set_draggable(True)
a_plot.set_xlabel(r"$\rho_p$")
a_plot.set_ylabel(r"$n_z$ [$cm^{-3}$]")
a_legend.axis("off")
fig.suptitle("Algorithm comparison")


######## images/plots of total density ############
fig, ax = plt.subplots()
pcm = ax.pcolormesh(
    asim.rhop_grid,
    asim.time_out,
    nz.sum(axis=1).T,
    cmap="plasma",
    vmax=nz.sum(axis=1)[0, :].max(),
    shading="auto",
)
fig.colorbar(pcm)
ax.set_xlabel(r"$\rho_p$")
ax.set_ylabel(r"$t$ [s]")
ax.set_title("Tot impurity density (Finite differences)")
ax.set_xlim([0.0, 1.0])

fig, ax = plt.subplots()
pcm = ax.pcolormesh(
    asim.rhop_grid,
    asim.time_out,
    nz_2.sum(axis=1).T,
    cmap="plasma",
    vmax=nz_2.sum(axis=1)[0, :].max(),
    shading="auto",
)
fig.colorbar(pcm, extend="max")
ax.set_xlabel(r"$\rho_p$")
ax.set_ylabel(r"$t$ [s]")
ax.set_title("Tot impurity density (Finite volumes)")
ax.set_xlim([0.0, 1.0])
