import os
import dotenv
import argparse
import sys
import pandas as pd
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


FILE = os.getenv('FILE_GEN')


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

    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]}")

    file_manager = FileManager(file = FILE, tree = "Events", 
                               branches = ["q2", "costhetak", "costhetal", "phi"])
    data         = file_manager.extract_data()
    q2, costhetak, costhetal, phi = data[0], data[1], data[2], data[3]

    q2, indices  = select_feature(feature = q2, limits = bin_ranges[Q2BIN])
    costhetak    = costhetak[indices]
    costhetal    = costhetal[indices]
    phi          = phi[indices]

    lower_costhetak_cut = float(os.getenv('LOWER_COSTHETAK_CUT'))
    upper_costhetak_cut = float(os.getenv('UPPER_COSTHETAK_CUT'))
    costhetak, indices  = select_feature(feature = costhetak, limits = (lower_costhetak_cut, upper_costhetak_cut))
    q2                  = q2[indices]
    costhetal           = costhetal[indices]
    phi                 = phi[indices]
        
    angular_data = pd.DataFrame({'ctl': costhetal, 'ctk': costhetak, 'phi': phi})
    angular_data.to_csv(f"angular_fit_generator_bin_{Q2BIN}.csv", index = False)
    x            = B2Kstll('ctl','ctk','phi')
    obs, pdf, _  = x.get_pdf('PWave')
    datazfit     = zfit.Data.from_pandas(df = angular_data, obs = obs)
    nll          = zfit.loss.UnbinnedNLL(model = pdf, data = datazfit)
    minimizer    = zfit.minimize.Minuit()

    result = minimizer.minimize(nll)
    print(result)
    print(result.params, pdf)

    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_generator_level_bin_{Q2BIN}")


if __name__ == "__main__":
    main()