# REMINDER: this code needs to be rewritten for the new framework. Remove this comment when the code is fully converted.

import json
import random
from lm_eval.base import Dataset
from ..utils import sh

class TriviaQA(Dataset):
    def __init__(self):
        self.download()
    def download(self):
        #pass
        #TODO: don't download if files already there
        sh("""
           mkdir -p data/triviaqa
           wget http://nlp.cs.washington.edu/triviaqa/data/triviaqa-unfiltered.tar.gz -O data/triviaqa/trivia_qa-unfiltered.tar.gz
           tar -xf data/triviaqa/trivia_qa-unfiltered.tar.gz
           mv triviaqa-unfiltered/ data/triviaqa/
           """)

    def has_training_docs(self):
        return True

    def has_validation_docs(self):
        return True

    def has_test_docs(self):
        return True

    def training_docs(self):
        return json.load(open('data/triviaqa/triviaqa-unfiltered/unfiltered-web-train.json'))['Data']

    def validation_docs(self):
        return  json.load(open('data/triviaqa/triviaqa-unfiltered/unfiltered-web-dev.json'))['Data']

    def test_docs(self):
        return  json.load(open('data/triviaqa/triviaqa-unfiltered/unfiltered-web-test.json'))['Data']     
    
    def fewshot_description(self):
        pass
    
    def doc_to_text(self, doc):
        return ''.join(['Q: ', doc['Question'], '\n\n','A: '])

    def doc_to_target(self, doc):
        return doc['Answer']['Aliases'][0]

    # TODO: Implement evaluation code

    # ***IMPORTANT***: this evaluation function needs to be written for the new framework. 
    # For more info, check out the interface in base.py and the example BoolQ implementation in superglue.py. 
    # Remove this comment when the evaluation code is implemented.