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

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

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

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
from pytplot import timebar, degap, options

from matplotlib.pyplot import Axes

# %% ../../notebooks/utils/01_plotting.ipynb 2
try:
    import scienceplots
    plt.style.use(['science', 'nature', 'notebook'])
except ImportError:
    pass

# %% ../../notebooks/utils/01_plotting.ipynb 3
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)

    degap(tvar2plot)
    return tvar2plot

# %% ../../notebooks/utils/01_plotting.ipynb 5
from pyspedas.analysis.deriv_data import deriv_data
from ..integration import J_FACTOR
import pytplot
from pytplot import split_vec

def plot_candidate(
    event: dict,
    data: xr.DataArray,
    add_ids_properties=True,
    plot_current_density=False,
    plot_fit_data=False,
    add_timebars=True,
    add_plasma_params=False,
    **kwargs,
):
    if pd.notnull(event.get("t.d_start")) and pd.notnull(event.get("t.d_end")):
        tvar = setup_mva_plot(
            data,
            event["tstart"],
            event["tstop"],
            event["t.d_start"],
            event["t.d_end"],
        )
    else:
        tvar = setup_mva_plot(data, event["tstart"], event["tstop"])
    
    tvars2plot = [tvar]

    if plot_current_density:
        Bl = split_vec(tvar)[0]
        dBldt = deriv_data(Bl)[0]
        v_k = event.get("v_k")
        pytplot.data_quants[dBldt] = pytplot.data_quants[dBldt] / v_k * J_FACTOR.value
        tvars2plot.append(dBldt)
        
        options(dBldt, "ytitle", "$J$")
        options(dBldt, "ysubtitle", "[nA/m$^2$]")
        options(dBldt, "legend_names", "$J_m$")

    if add_timebars:
        d_time = event.get("t.d_time")
        d_start = event.get("t.d_start")
        d_stop = event.get("t.d_end")

        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_ids_properties:
        thickness = event.get("L_k")
        current_density = event.get("j0_k")
        title += "#Discontinuity properties# "
        if thickness:
            title += rf"$L: {thickness:.2f} \mathrm{{km}}$"
        if current_density:
            title += rf", $j: {current_density:.2f} \mathrm{{nA/m}}^2$"
    
    if add_plasma_params:
        plasma_speed = event.get("plasma_speed")
        plasma_density = event.get("plasma_density")
        plasma_temperature = event.get("plasma_temperature")

        title += "\n#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)
        
    for tvar2plot in tvars2plot:
        options(tvar2plot, "thick", 2)
        options(tvar2plot, "char_size", 16)

    fig, axes = tplot(tvars2plot, return_plot_objects=True)
    if isinstance(axes, Axes):
        axes = [axes]

    base_axis = axes[0]

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

        c = event.get("fit.vars.c")
        amp = event.get("fit.vars.amplitude")
        sigma = event.get('fit.vars.sigma')
        
        d_star = event.get("d_star")
        rsquared = event.get("fit.stat.rsquared")
        chisqr = event.get("fit.stat.chisqr")
        
        base_axis.plot(d_time, c + amp / 2, marker="o", markersize=10, color="red")
        if fit_time is not None and fit_data is not None:
            base_axis.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}"
        title += f"\n#Fit# $c$: {c:.2f}, $Amp$: {amp:.2f}, $\Sigma$: {sigma:.2f}"

    # add title to the first plot
    base_axis.set_title(title)

    return fig, axes
