import os
import argparse
import torch
import json
import evaluation
import numpy as np

from models.clip import clip
from evaluation import PACScore, RefPACScore
from models import open_clip

_MODELS = {
    "ViT-B/32": "checkpoints/clip_ViT-B-32.pth",
    "open_clip_ViT-L/14": "checkpoints/openClip_ViT-L-14.pth"
}


def compute_scores(model, preprocess, image_ids, candidates, references, args):
    gen = {}
    gts = {}

    ims_cs = list()
    gen_cs = list()
    gts_cs = list()
    all_scores = dict()
    model.eval()

    for i, (im_i, gts_i, gen_i) in enumerate(zip(image_ids, references, candidates)):
        gen['%d' % (i)] = [gen_i, ]
        gts['%d' % (i)] = gts_i

        ims_cs.append(im_i)
        gen_cs.append(gen_i)
        gts_cs.append(gts_i)

    gts = evaluation.PTBTokenizer.tokenize(gts)
    gen = evaluation.PTBTokenizer.tokenize(gen)

    all_scores_metrics = evaluation.get_all_metrics(gts, gen)
    
    for k, v in all_scores_metrics.items():
        if k == 'BLEU':
            all_scores['BLEU-1'] = v[0]
            all_scores['BLEU-4'] = v[-1]
        else:
            all_scores[k] = v

    # PAC-S
    _, pac_scores, candidate_feats, len_candidates = PACScore(
        model, preprocess, ims_cs, gen_cs, device, w=2.0)
    all_scores['PAC-S'] = np.mean(pac_scores)

    # RefPAC-S
    if args.compute_refpac:
        _, per_instance_text_text = RefPACScore(
            model, gts_cs, candidate_feats, device, torch.tensor(len_candidates))
        refpac_scores = 2 * pac_scores * per_instance_text_text / \
            (pac_scores + per_instance_text_text)
        all_scores['RefPAC-S'] = np.mean(refpac_scores)

    return all_scores


if __name__ == '__main__':
    # Argument parsing
    parser = argparse.ArgumentParser(description='PAC-S evaluation')
    parser.add_argument('--clip_model', type=str, default='ViT-B/32',
                        choices=['ViT-B/32', 'open_clip_ViT-L/14'])
    parser.add_argument('--image_dir', type=str, default='example/images')
    parser.add_argument('--candidates_json', type=str,
                        default='example/good_captions.json')
    parser.add_argument('--references_json', type=str, default='example/refs.json')
    parser.add_argument('--compute_refpac', action='store_true')

    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"

    image_ids = [img_id for img_id in os.listdir(args.image_dir)]

    with open(args.candidates_json) as f:
        candidates = json.load(f)
    candidates = [candidates[cid.split('.')[0]] for cid in image_ids]

    with open(args.references_json) as f:
        references = json.load(f)
        references = [references[cid.split('.')[0]] for cid in image_ids]

    image_ids = [os.path.join(args.image_dir, img_id) for img_id in image_ids]

    if args.clip_model.startswith('open_clip'):
        print("Using Open CLIP Model: " + args.clip_model)
        model, _, preprocess = open_clip.create_model_and_transforms(
            'ViT-L-14', pretrained='laion2b_s32b_b82k')
    else:
        print("Using CLIP Model: " + args.clip_model)
        model, preprocess = clip.load(args.clip_model, device=device)

    model = model.to(device)
    model = model.float()

    checkpoint = torch.load(_MODELS[args.clip_model])
    model.load_state_dict(checkpoint['state_dict'])
    model.eval()

    all_scores = compute_scores(
        model, preprocess, image_ids, candidates, references, args)
    
    for k, v in all_scores.items():
        print('%s: %.4f' % (k, v))