import fire
import json
import pickle
from pathlib import Path
from src.data import load_chia, load_fb
from src.prompt import few_shot_entity_recognition
from tqdm import tqdm


def process_chia(n: int = None, random: bool = False):
    """Processes the Chia dataset

    Args:
        n (int, optional): Number of rows to read. Defaults to None.
        random (bool, optional): Whether to read rows randomly. Defaults to False.
    """
    df = load_chia()

    if random:
        for _, row in df.sample(frac=1.)[:n].iterrows():
            print(row["criteria"])
            print("TRUE: ", row["drugs"], row["persons"], row["conditions"])
            print("PREDICTED: ", few_shot_entity_recognition(row["criteria"]))
            print("-" * 100)
    else:
        # iterate over rows of the dataframe
        for _, row in df[:n].iterrows():
            print(row["criteria"])
            print(row["drugs"], row["persons"], row["conditions"])
            print(few_shot_entity_recognition(row["criteria"]))
            print("-" * 100)


def ner_fb(entity: str, n: int = None, random: bool = False, verbose: bool = False):
    """Applies the LLM prompting to extract NERs from the FB dataset

    Args:
        entity (str): Entity type
        n (int, optional): Number of rows to read. Defaults to None.
        random (bool, optional): Whether to read rows randomly. Defaults to False.
        verbose (bool, optional): Whether to print the results. Defaults to False.
    """
    df = load_fb()["test"]

    results = []

    few_shot_examples = Path("data/few-shots.json")
    with open(few_shot_examples, "r") as f:
        few_shot_examples = json.load(f)[entity]

    if random:
        for _, row in tqdm(df.sample(frac=1.)[:n].iterrows()):
            criterion = row["criterion"]
            ent_true = row[entity]
            ent_pred = few_shot_entity_recognition(few_shot_examples, criterion, entity)

            results.append((entity, criterion, ent_true, ent_pred))
    else:
        for _, row in tqdm(df[:n].iterrows()):
            criterion = row["criterion"]
            ent_true = row[entity]
            ent_pred = few_shot_entity_recognition(few_shot_examples, criterion, entity)

            results.append((entity, criterion, ent_true, ent_pred))

    output_file = Path(f"data/{entity}_ner_results.pkl")
    with open(output_file, "wb") as f:
        pickle.dump(results, f)

    if verbose:
        for entity, criterion, ent_true, ent_pred in results:
            print(criterion)
            print("TRUE: ", ent_true)
            print("PREDICTED: ", ent_pred)
            print("-" * 100)


if __name__ == "__main__":
    fire.Fire()