import argparse
import os
import lpips
import csv
import numpy as np  # Missing import for np

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('-d0', '--dir1', type=str, default='./imgs/ex_dir0')
parser.add_argument('-d1', '--dir2', type=str, default='./imgs/ex_dir1')
parser.add_argument('-o', '--out', type=str, default='./imgs/example_dists.txt')

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')
parser.add_argument('--use_gpu', action='store_true', default=True, help='turn on flag to use GPU')
parser.add_argument('-v', '--version', type=str, default='0.1')
opt = parser.parse_args()

# Initializing the model
loss_fn = lpips.LPIPS(net='alex', version=opt.version)
if opt.use_gpu:
    loss_fn.cuda()

# Crawl directories
f = open(opt.results_dir + 'lpips_values.csv', 'w')
files = os.listdir(opt.dir1)

fields_list = opt.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'] = os.path.basename(os.path.dirname(opt.dir1))
extracted_fields['metric'] = 'lpips'

lpips_values = []
for file in files:
    if os.path.exists(os.path.join(opt.dir2, file)):
        # Load images
        img0 = lpips.im2tensor(lpips.load_image(os.path.join(opt.dir1, file)))  # RGB image from [-1, 1]
        img1 = lpips.im2tensor(lpips.load_image(os.path.join(opt.dir2, file)))

        if opt.use_gpu:
            img0 = img0.cuda()
            img1 = img1.cuda()

        # Compute distance
        dist01 = loss_fn.forward(img0, img1)
        print('%s, %.3f' % (file, dist01))
        f.writelines('%s, %.6f\n' % (file, dist01))
        lpips_values.append(dist01.cpu().detach().numpy())

f.close()

similarities = np.array(lpips_values)

mean_similarity = np.mean(similarities)
std_similarity = np.std(similarities)
extracted_fields['value'] = f"{mean_similarity:.4f}±{std_similarity:.4f}"

with open(opt.results_dir + 'metrics.csv', 'a', newline='') as csvfile:
    writer = csv.DictWriter(csvfile, fieldnames=extracted_fields.keys())
    writer.writerow(extracted_fields)  # Write the data row
