"""
This module provides routines for calculating and applying quality cuts
"""

from __future__ import annotations

import glob
import json
import logging
import os

import lgdo.lh5 as lh5
import numpy as np
import pandas as pd
from lgdo.types import Table
from scipy import stats

import pygama.math.histogram as pgh
import pygama.math.peak_fitting as pgf
import pygama.pargen.energy_cal as pgc

log = logging.getLogger(__name__)


def get_keys(in_data, parameters):
    out_params = []
    if isinstance(in_data, dict):
        possible_keys = in_data.keys()
    elif isinstance(in_data, list):
        possible_keys = in_data
    for param in parameters:
        for key in possible_keys:
            if key in param:
                out_params.append(key)
    return np.unique(out_params).tolist()


def generate_cuts(
    data: dict[str, np.ndarray], parameters: dict[str, int], rounding: int = 4
) -> dict:
    """
    Finds double sided cut boundaries for a file for the parameters specified

    Parameters
    ----------
    data : lh5 table or dictionary of arrays
                data to calculate cuts on
    parameters : dict
                 dictionary with the parameter to be cut and the number of sigmas to cut at
    """

    output_dict = {}
    if isinstance(data, pd.DataFrame):
        pass
    elif isinstance(data, Table):
        data = {entry: data[entry].nda for entry in get_keys(data, parameters)}
        data = pd.DataFrame.from_dict(data)
    elif isinstance(data, dict):
        data = pd.DataFrame.from_dict(data)
    for par in parameters.keys():
        if isinstance(parameters[par], dict):
            if "Lower Boundary" in list(parameters[par]) or "Upper Boundary" in list(
                parameters[par]
            ):
                output_dict[par] = parameters[par].copy()
                if "Lower Boundary" not in parameters[par]:
                    output_dict[par]["Lower Boundary"] = -np.inf
                if "Upper Boundary" not in parameters[par]:
                    output_dict[par]["Upper Boundary"] = np.inf
                continue
        num_sigmas = parameters[par]
        try:
            all_par_array = data[par].to_numpy()
        except KeyError:
            all_par_array = data.eval(par).to_numpy()
        idxs = (all_par_array > np.nanpercentile(all_par_array, 1)) & (
            all_par_array < np.nanpercentile(all_par_array, 99)
        )
        par_array = all_par_array[idxs]
        bin_width = (
            np.nanpercentile(par_array, 70) - np.nanpercentile(par_array, 50)
        ) / 5

        counts, start_bins, var = pgh.get_hist(
            par_array, range=(np.nanmin(par_array), np.nanmax(par_array)), dx=bin_width
        )
        max_idx = np.argmax(counts)
        mu = start_bins[max_idx]
        try:
            fwhm = pgh.get_fwhm(counts, start_bins)[0]
            guess_sig = fwhm / 2.355

            lower_bound = mu - 10 * guess_sig

            upper_bound = mu + 10 * guess_sig

        except:
            bin_range = 1000

            if max_idx < bin_range:
                lower_bound_idx = 0
            else:
                lower_bound_idx = max_idx - bin_range
            lower_bound = start_bins[lower_bound_idx]

            if max_idx > len(start_bins) - bin_range:
                upper_bound_idx = -1
            else:
                upper_bound_idx = max_idx + bin_range

            upper_bound = start_bins[upper_bound_idx]

        if (lower_bound < np.nanmin(par_array)) or (lower_bound > np.nanmax(par_array)):
            lower_bound = np.nanmin(par_array)
        if (upper_bound > np.nanmax(par_array)) or (upper_bound < np.nanmin(par_array)):
            upper_bound = np.nanmax(par_array)

        try:
            counts, bins, var = pgh.get_hist(
                par_array, bins=200, range=(lower_bound, upper_bound)
            )

            bin_centres = pgh.get_bin_centers(bins)

            fwhm = pgh.get_fwhm(counts, bins)[0]
            mean = float(bin_centres[np.argmax(counts)])
            pars, cov = pgf.gauss_mode_width_max(
                counts,
                bins,
                mode_guess=mean,
                n_bins=20,
                cost_func="Least Squares",
                inflate_errors=False,
                gof_method="var",
            )
            mean = pars[0]
            std = fwhm / 2.355

            if mean < np.nanmin(bins) or mean > np.nanmax(bins):
                raise IndexError
        except IndexError:
            bin_range = 5000

            if max_idx < bin_range:
                lower_bound_idx = 0
            else:
                lower_bound_idx = max_idx - bin_range
            lower_bound = start_bins[lower_bound_idx]

            if max_idx > len(start_bins) - bin_range:
                upper_bound_idx = -1
            else:
                upper_bound_idx = max_idx + bin_range
            upper_bound = start_bins[upper_bound_idx]
            counts, bins, var = pgh.get_hist(
                par_array, bins=200, range=(lower_bound, upper_bound)
            )

            bin_centres = pgh.get_bin_centers(bins)

            fwhm = pgh.get_fwhm(counts, bins)[0]
            mean = float(bin_centres[np.argmax(counts)])
            std = fwhm / 2.355

        if isinstance(num_sigmas, (int, float)):
            num_sigmas_left = num_sigmas
            num_sigmas_right = num_sigmas
        elif isinstance(num_sigmas, dict):
            if "left" in num_sigmas:
                num_sigmas_left = num_sigmas["left"]
            else:
                num_sigmas["left"] = np.inf
                num_sigmas_left = np.inf
            if "right" in num_sigmas:
                num_sigmas_right = num_sigmas["right"]
            else:
                num_sigmas["right"] = np.inf
                num_sigmas_right = np.inf
        upper = float((num_sigmas_right * std) + mean)
        lower = float((-num_sigmas_left * std) + mean)
        output_dict[par] = {
            "Mean Value": round(mean, rounding),
            "Sigmas Cut": num_sigmas,
            "Upper Boundary": round(upper, rounding),
            "Lower Boundary": round(lower, rounding),
        }
    return output_dict


def get_cut_indexes(
    all_data: dict[str, np.ndarray], cut_dict: dict, energy_param: str = "trapTmax"
) -> list[int]:
    """
    Returns a mask of the data, for a single file, that passes cuts based on dictionary of cuts
    in form of cut boundaries above
    Parameters
    ----------
    File : dict or lh5_table
           dictionary of parameters + array such as load_nda or lh5 table of params
    Cut_dict : string
                Dictionary file with cuts
    """

    indexes = None
    keys = cut_dict.keys()
    if isinstance(all_data, pd.DataFrame):
        pass
    elif isinstance(all_data, Table):
        cut_keys = list(cut_dict)
        cut_keys.append(energy_param)
        all_data = {
            entry: all_data[entry].nda for entry in get_keys(all_data, cut_keys)
        }
        all_data = pd.DataFrame.from_dict(all_data)
    elif isinstance(all_data, dict):
        all_data = pd.DataFrame.from_dict(all_data)
    for cut in keys:
        try:
            data = all_data[cut]
        except KeyError:
            data = all_data.eval(cut).to_numpy()
        upper = cut_dict[cut]["Upper Boundary"]
        lower = cut_dict[cut]["Lower Boundary"]
        idxs = (data < upper) & (data > lower) & (~np.isnan(data))
        percent = 100 * len(np.where(idxs)[0]) / len(idxs)
        log.info(f"{percent:.2f}% passed {cut} cut")

        # Combine masks
        if indexes is not None:
            indexes = indexes & idxs

        else:
            indexes = idxs
        log.debug(f"{cut} loaded")
    percent = 100 * len(np.where(indexes)[0]) / len(indexes)
    log.info(f"{percent:.2f}% passed all cuts")
    return indexes


def cut_dict_to_hit_dict(cut_dict, final_cut_field="is_valid_cal"):
    out_dict = {}
    symbols = "/-+*"
    replacewith = "_"
    for i, param in enumerate(cut_dict):
        out_dict[
            f"{''.join(replacewith  if c in symbols else c for c in param).replace('(','').replace(')','')}_cut"
        ] = {
            "expression": f"(a<({param}))&(({param})<b)",
            "parameters": {
                "a": cut_dict[param]["Lower Boundary"],
                "b": cut_dict[param]["Upper Boundary"],
            },
        }
    quality_cut_exp = ""
    for par in list(cut_dict)[:-1]:
        quality_cut_exp += f"({''.join(replacewith  if c in symbols else c for c in par).replace('(','').replace(')','')}_cut)&"
    quality_cut_exp += f"({''.join(replacewith  if c in symbols else c for c in list(cut_dict)[-1]).replace('(','').replace(')','')}_cut)"
    out_dict[final_cut_field] = {"expression": quality_cut_exp, "parameters": {}}
    return out_dict


def find_pulser_properties(df, energy="daqenergy"):
    if np.nanmax(df[energy]) > 8000:
        hist, bins, var = pgh.get_hist(
            df[energy], dx=1, range=(1000, np.nanmax(df[energy]))
        )
        allowed_err = 200
    else:
        hist, bins, var = pgh.get_hist(
            df[energy], dx=0.2, range=(500, np.nanmax(df[energy]))
        )
        allowed_err = 50
    if np.any(var == 0):
        var[np.where(var == 0)] = 1
    imaxes = pgc.get_i_local_maxima(hist / np.sqrt(var), 3)
    peak_energies = pgh.get_bin_centers(bins)[imaxes]
    pt_pars, pt_covs = pgc.hpge_fit_E_peak_tops(
        hist, bins, var, peak_energies, n_to_fit=10
    )
    peak_e_err = pt_pars[:, 1] * 4

    allowed_mask = np.ones(len(peak_energies), dtype=bool)
    for i, e in enumerate(peak_energies[1:-1]):
        i += 1
        if peak_e_err[i] > allowed_err:
            continue
        if i == 1:
            if (
                e - peak_e_err[i] < peak_energies[i - 1] + peak_e_err[i - 1]
                and peak_e_err[i - 1] < allowed_err
            ):
                overlap = (
                    peak_energies[i - 1]
                    + peak_e_err[i - 1]
                    - (peak_energies[i] - peak_e_err[i])
                )
                peak_e_err[i] -= overlap * (
                    peak_e_err[i] / (peak_e_err[i] + peak_e_err[i - 1])
                )
                peak_e_err[i - 1] -= overlap * (
                    peak_e_err[i - 1] / (peak_e_err[i] + peak_e_err[i - 1])
                )

        if (
            e + peak_e_err[i] > peak_energies[i + 1] - peak_e_err[i + 1]
            and peak_e_err[i + 1] < allowed_err
        ):
            overlap = (e + peak_e_err[i]) - (peak_energies[i + 1] - peak_e_err[i + 1])
            total = peak_e_err[i] + peak_e_err[i + 1]
            peak_e_err[i] -= (overlap) * (peak_e_err[i] / total)
            peak_e_err[i + 1] -= (overlap) * (peak_e_err[i + 1] / total)

    out_pulsers = []
    for i, e in enumerate(peak_energies[allowed_mask]):
        if peak_e_err[i] > allowed_err:
            continue

        try:
            e_cut = (df[energy] > e - peak_e_err[i]) & (df[energy] < e + peak_e_err[i])
            df_peak = df[e_cut]

            time_since_last = (
                df_peak.timestamp.values[1:] - df_peak.timestamp.values[:-1]
            )

            tsl = time_since_last[
                (time_since_last >= 0)
                & (time_since_last < np.percentile(time_since_last, 99.9))
            ]

            bins = np.arange(0.1, 5, 0.001)
            bcs = pgh.get_bin_centers(bins)
            hist, bins, var = pgh.get_hist(tsl, bins=bins)

            maxs = pgc.get_i_local_maxima(hist, 45)
            maxs = maxs[maxs > 20]

            super_max = pgc.get_i_local_maxima(hist, 500)
            super_max = super_max[super_max > 20]
            if len(maxs) < 2:
                continue
            else:
                max_locs = np.array([0.0])
                max_locs = np.append(max_locs, bcs[np.array(maxs)])
                if (
                    len(np.where(np.abs(np.diff(np.diff(max_locs))) <= 0.001)[0]) > 1
                    or (np.abs(np.diff(np.diff(max_locs))) <= 0.001).all()
                    or len(super_max) > 0
                ):
                    pulser_e = e
                    period = stats.mode(tsl).mode[0]
                    if period > 0.1:
                        out_pulsers.append((pulser_e, peak_e_err[i], period, energy))

                else:
                    continue
        except:
            continue
    return out_pulsers


def tag_pulsers(df, chan_info, window=0.01):
    df["isPulser"] = 0

    if isinstance(chan_info, tuple):
        chan_info = [chan_info]
    final_mask = None
    for chan_i in chan_info:
        pulser_energy, peak_e_err, period, energy_name = chan_i

        e_cut = (df[energy_name] < pulser_energy + peak_e_err) & (
            df[energy_name] > pulser_energy - peak_e_err
        )
        df_pulser = df[e_cut]

        time_since_last = np.zeros(len(df_pulser))
        time_since_last[1:] = (
            df_pulser.timestamp.values[1:] - df_pulser.timestamp.values[:-1]
        )

        mode_idxs = (time_since_last > period - window) & (
            time_since_last < period + window
        )

        pulser_events = np.count_nonzero(mode_idxs)
        # print(f"pulser events: {pulser_events}")
        if pulser_events < 3:
            return df
        df_pulser = df_pulser[mode_idxs]

        ts = df_pulser.timestamp.values
        diff_zero = np.zeros(len(ts))
        diff_zero[1:] = np.around(np.divide(np.subtract(ts[1:], ts[:-1]), period))
        diff_cum = np.cumsum(diff_zero)
        z = np.polyfit(diff_cum, ts, 1)
        p = np.poly1d(z)

        period = z[0]
        phase = z[1]
        pulser_mod = np.abs(df_pulser.timestamp - phase) % period
        mod = np.abs(df.timestamp - phase) % period

        period_cut = (mod < 0.1) | ((period - mod) < 0.1)  # 0.1)

        if final_mask is None:
            final_mask = e_cut & period_cut
        else:
            final_mask = final_mask | (e_cut & period_cut)

    df.loc[final_mask, "isPulser"] = 1

    return df
