from llm_compiler.utils.util import load_json

class DataType:
    def __init__(self, dataset) -> None:
        self.data_type_path = "data/datatype/" + dataset + "_datatype.json"
        content = load_json(self.data_type_path)
        self.type = content["datatype"]
        self.examples = content["example"]

    def data_annotated_check(self, annotated_data_type):
        return self.type == annotated_data_type
    
    def get_example(self, new_examples):
        examples = self.examples + new_examples
        return "\n".join(examples)
    
    def get_type(self):
        datatype = ""
        for i, t in enumerate(sorted(self.type)):
            discription = self.type[t]
            datatype = datatype + ("" if i != len(self.type)-1 else "and ")
            datatype = datatype + discription + " " + "\"" + t + "\""
            datatype = datatype + (", " if i != len(self.type)-1 else "")
        return datatype