cutFinder

#!/usr/bin/env python
import multiprocessing
import typer
from typing_extensions import Annotated

app = typer.Typer(
    pretty_exceptions_show_locals=False, rich_markup_mode="rich", add_completion=False
)

@app.command()
def main(
    path: Annotated[
        str, typer.Option("-c", "--config", help="Path to the config file.")
    ],
    output: Annotated[
        str, typer.Option("-o", "--output", help="Path to the output folder.")
    ],
    ncpus: int = typer.Option(
        multiprocessing.cpu_count(),
        "-j",
        "--ncpus",
        help="Number of CPUs to use.",
    ),
    regressor_only: bool = typer.Option(
        False,
        "--regressor-only",
        help="Only perform the regression on pre-computed cuts.",
    ),
):

    from CutFinder.configs import ConfigRef
    from CutFinder.plots import Plotter
    from CutFinder.readers import ConfigReader

    import os
    import json

    import numpy as np

    import ROOT

    ROOT.gInterpreter.Declare('#include "include/functions.cpp"')

    if ncpus > 1:
        ROOT.EnableImplicitMT(ncpus)

    config_reader = ConfigReader(path)
    glob = config_reader.glob
    objs = config_reader.objs
    refs = config_reader.refs
    for ref in refs:
        for obj in objs:
            if obj.refs is not None:
                if ref.name not in obj.refs:
                    continue
            if ((obj.scaling is not None and ref.scaling is None) or
                (obj.scaling is None and ref.scaling is not None)):
                raise ValueError(f"You are comparing Offline and Online pT between obj {obj.name} and ref {ref.name}. Please provide scaling functions for both configurations or none.")

            #inverse scale if scaling is applied
            if obj.scaling is not None:
                cut_bins = obj.inverse_scaling(np.array(glob.pt_bins))
            else:
                cut_bins = glob.pt_bins

            if not regressor_only:
                # if already computed, compute does nothing
                ref.compute()
                obj.compute()
                cuts, cuts_err, rate = glob.algo(ref, obj, glob, **glob.algo_kwargs)

            else:
                #load json record to get cuts
                with open(f"{output}/{obj.name}/records.json","r") as fp:
                    records = json.load(fp)
                ref_rate = records[ref.name]["ref_rate"]
                cut_record = records[ref.name]["full"]
                cut_bins = np.array(cut_record["bins"])
                cuts = np.array(cut_record["cuts"])
                rate= np.array(cut_record["rate"])
                cuts_err = np.array(cut_record.get("cuts_err", np.ones_like(cuts)))

            #create record and save cuts
            obj.addToRecord(ref, "full", cut_bins, cuts, rate, cuts_err=cuts_err)

            # in regressor only mode, save also the ref rate
            if regressor_only:
                obj.records[ref.name]["ref_rate"] = ref_rate

            fitted_cut_bins, fitted_cuts, chi2 = glob.regressor(
                cut_bins,
                cuts,
                sigma = cuts_err,
                fitrange=glob.fitrange,
                **glob.regressor_kwargs,
            )
            fitted_clone = obj.clone(
                ConfigRef, WP=[fitted_cut_bins, fitted_cuts]
            )
            fitted_clone.compute()
            fitted_clone.makeRate(glob.pt_bins, glob.maxRate)

            #create record and save cuts
            obj.addToRecord(
                ref,
                "fitted",
                fitted_cut_bins,
                fitted_cuts,
                fitted_clone.rate,
                chi2=chi2,
            )

    #prepare output folder
    os.makedirs(output, exist_ok=True)
    os.system(f"cp -f {path} {output}/{path.split('/')[-1]}")

    plotter = Plotter()
    for obj in objs:
        os.makedirs(f"{output}/{obj.name}", exist_ok=True)
        #copy index.php
        os.system(f"cp {os.path.join(os.path.dirname(__file__),"externals/index.php")} {output}/{obj.name}/index.php")

        #plots
        plotter.plot_rates(glob, obj, output=output)
        plotter.plot_cuts(obj, output=output)

        #add glob pt_bins to record and save records.json
        if isinstance(glob.pt_bins, np.ndarray):
            obj.records["pt_bins"] = glob.pt_bins.tolist()
        else:
            obj.records["pt_bins"] = glob.pt_bins
        with open(f"{output}/{obj.name}/records.json","w") as fp:
            json.dump(obj.records, fp, indent=4)



if __name__ == "__main__":
    app()