import os, copy, json
from encode_experts import encode_expert_dict

for method in ["ascii", "caesar", "morse", "atbash"]:
    expert = encode_expert_dict[method]
    
    model_answer_dir = f"../model_answer_{method}"
    
    for f in os.listdir(model_answer_dir):
        if f.endswith(".jsonl") and not f.endswith("_decoded.jsonl"):# and os.path.notexists(f"{model_answer_dir}/{f.split(".jsonl")[0]}_decoded.jsonl"):
            print(f"Decoding {model_answer_dir}/{f}...")
            with open (f"{model_answer_dir}/{f.split('.jsonl')[0]}_decoded.jsonl", "w") as outfile:
                with open(f"{model_answer_dir}/{f}", "r") as infile:
                    for line in infile:
                        line = json.loads(line)
                        line_decoded = copy.deepcopy(line)
                        content_decoded = expert.decode(line["choices"][0]["turns"][0])
                        line_decoded["choices"][0]["turns"][0] = content_decoded
                        outfile.write(json.dumps(line_decoded) + "\n")





for method in ["translate-ml", "translate-ta", "translate-mr", "translate-zh-cn", "translate-fr"]:
    from googletrans import Translator # Need to create a new environment for "pip install googletrans==3.1.0a0", since "googletrans" is not compatiable with "openai"
    translator = Translator()

    model_answer_dir = f"../model_answer_{method}"
    
    for f in os.listdir(model_answer_dir):
        if f.endswith(".jsonl") and not f.endswith("_translated_to_en.jsonl"):# and os.path.notexists(f"{model_answer_dir}/{f.split(".jsonl")[0]}_decoded.jsonl"):
            print(f"Translating {model_answer_dir}/{f} to English...")
            outfile_name = f"{model_answer_dir}/{f.split('.jsonl')[0]}_translated_to_en.jsonl"
            if os.path.exists(outfile_name):
                print(f"\t{outfile_name} already exists, skipping...")
                continue
            with open (outfile_name, "w") as outfile:
                with open(f"{model_answer_dir}/{f}", "r") as infile:
                    for line in infile:
                        line = json.loads(line)
                        line_decoded = copy.deepcopy(line)
                        content_decoded = translator.translate(line["choices"][0]["turns"][0], dest='en', src=method.split("translate-")[1]).text
                        line_decoded["choices"][0]["turns"][0] = content_decoded
                        outfile.write(json.dumps(line_decoded) + "\n")