import os
import pandas as pd
from idmtools_calibra.plotters.site_data_plotter import SiteDataPlotter


def generate_ll_all(calib_manager, num_to_plot=5, iteration=None, ll_all_name=None):
    from idmtools_calibra.process_state import StatusPoint
    from idmtools_calibra.cli.utils import read_calib_data
    calib_data = read_calib_data(calib_manager.calibration_path)
    # in case environment has been changed and new suite_id & suites are generated

    # load all_results
    results = calib_data.get('results')
    all_results = pd.DataFrame()
    if isinstance(results, dict):
        all_results = pd.DataFrame.from_dict(results, orient='columns')

    all_results.reset_index(inplace=True, drop=True)
    top_columns = ['iteration', 'sample', 'total']
    all_results = all_results.reindex(
        columns=(top_columns + list([a for a in all_results.columns if a not in top_columns])))

    # restore last time location
    _iteration = calib_data['iteration']
    if iteration is None:
        current_iteration = calib_data['iteration']
    else:
        current_iteration = iteration

    if current_iteration > _iteration:
        raise Exception(f"LL_all for iteration = {iteration} is not available!")

    # validate current iteration status
    it = calib_manager.state_for_iteration(current_iteration)
    latest_step = it.status if isinstance(it.status, StatusPoint) else StatusPoint[it.status]

    if latest_step is None or latest_step.value < StatusPoint.done.value:
        raise Exception(f"LL_all for iteration = {current_iteration} is not available!")

    # build ll_all.csv
    if ll_all_name is None:
        ll_all_name = "ll_all_final.csv"

    # clean existing one
    ll_all_path = os.path.join(calib_manager.directory, '_plots', ll_all_name)
    if os.path.exists(ll_all_path):
        os.remove(ll_all_path)

    for iteration in range(current_iteration + 1):
        it = calib_manager.state_for_iteration(iteration)
        all_results = all_results.copy(deep=True)
        it.all_results = all_results[all_results.iteration <= iteration]
        it.all_results.set_index('sample', inplace=True)  # important to keep the same format as original

        sp = SiteDataPlotter(num_to_plot=num_to_plot, combine_sites=True)
        sp.iteration_state = it
        sp.write_LL_csv(ll_all_name=ll_all_name)
