import os
from eval import read_run_dir
from plot_utils import collect_results

run_dir = "../models_skill"

df = read_run_dir(run_dir)

# task = "cot_skill_final"
# # run_id = 'b8fc4ae3-d360-461b-a404-ede7b97bb503'
# run_id = 'cda7dda2-77fb-4aca-9c60-6e8cc2a81874'
# run_id = '79f76c17-ab2b-480f-9375-971a2f972435'

# # run_id = '1380be44-09bf-42ed-89ff-7d627b012741'
# # run_id = '30b8b550-b4e3-4aed-86bc-ce03c6b0a4e1'
# # run_id = '856bb9d7-29dd-4ee9-980d-d13f4296a77c'

# run_id = '0924cf52-d8db-4164-8cc2-5a9b25ce4574'
# run_id = 'random_f10'
# run_id = '76f2ab8f-3706-4ec7-8b62-1f307a9a75d3'
# run_id = '2daf5935-3ca1-46bf-88ea-55e551d79008'
# run_id = '43f05441-2965-442e-b1c6-08284aab1a99'

task = "cot_skill_supp"
# run_id = 'cot1_d5_f4c6_final'
# run_id = 'cot2_d5_f4c6_final'
# run_id = 'cot3_d5_f4c6_final'
# run_id = 'cot6_d5_f4c6_final'
run_ids = ['s10f1_random','s10f2_random','s10f3_random','s10f4_random']
# steps = [5000 * i for i in range(21)]
steps = [-1]
# steps = [10000]
for run_id in run_ids:
    for step in steps:
        run_path = os.path.join(run_dir, task, run_id)

        def valid_row(r):
            return r.task == task and r.run_id == run_id

        metrics = collect_results(run_dir, df, valid_row=valid_row, step=step)

        print('DONE!')

"""
@ks: TODO: script to evaluate models and move files to the write places
task = 'cot_4nn_deep'
run_ids = df[df['task'] == task]['run_id'].tolist()

for run_id in run_ids:
    # TODO: update this to do the directory things also
    run_path = os.path.join(run_dir, task, run_id)

    def valid_row(r):
        return r.task == task and r.run_id == run_id

    metrics = collect_results(run_dir, df, valid_row=valid_row)

    new_dir_name = df[df['run_id'] == run_id]['run_name'].values[0].replace(", ", "_").replace("=", "")
    src_dir = "../models/{}/{}".format(task, run_id)
    dest_dir = "../models/cot_deep/{}".format(new_dir_name)
    new_dir = "../models/cot_deep/{}".format(new_dir_name)

    if not os.path.exists(new_dir):
        os.mkdir(new_dir)
    print("Copying config and results from {} to {}".format(src_dir, dest_dir))

    shutil.copy2("{}/config.yaml".format(src_dir), "{}/config.yaml".format(dest_dir))
    shutil.copy2("{}/metrics.json".format(src_dir), "{}/metrics.json".format(dest_dir))

print('DONE!')
"""
