# AUTOGENERATED! DO NOT EDIT! File to edit: ../../notebooks/utils/01_plotting.ipynb.

# %% auto 0
__all__ = ['savefig', 'time_stamp', 'setup_mva_plot', 'plot_candidate']

# %% ../../notebooks/utils/01_plotting.ipynb 1
import matplotlib.pyplot as plt

# %% ../../notebooks/utils/01_plotting.ipynb 2
def savefig(name, **kwargs):
    plt.savefig(f"../figures/{name}.png", bbox_inches="tight", **kwargs)
    plt.savefig(f"../figures/{name}.pdf", bbox_inches="tight", **kwargs)

# %% ../../notebooks/utils/01_plotting.ipynb 4
import xarray as xr
import pandas as pd
from datetime import datetime

from pyspedas.cotrans.minvar_matrix_make import minvar_matrix_make
from pyspedas import tvector_rotate
from pyspedas.analysis.tvectot import tvectot

from pytplot import tplot
from pytplot import store_data, get_data, split_vec, join_vec
from pytplot import timebar, highlight, degap, options


# %% ../../notebooks/utils/01_plotting.ipynb 5
import matplotlib.pyplot as plt
from .plot import savefig

try:
    import scienceplots
    plt.style.use(['science', 'nature', 'notebook'])
except ImportError:
    pass

# %% ../../notebooks/utils/01_plotting.ipynb 6
def time_stamp(ts):
    "Return POSIX timestamp as float."
    return pd.Timestamp(ts, tz="UTC").timestamp()


def setup_mva_plot(
    data: xr.DataArray,
    tstart: datetime,
    tstop: datetime,
    mva_tstart: datetime = None,
    mva_tstop: datetime = None,
):
    if mva_tstart is None:
        mva_tstart = tstart
    if mva_tstop is None:
        mva_tstop = tstop

    mva_b = data.sel(time=slice(mva_tstart, mva_tstop))
    store_data("fgm", data={"x": mva_b.time, "y": mva_b})
    minvar_matrix_make("fgm")  # get the MVA matrix

    temp_b = data.sel(time=slice(tstart, tstop))
    store_data("fgm", data={"x": temp_b.time, "y": temp_b})
    tvar = tvector_rotate("fgm_mva_mat", "fgm")[0]
    ysubtitle = "[nT LMN]"
    legend_names = [r"$B_l$", r"$B_m$", r"$B_n$"]
    
    tvar2plot = tvectot(tvar, join_component=True)
    legend_names = legend_names + [r"$B_{total}$"]

    options(tvar2plot, "ytitle", "$B$")
    options(tvar2plot, "ysubtitle", ysubtitle)
    options(tvar2plot, "legend_names", legend_names)
    
    options(tvar2plot, "thick", 2)
    options(tvar2plot, "char_size", 16)
    
    # tstart_ts = time_stamp(tstart)
    # tstop_ts = time_stamp(tstop)
    # highlight(tvar2plot, [tstart_ts, tstop_ts])
    degap(tvar2plot)
    return tvar2plot

# %% ../../notebooks/utils/01_plotting.ipynb 8
from matplotlib import lines


def plot_candidate(
    candidate: dict,
    data: xr.DataArray,
    add_timebars=True,
    add_plasma_params=False,
    plot_fit_data=False,
    **kwargs,
):
    if pd.notnull(candidate.get("d_tstart")) and pd.notnull(candidate.get("d_tstop")):
        tvar = setup_mva_plot(
            data,
            candidate["tstart"],
            candidate["tstop"],
            candidate["d_tstart"],
            candidate["d_tstop"],
        )
    else:
        tvar = setup_mva_plot(data, candidate["tstart"], candidate["tstop"])

    if add_timebars:
        d_time = candidate.get("d_time")
        d_start = candidate.get("d_tstart")
        d_stop = candidate.get("d_tstop")

        if d_time:
            timebar(time_stamp(d_time), color="red")
        if d_start and pd.notnull(d_start):
            timebar(time_stamp(d_start))
        if d_stop and pd.notnull(d_stop):
            timebar(time_stamp(d_stop))

    title = ""
    
    if add_plasma_params:
        plasma_speed = candidate.get("plasma_speed")
        plasma_density = candidate.get("plasma_density")
        plasma_temperature = candidate.get("plasma_temperature")

        title += "#Plasma parameters# "
        if plasma_speed:
            title += rf"$V_i: {plasma_speed:.2f} \mathrm{{km/s}}$"
        if plasma_density:
            title += rf", $n_i: {plasma_density:.2f} \mathrm{{cm}}^{{-3}}$"
        if plasma_temperature:
            title += rf", $T_i: {plasma_temperature:.2f} \mathrm{{eV}}$"

        # options(tvar, "title", title)

    fig, axes = tplot(tvar, return_plot_objects=True, **kwargs)

    if plot_fit_data:
        fit_data = candidate.get("fit.best_fit")
        fit_time = candidate.get("fit.time")

        c = candidate.get("fit.vars.c")
        amp = candidate.get("fit.vars.amplitude")
        
        d_star = candidate.get("d_star")
        rsquared = candidate.get("fit.stat.rsquared")
        chisqr = candidate.get("fit.stat.chisqr")

        plt.plot(d_time, c + amp / 2, marker="o", markersize=10, color="red")
        plt.plot(fit_time, fit_data, label="Fit", color="black", linestyle="--")
        
        title += f"\n#Fit# $\max dB/dt$: {d_star:.2f}, $R^2$: {rsquared:.2f}, $\chi^2$: {chisqr:.2f}"
    
    plt.title(title)

    return fig, axes
