import os
import json
from instructions import Instruction

prompts = [
    "A wug is a robin. A dax is a penguin. Question: Which of these can fly? Answer: wug",
    "A fep is an elephant. A blicket is a mouse. Question: Which of these has a trunk? Answer: fep",
    "A toma is a cat. A bova is a mouse. Question: Which of these can meow? Answer: toma",
]
prompts = "\n".join(prompts)
prompts = "{examples}"

test_stimulus = "A wif is a mouse. A gek is a beaver. Question: Which of these can swim? Answer:"
test_stimulus = "{test_stimulus}"

for comps in ["comps", "comps-qa"]:
    print("="*100)
    print(f"COMPS DATASET: {comps}")
    PATH = f"data/instructions/{comps}"

    for file in os.listdir(PATH):
        filepath = os.path.join(PATH, file)
        with open(filepath, "r") as f:
            instruction_dict = json.load(f)
            print("="*40)
            print("Filename: ", file)
            for domain, instruction_obj in instruction_dict.items():
                instruction = Instruction(**instruction_obj["incontext"])
                instruction.construct(prompts, test_stimulus)
                print(f"\n{'='*10} In-context {domain} {'='*10}")
                print(instruction.string)
                print(f"\n{'='*10} Zero-shot {domain} {'='*10}")
                instruction.zero_shot(**instruction_obj["zero-shot"])
                instruction.construct("", test_stimulus)
                print(instruction.string)
                # print("\n" + "="*40)

    print("="*100)
    for file in os.listdir(PATH):
        filepath = os.path.join(PATH, file)
        with open(filepath, "r") as f:
            instruction_dict = json.load(f)
        instruction_obj = instruction_dict['indomain']
        instruction = Instruction(**instruction_obj['incontext'])

        instruction.construct("{prompt}", "{test_stimulus}")
        instruction.zero_shot(**instruction_obj["zero-shot"])
        instruction.construct("", "{test_stimulus}")
        print(instruction.string)

