"""
You must run rtmodel_wrapper_for_roman_simulations.py first to generate the input CSV file.
"""
import pandas as pd
import numpy as np


def chi2_cut(single_lens_chi2, binary_lens_chi2, threshold=500):
    """
    Check if the difference between the single lens model chi2 and the best binary lens model chi2 exceeds the threshold.
    """
    return (single_lens_chi2 - binary_lens_chi2) > threshold


def acceptable_mass_ratio_range(model_mass_ratio, true_mass_ratio, log10_threshold=0.301):
    """
    Check if the absolute log10 difference between the model mass ratio and the true mass ratio is within the threshold 0.301
    10^0.301 ~ 2 . So model_q should not exceed 2*true_q, and vice versa.
    """
    ratio = model_mass_ratio / true_mass_ratio
    return np.abs(np.log10(ratio)) <= log10_threshold


def best_model_generator(dataframe):
    # best chi2 of Ls and LX (Exclude LO cause crazy orbits for now)
    dataframe['best_chi2'] = dataframe[['LS_chi2', 'LX_chi2']].min(axis=1)
    # name of the lowest chi2 models:
    dataframe['best_chi2_name'] = dataframe[['LS_chi2', 'LX_chi2']].idxmin(axis=1)
    # save the best mass ratio according to the lowest chi2 model:
    dataframe['best_mass_ratio'] = dataframe.apply(
        lambda row: row[f"{str(row['best_chi2_name']).replace('_chi2', '')}_mass_ratio"],
        axis=1)
    return dataframe


def main(file_path_, run_version_, root_path_, chi2_thresholds=None, saving_new_df=False):
    if chi2_thresholds is None:
        chi2_thresholds = [None, 0, 100, 200, 300, 400, 500, 600, 900, 1000]
    df = pd.read_csv(file_path_)
    complete_df = best_model_generator(df)
    print('Total number of events in the original DataFrame:', len(complete_df))
    for chi2_threshold in chi2_thresholds:
        if chi2_threshold is None:
            passed_chi2_df = complete_df.copy()
        else:
            complete_df[f'chi2_cut_{chi2_threshold}'] = chi2_cut(complete_df['PS_chi2'], complete_df['best_chi2'],
                                                                 threshold=chi2_threshold)
            # NEW DF for only those that passed the chi2 cut
            passed_chi2_df = complete_df[complete_df[f'chi2_cut_{chi2_threshold}']].copy()
            print('Number of events that had a best_PS_fit chi2 - best_binary_lens_fit chi2 > ', chi2_threshold, ' :',
                  len(passed_chi2_df))
            print('Number of events that did NOT pass the chi2 cut', chi2_threshold, ' :',
                  len(complete_df[~complete_df[f'chi2_cut_{chi2_threshold}']].copy()))
        # Now apply mass ratio cut for those that passed chi2 cut
        passed_chi2_df[f'acceptable_mass_ratio_range_chi2_{chi2_threshold}'] = acceptable_mass_ratio_range(
            passed_chi2_df['best_mass_ratio'],
            passed_chi2_df['true_mass_ratio'],
            log10_threshold=0.301
        )
        print('Of them, number of events within the mass ratio acceptable range:',
              len(passed_chi2_df[passed_chi2_df[f'acceptable_mass_ratio_range_chi2_{chi2_threshold}']]),
              f", {100 * len(passed_chi2_df[passed_chi2_df[f'acceptable_mass_ratio_range_chi2_{chi2_threshold}']]) / len(passed_chi2_df):.2f}%")
        print('Of them, number of events that exceed the mass ratio ratio cut / not acceptable:',
              len(passed_chi2_df[~passed_chi2_df[f'acceptable_mass_ratio_range_chi2_{chi2_threshold}']]),
              f", {100 * len(passed_chi2_df[~passed_chi2_df[f'acceptable_mass_ratio_range_chi2_{chi2_threshold}']]) / len(passed_chi2_df):.2f}%")
        print()
        if saving_new_df:
            new_df = passed_chi2_df[['event_name', 'true_mass_ratio', 'LS_mass_ratio', 'LX_mass_ratio',
                                     'PS_chi2', 'PX_chi2', 'LS_chi2', 'LX_chi2']].copy()

            output_csv = f"{root_path_}/{run_version_}_chi2_threshold_{chi2_threshold}.csv"
            new_df.to_csv(output_csv, index=False)
            print(f"Saved filtered DataFrame to {output_csv}")
            output_csv_2 = f"{root_path_}/{run_version_}_chi2_threshold_{chi2_threshold}_acceptable_mass_ratio.csv"
            acceptable_q_df = new_df[passed_chi2_df[f'acceptable_mass_ratio_range_chi2_{chi2_threshold}']]
            acceptable_q_df.to_csv(output_csv_2, index=False)
            print(f"Saved filtered DataFrame to {output_csv_2}")
            output_csv_3 = f"{root_path_}/{run_version_}_chi2_threshold_{chi2_threshold}_not_acceptable_mass_ratio.csv"
            not_acceptable_q_df = new_df[~passed_chi2_df[f'acceptable_mass_ratio_range_chi2_{chi2_threshold}']]
            not_acceptable_q_df.to_csv(output_csv_3, index=False)
            print(f"Saved filtered DataFrame to {output_csv_3}")


if __name__ == "__main__":
    # Load the CSV file generated by `rtmodel_wrapper_for_roman_simulations.py`
    run_version = 'v31'
    print(run_version, '\n')
    # ADAPT THIS
    if run_version == 'v31':
        root_path = '/Users/stela/Documents/Scripts/orbital_task/RTModel_runs/sample_rtmodel_v3.1'
        file_path = f'{root_path}/v3.1_true_and_rtmodel_fits.csv'
    elif run_version == 'v24':
        root_path = '/Users/stela/Documents/Scripts/orbital_task/RTModel_runs/sample_rtmodel_v2.4'
        file_path = f'{root_path}/true_and_rtmodel_fits.csv'
    elif run_version == 'ICGS':
        root_path = '/Users/stela/Documents/Scripts/orbital_task/RTModel_runs/154_failures_v24_v31'
        file_path = f'{root_path}/ICGS_true_and_rtmodel_fits_done.csv'
    else:
        raise ValueError("Invalid run_version. Use 'v31' or 'v24' or 'ICGS'.")
    main(file_path, run_version, root_path, saving_new_df=True)
