"""
This module is for extracting a single pole zero constant from the decay tail
"""

from __future__ import annotations

import json
import logging
import os
import pathlib
import pickle as pkl
from collections import OrderedDict

import matplotlib as mpl

mpl.use("agg")
import lgdo
import lgdo.lh5_store as lh5
import matplotlib.pyplot as plt
import numpy as np

import pygama.math.histogram as pgh
import pygama.math.peak_fitting as pgf
import pygama.pargen.cuts as cts
import pygama.pargen.dsp_optimize as opt
import pygama.pargen.energy_optimisation as om

log = logging.getLogger(__name__)


def load_data(
    raw_file: list[str],
    lh5_path: str,
    pulser_mask=None,
    n_events: int = 10000,
    threshold: int = 5000,
    wf_field: str = "waveform",
) -> lgdo.Table:
    sto = lh5.LH5Store()
    df = lh5.load_dfs(raw_file, ["daqenergy", "timestamp"], lh5_path)

    if pulser_mask is None:
        pulser_props = cts.find_pulser_properties(df, energy="daqenergy")
        if len(pulser_props) > 0:
            final_mask = None
            for entry in pulser_props:
                e_cut = (df.daqenergy.values < entry[0] + entry[1]) & (
                    df.daqenergy.values > entry[0] - entry[1]
                )
                if final_mask is None:
                    final_mask = e_cut
                else:
                    final_mask = final_mask | e_cut
            ids = final_mask
            log.debug(f"pulser found: {pulser_props}")
        else:
            log.debug("no_pulser")
            ids = np.zeros(len(df.daqenergy.values), dtype=bool)
    else:
        ids = pulser_mask

    cuts = np.where((df.daqenergy.values > threshold) & (~ids))[0]

    waveforms = sto.read_object(
        f"{lh5_path}/{wf_field}", raw_file, idx=cuts, n_rows=n_events
    )[0]
    baseline = sto.read_object(
        f"{lh5_path}/baseline", raw_file, idx=cuts, n_rows=n_events
    )[0]
    tb_data = lh5.Table(col_dict={f"{wf_field}": waveforms, "baseline": baseline})
    return tb_data


def get_decay_constant(
    slopes: np.array, wfs: lgdo.WaveformTable, display: int = 0
) -> dict:
    """
    Finds the decay constant from the modal value of the tail slope after cuts
    and saves it to the specified json.

    Parameters
    ----------
    slopes : array
        tail slope array

    dict_file : str
        path to json file to save decay constant value to.
        It will be saved as a dictionary of form {'pz': {'tau': decay_constant}}

    Returns
    -------
    tau_dict : dict
    """
    tau_dict = {}

    pz = tau_dict.get("pz")

    counts, bins, var = pgh.get_hist(slopes, bins=100000, range=(-0.01, 0))
    bin_centres = pgh.get_bin_centers(bins)
    high_bin = bin_centres[np.argmax(counts)]
    try:
        pars, cov = pgf.gauss_mode_width_max(
            counts,
            bins,
            n_bins=10,
            cost_func="Least Squares",
            inflate_errors=False,
            gof_method="var",
        )
        if np.abs(np.abs(pars[0] - high_bin) / high_bin) > 0.05:
            raise ValueError
        high_bin = pars[0]
    except:
        pass
    tau = round(-1 / (high_bin), 1)

    sampling_rate = wfs["dt"].nda[0]
    units = wfs["dt"].attrs["units"]
    tau = f"{tau*sampling_rate}*{units}"

    tau_dict["pz"] = {"tau": tau}
    if display > 0:
        out_plot_dict = {}
        plt.rcParams["figure.figsize"] = (10, 6)
        plt.rcParams["font.size"] = 8
        fig, ax = plt.subplots()
        bins = np.linspace(-0.01, 0, 100000)  # change if needed
        counts, bins, bars = ax.hist(slopes, bins=bins, histtype="step")
        plot_max = np.argmax(counts)
        in_min = plot_max - 20
        if in_min < 0:
            in_min = 0
        in_max = plot_max + 21
        if in_max >= len(bins):
            in_min = len(bins) - 1
        plt.xlabel("Slope")
        plt.ylabel("Counts")
        plt.yscale("log")
        axins = ax.inset_axes([0.5, 0.45, 0.47, 0.47])
        axins.hist(
            slopes[(slopes > bins[in_min]) & (slopes < bins[in_max])],
            bins=200,
            histtype="step",
        )
        axins.axvline(high_bin, color="red")
        axins.set_xlim(bins[in_min], bins[in_max])
        labels = ax.get_xticklabels()
        ax.set_xticklabels(labels=labels, rotation=45)
        out_plot_dict["slope"] = fig
        if display > 1:
            plt.show()
        else:
            plt.close()
        return tau_dict, out_plot_dict
    else:
        return tau_dict


def fom_dpz(tb_data, verbosity=0, rand_arg=None):
    std = tb_data["pz_std"].nda
    counts, start_bins, var = pgh.get_hist(std, dx=0.1, range=(0, 400))
    max_idx = np.argmax(counts)
    mu = start_bins[max_idx]
    try:
        pars, cov = pgf.gauss_mode_width_max(
            counts,
            start_bins,
            mode_guess=mu,
            n_bins=10,
            cost_func="Least Squares",
            inflate_errors=False,
            gof_method="var",
        )

        mu = pars[0]

    except:
        mu = start_bins[max_idx]

    return {"y_val": np.abs(mu)}


def get_dpz_consts(grid_out, opt_dict):
    std_grid = np.ndarray(shape=grid_out.shape)
    for i in range(grid_out.shape[0]):
        for j in range(grid_out.shape[1]):
            std_grid[i, j] = grid_out[i, j]["y_val"]
    min_val = np.amin(std_grid)
    min_point = np.where(std_grid == np.amin(std_grid))

    opt_name = list(opt_dict.keys())[0]
    keys = list(opt_dict[opt_name].keys())
    param_list = []
    shape = []
    db_dict = {}
    for key in keys:
        param_dict = opt_dict[opt_name][key]
        grid_axis = np.arange(
            param_dict["start"], param_dict["end"], param_dict["spacing"]
        )
        unit = param_dict.get("unit")
        param_list.append(grid_axis)
        shape.append(len(grid_axis))
    for i, key in enumerate(keys):
        unit = opt_dict[opt_name][key].get("unit")

        if unit is not None:
            try:
                db_dict[opt_name].update(
                    {key: f"{param_list[i][min_point[i]][0]}*{unit}"}
                )
            except:
                db_dict[opt_name] = {key: f"{param_list[i][min_point[i]][0]}*{unit}"}
        else:
            try:
                db_dict[opt_name].update({key: f"{param_list[i][min_point[i]][0]}"})
            except:
                db_dict[opt_name] = {key: f"{param_list[i][min_point[i]][0]}"}
    return db_dict


def dsp_preprocess_decay_const(
    tb_data,
    dsp_config: dict,
    double_pz: bool = False,
    display: int = 0,
    opt_dict: dict = None,
    wf_field: str = "waveform",
    wf_plot: str = "wf_pz",
    norm_param: str = "pz_mean",
    cut_parameters: dict = {"bl_mean": 4, "bl_std": 4, "bl_slope": 4},
) -> dict:
    """
    This function calculates the pole zero constant for the input data

    Parameters
    ----------
    f_raw : str
        The raw file to run the macro on
    dsp_config: str
        Path to the dsp config file, this is a stripped down version which just includes cuts and slope of decay tail
    channel:  str
        Name of channel to process, should be name of lh5 group in raw files

    Returns
    -------
    tau_dict : dict
    """

    tb_out = opt.run_one_dsp(tb_data, dsp_config)
    log.debug("Processed Data")
    cut_dict = cts.generate_cuts(tb_out, parameters=cut_parameters)
    log.debug("Generated Cuts:", cut_dict)
    idxs = cts.get_cut_indexes(tb_out, cut_dict)
    log.debug("Applied cuts")
    log.debug(f"{len(idxs)} events passed cuts")
    slopes = tb_out["tail_slope"].nda
    log.debug("Calculating pz constant")
    if display > 0:
        tau_dict, plot_dict = get_decay_constant(
            slopes[idxs], tb_data[wf_field], display=display
        )
    else:
        tau_dict = get_decay_constant(slopes[idxs], tb_data[wf_field])
    if double_pz == True:
        log.debug("Calculating double pz constants")
        pspace = om.set_par_space(opt_dict)
        grid_out = opt.run_grid(
            tb_data, dsp_config, pspace, fom_dpz, tau_dict, fom_kwargs=None
        )
        out_dict = get_dpz_consts(grid_out, opt_dict)
        tau_dict["pz"].update(out_dict["pz"])
    if display > 0:
        tb_out = opt.run_one_dsp(tb_data, dsp_config, db_dict=tau_dict)
        wfs = tb_out[wf_plot]["values"].nda[idxs]
        wf_idxs = np.random.choice(len(wfs), 100)
        if norm_param is not None:
            means = tb_out[norm_param].nda[idxs][wf_idxs]
            wfs = np.divide(wfs[wf_idxs], np.reshape(means, (len(wf_idxs), 1)))
        else:
            wfs = wfs[wf_idxs]
        fig2 = plt.figure()
        for wf in wfs:
            plt.plot(np.arange(0, len(wf), 1), wf)
        plt.axhline(1, color="black")
        plt.axhline(0, color="black")
        plt.xlabel("Samples")
        plt.ylabel("ADU")
        plot_dict["waveforms"] = fig2
        if display > 1:
            plt.show()
        else:
            plt.close()
        return tau_dict, plot_dict
    else:
        return tau_dict
