from transformers import pipeline
import pandas as pd
from tqdm import tqdm
import json
import sys


def main(boolq_path):
    translate = pipeline(task='translation', model='Helsinki-NLP/opus-mt-en-sk', device=0)

    nones = 0
    errors = 0
    too_long = 0

    for split in ['dev', 'train']:
        json_lines = pd.read_json(f'{boolq_path}/{split}.jsonl', lines=True, orient='records')
        with open(f'{boolq_path}/{split}_sk.jsonl', 'w') as f_out:
            for index, json_line in tqdm(list(json_lines.iterrows())):
                if json_line['passage'] is None or json_line['question'] is None:
                    nones += 1
                    print(f'Nones in example: {json_line}')
                    continue

                if len(json_line['passage']) > 512 or len(json_line['question']) > 512:
                    too_long += 1
                    continue

                json_sk = dict()
                try:
                    json_sk['question'] = translate(json_line['question'])[0]['translation_text']\
                        .replace("City name (optional, probably does not need a translation)", "")
                    json_sk['passage'] = translate(json_line['passage'])[0]['translation_text']\
                        .replace("City name (optional, probably does not need a translation)", "")
                    json_sk['answer'] = json_line['answer']
                except:
                    errors += 1
                    print(f"Error in translation: {json_line}")
                    continue

                f_out.write(f"{json.dumps(json_sk)}\n")

        print(f"Errors: {errors}, Nones: {nones}, Too long: {too_long}")


if __name__ == '__main__':
    main(sys.argv[1])
