import os
import argparse
from cleanfid import fid
import re
import csv

def main(args):
    score = fid.compute_fid(args.dir1, args.dir2)
    fields_list = args.row_prefix.split(',')
    extracted_fields = {}
    # save as excel file
    extracted_fields['algo_name'] = fields_list[0]
    extracted_fields['change'] = fields_list[1]
    extracted_fields['task'] = fields_list[2]
    extracted_fields['config'] = fields_list[3]
    extracted_fields['finetune_algo'] = fields_list[4]
    extracted_fields['finetune_task'] = fields_list[5]
    extracted_fields['finetune_config'] = fields_list[6]
    extracted_fields['prompts_csv'] = fields_list[7]
    extracted_fields['random_seed'] = fields_list[8]
    extracted_fields['type'] = "mscoco-10k"
    extracted_fields['metric'] = 'FID'
    extracted_fields['value'] = score

    file_exists = os.path.isfile(args.results_dir + 'metrics.csv')
    with open(args.results_dir + 'metrics.csv', 'a', newline='') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=extracted_fields.keys())
        if not file_exists:
            writer.writeheader()  # write the header
        writer.writerow(extracted_fields)  # write the data row 
          
    print(f'FID score: {score}')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Compute FID score between two directories.')
    parser.add_argument('--dir1', type=str, required=True, help='Path to the first directory')
    parser.add_argument('--dir2', type=str, required=True, help='Path to the second directory')
    parser.add_argument("--results_dir", type=str, default='/data/cluster_name/scratch/$(whoami)/projects/MACE-Update/experiments/experimental_results.csv')
    parser.add_argument("--row_prefix", type=str, default='GCD')
    args = parser.parse_args()
    main(args)
