import uproot
import awkward as ak
from pathlib import Path
from scipy.optimize import curve_fit
import numpy as np
import pandas as pd
from math import ceil
import argparse


def fastSigmoid(x, p0, p1, p2):
    return p0 + p1 * x / (1 + abs(p2 * x))


def parameterise_hough_histogram(
    input_file: str = "data/param_data_selected_all_p.root",
    tree_name: str = "Selected",
    n_bins_start: int = 900,
    hist_range: tuple[float, float] = (-3000.0, 3000.0),
    first_bin_center: float = 2.5,
) -> Path:
    """Function to parameterise the binning of the Hough histogram using the occupancy on the reference plane.

    Args:
        input_file (str, optional): Defaults to "data/param_data_selected_all_p.root".
        tree_name (str, optional): Defaults to "Selected".
        n_bins_start (int, optional): Starting (minimal) number of bins in histogram. Defaults to 900.
        hist_range (tuple[float, float], optional): Range in mm the histogram covers. Defaults to (-3000.0, 3000.0).
        first_bin_center (float, optional): Calculated bin center at lower range. Defaults to 2.5.

    Returns:
        Path: Path to cpp code file.
    """
    input_tree = uproot.open({input_file: tree_name})
    # this is an event list of dictionaries containing awkward arrays
    array = input_tree.arrays()
    array = array[[field for field in ak.fields(array) if "scifi_hit" not in field]]
    df = ak.to_dataframe(array)
    selection = (df["x_ref"] > hist_range[0]) & (df["x_ref"] < hist_range[1])
    data = df.loc[selection, "x_ref"]
    _, equal_bins = pd.qcut(data, q=n_bins_start, retbins=True)
    bin_numbering = np.arange(0, n_bins_start + 1)
    equalbins_center = equal_bins[int(n_bins_start / 10) : int(9 * n_bins_start / 10)]
    bin_numbering_center = bin_numbering[
        int(n_bins_start / 10) : int(9 * n_bins_start / 10)
    ]
    func = fastSigmoid
    popt, _ = curve_fit(func, xdata=equalbins_center, ydata=bin_numbering_center)
    print("Parameterisation for central occupancy:")
    print("fastSigmoid(x,", *popt, ")")
    print("Scan shift to match first bin center ...")
    shift = 0.0
    while func(hist_range[0], popt[0] + shift, *popt[1:]) < first_bin_center:
        shift += 0.1
    popt[0] += shift - 0.1
    popt[2] = abs(popt[2])
    print("shifted: fastSigmoid(x,", *popt, ")")
    n_bins_final = ceil(func(hist_range[1], *popt))
    print(
        "Final number of bins:",
        n_bins_final,
        "including offset of",
        int(first_bin_center),
    )
    comment = f"// p[0] + p[1] * x / (1 + p[2] * abs(x)) for nBins = {n_bins_final}\n"
    cpp = (
        "constexpr auto p = std::array{"
        + ", ".join([str(p) + "f" for p in popt])
        + "};"
    )
    outpath = Path("parameterisations/result/hough_histogram_binning_params.hpp")
    outpath.parent.mkdir(parents=True, exist_ok=True)
    with open(outpath, "w") as result:
        result.writelines([comment, cpp])
    return outpath


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input-file",
        type=str,
        help="Path to the input file",
        required=False,
    )
    parser.add_argument(
        "--tree-name",
        type=str,
        help="Path to the input file",
        required=False,
    )
    args = parser.parse_args()
    args_dict = {arg: val for arg, val in vars(args).items() if val is not None}
    outfile = parameterise_hough_histogram(**args_dict)