import os
import dotenv
import sys
import argparse
import pandas as pd
import mplhep
import zfit
from prettytable import PrettyTable

dotenv.load_dotenv('../properties.env')

sys.path.insert(0, os.getenv('SYS_PATH'))

from b2kstll.models.angular import B2Kstll
from b2kstll.plot import plot_distributions

from hep_analytics.processing.extract import FileManager
from hep_analytics.processing.transform import select_feature, reweight_feature
from hep_analytics.processing.visualisation import reweight_comparing_plot

FILE_MC  = os.getenv('MC_FILE')
FILE_GEN = os.getenv('GEN_FILE')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--q2bin', dest ='q2bin', default = 0)
    args   = parser.parse_args()
    Q2BIN  = int(args.q2bin)

    mplhep.style.use("LHCb2")

    bin_ranges = [(0.25, 4.00), (4.00, 8.00), (11.00, 12.50), (15.00, 18.00), (1.10, 6.00), (1.1, 2.5), (2.5, 4.0), (4.0, 6.0), (6.0, 8.0)]
    print(f"Selected Q2 Bin Range is {bin_ranges[Q2BIN]}")

    bin_labels = [r"$0.25 < q^2 < 4.00$", r"$4.00 < q^2 < 8.00$", r"$11.00 < q^2 < 12.50$",
                  r"$15.00 < q^2 < 18.00$", r"$1.10 < q^2 < 6.00$",
                  r"$1.1 < q^2 < 2.5$", r"$2.5 < q^2 < 4.0$", r"$4.0 < q^2 < 6.0$", r"$6.0 < q^2 < 8.0$"]

    filemanager    = FileManager(file = FILE_MC, tree = "Events", branches = ["q2", "costhetak", "costhetal", "phi"])
    mc_data        = filemanager.extract_data()
    q2_mc, theta_k_mc, theta_l_mc, phi_mc = mc_data[0], mc_data[1], mc_data[2], mc_data[3]
    q2_mc, indices = select_feature(feature = q2_mc, limits = bin_ranges[Q2BIN])
    phi_mc         = phi_mc[indices]
    theta_l_mc     = theta_l_mc[indices]
    theta_k_mc     = theta_k_mc[indices]
    lower_costhetak_cut = float(os.getenv('LOWER_COSTHETAK_CUT'))
    upper_costhetak_cut = float(os.getenv('UPPER_COSTHETAK_CUT'))
    theta_k_mc, indices  = select_feature(feature = theta_k_mc, limits = (lower_costhetak_cut, upper_costhetak_cut))
    q2_mc          = q2_mc[indices]
    phi_mc         = phi_mc[indices]
    theta_l_mc     = theta_l_mc[indices]

    filemanager = FileManager(file = FILE_GEN, tree = "Events", branches = ["q2", "costhetak", "costhetal", "phi"])
    gen_data    = filemanager.extract_data()
    q2_gen, theta_k_gen, theta_l_gen, phi_gen = gen_data[0], gen_data[1], gen_data[2], gen_data[3]
    q2_gen, indices = select_feature(feature = q2_gen, limits = bin_ranges[Q2BIN])
    phi_gen     = phi_gen[indices]
    theta_l_gen = theta_l_gen[indices]
    theta_k_gen = theta_k_gen[indices]
    theta_k_gen, indices  = select_feature(feature = theta_k_gen, limits = (lower_costhetak_cut, upper_costhetak_cut))
    q2_gen      = q2_gen[indices]
    phi_gen     = phi_gen[indices]
    theta_l_gen = theta_l_gen[indices]
    
    angular_data     = pd.DataFrame({'ctl': theta_l_mc, 'ctk': theta_k_mc, 'phi': phi_mc})
    angular_data.to_csv(f"ang_fit_mc_q2_bins_{Q2BIN}_bin.csv", index = False)
    x                = B2Kstll('ctl','ctk','phi')
    x.set_acc(f"./acc_3d_JpsiKstMC_reweighted_4_bin.yaml")
    obs, pdf, _      = x.get_pdf('PWave')

    q2_mc_weights = reweight_feature(original_feature = q2_mc, target_feature = q2_gen, n_bins = 25)
    reweight_comparing_plot(original_feature = q2_mc, target_feature = q2_gen, weights = q2_mc_weights,
        n_bins = 25, suptitle = f"{bin_labels[Q2BIN]}", titles = ["Q2 MC", "Q2 MC (reweighted)", "Q2 Gen"], save = "reweighted_q2_gen_mc.png")

    datazfit        = zfit.Data.from_pandas(df = angular_data, obs = obs, weights = q2_mc_weights)
    nll             = zfit.loss.UnbinnedNLL(model = pdf, data = datazfit)
    minimizer       = zfit.minimize.Minuit()
    result          = minimizer.minimize(nll)
    param_errors, _ = result.errors()
    print(param_errors)

    info_table = PrettyTable(["Variable", "Value", "Lower Error", "Upper Error"])
    fit_labels = ["AFB", "FL", "S3", "S4", "S5", "S7", "S8", "S9"]
    for name in fit_labels:
        value = result.params[name]["value"]
        lower = result.params[name]["minuit_minos"]["lower"]
        upper = result.params[name]["minuit_minos"]["upper"]
        info_table.add_row([name, value, lower, upper])

    print(info_table)

    plot_distributions(result, suffix = f"accPwavePDF_{Q2BIN}_bin_reweighted_mc")

    
if __name__ == "__main__":
    main()