"""Script to output latex code for highlighting tokens in a dialogue."""

from functools import cache
from pathlib import Path

import fire

from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.prompt_scorer import create_prompt_scorer
from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.utils import process_str_tokens


@cache
def get_experiment(exp_path: str | Path) -> Experiment:
    return Experiment(ExperimentConfig.from_path(Path(exp_path)))


def output_latex_coloring(
    experiment_path: str | Path,
    dataset_name: str,
    idx: int,
    extent: float = 6.0,
    max_alpha: float = 0.4,
):
    exp = get_experiment(experiment_path)
    dataset = [d for d in exp.eval_datasets + exp.control_datasets if d.name == dataset_name][0]
    tok_dataset = TokenizedDataset.from_dataset(dataset, exp.tokenizer)
    try:
        all_token_scores = exp.get_eval_scores(all_tokens=True)[dataset_name]
    except KeyError:
        all_token_scores = exp.get_control_scores(all_tokens=True)[dataset_name]
    scorer = create_prompt_scorer("mean", exp, [0.01], default_fpr=0.01)
    normalized_scores = scorer.normalize_scores(all_token_scores)

    str_tokens = process_str_tokens(tok_dataset.str_tokens[idx], escape_newlines=True)
    scores = normalized_scores.scores[idx][-len(str_tokens) :]
    assert tok_dataset.detection_mask is not None
    detect_mask = tok_dataset.detection_mask[idx, -len(str_tokens) :]

    def to_color(score: float) -> str:
        """map scores on diverging color map.
        0=white, -extent=blue!max_alpha, extent=red!max_alpha"""
        alpha = int(min((abs(score) / extent) * max_alpha, max_alpha) * 100)
        color = "red" if score > 0 else "blue"
        return f"{color}!{alpha}"

    dialogue = dataset.dialogues[idx]
    for m in dialogue:
        print(f"[{m.role}]")
        print(m.content)
        print("")
    print("-" * 100)
    for token, score, detect in zip(str_tokens, scores, detect_mask, strict=True):
        if detect:
            print(f"\\hlc[{to_color(score.item())}]{{{token}}}", end="")
        else:
            print(f"{token}", end="")
        if token == "\\n":
            print("")
        if token == "\\n\\n":
            print("\n")

    # print(str_tokens)


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