import os
import json
import numpy as np
import argparse
import pathlib
from collections import Counter

from load_aokvqa import load_aokvqa


parser = argparse.ArgumentParser()
parser.add_argument('--aokvqa-dir', type=pathlib.Path, required=True, dest='aokvqa_dir')
parser.add_argument('--split', type=str, choices=['train', 'val', 'test'], required=True)
parser.add_argument('--mc', action='store_true', dest='multiple_choice')
parser.add_argument('--out', type=argparse.FileType('w'), required=True, dest='output_file')
args = parser.parse_args()

np.random.seed(0)

train_set = load_aokvqa(args.aokvqa_dir, 'train')
train_freq = dict(Counter(
    [d['choices'][d['correct_choice_idx']] for d in train_set]
))

if args.multiple_choice is False:
    choices = list(train_freq.keys())
    probs = [f / len(train_set) for f in train_freq.values()]

##

predictions = {}

eval_set = load_aokvqa(args.aokvqa_dir, args.split)

for d in eval_set:
    if args.multiple_choice:
        choices = d['choices']
        probs = [train_freq.get(c, 0) for c in choices]
        if probs == [0, 0, 0, 0]:
            probs = [1, 1, 1, 1]
        probs = [p / sum(probs) for p in probs]

    q = d['question_id']
    predictions[q] = np.random.choice(choices, size=1, p=probs)[0]

json.dump(predictions, args.output_file)