# %%
import json
import os
import shutil
from functools import lru_cache
from pathlib import Path
from textwrap import dedent
from typing import cast

import fire
import yaml
from dotenv import load_dotenv
from goodfire import Client
from inspect_ai import Task, eval, task
from inspect_ai.dataset import Sample
from inspect_ai.scorer import match
from inspect_ai.solver import generate, prompt_template

from deception_detection.detectors import LogisticRegressionDetector
from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.utils import REPO_ROOT

# %%


def run_shift(base_path: str | Path, use_all_features: bool = False):
    with open(Path(base_path) / "sae_feature_labels.json") as file:
        feature_labels = json.load(file)

    whitelist = [
        int(idx)
        for idx, info in feature_labels.items()
        if info["deception_related"] or use_all_features
    ]

    print(f"Using {len(whitelist)} features")

    default_config_path = REPO_ROOT / "deception_detection/scripts/configs/sae_rp.yaml"
    with open(default_config_path) as file:
        config_dict = yaml.safe_load(file)
    config_dict["id"] = "shift"
    config_dict["sae_latent_whitelist"] = whitelist

    config = ExperimentConfig(**config_dict)
    shift_exp = Experiment(config)
    shift_exp.get_detector()
    shift_exp.get_eval_scores()
    shift_exp.get_control_scores()
    # copy feature labels from base to exp folder
    shutil.copy(
        Path(base_path) / "sae_feature_labels.json",
        shift_exp.cfg.folder / "sae_feature_labels.json",
    )
    return shift_exp


@lru_cache
def get_goodfire_client():
    load_dotenv()
    api_key = os.getenv("GOODFIRE_API_KEY")
    assert api_key is not None
    return Client(api_key)


def get_experiment(exp: Experiment | str | Path) -> Experiment:
    if isinstance(exp, Experiment):
        return exp
    return Experiment(ExperimentConfig.from_path(Path(exp)))


def save_feature_labels(exp: Experiment | str | Path, k: int = 1000):
    exp = get_experiment(exp)
    detector = exp.get_detector()
    assert isinstance(detector, LogisticRegressionDetector)
    assert detector.directions is not None
    weights = detector.directions[0]
    topk = weights.abs().topk(k=k)
    feature_idxs = cast(list[int], topk.indices.tolist())
    client = get_goodfire_client()
    result = client.features.lookup(indices=feature_idxs, model="meta-llama/Llama-3.3-70B-Instruct")
    results_to_save = {
        idx: {"weight": weights[idx].item(), "label": result[idx].label} for idx in feature_idxs
    }
    with open(exp.cfg.folder / "sae_feature_labels.json", "w") as file:
        json.dump(results_to_save, file)


PROMPT = dedent(
    """The following is a description of a "feature" from a sparse auto encoder.
    Please classify if the description relates to deception or honesty.
    Only output "Yes" or "No" without explanation.

    Feature description: {label}"""
)


def classify_feature_labels(exp: Experiment | str | Path):
    """Classify SAE feature labels as deception-related or not using an LLM."""
    exp = get_experiment(exp)
    load_dotenv()

    with open(exp.cfg.folder / "sae_feature_labels.json") as file:
        feature_labels = json.load(file)

    # Create samples for each feature
    samples = [
        Sample(
            input=PROMPT.format(label=data["label"]),
            target=["Yes", "No"],  # Accept either Yes or No as valid
            metadata={"idx": idx},
        )
        for idx, data in feature_labels.items()
    ]

    @task
    def classify_features():
        return Task(
            dataset=samples, solver=[prompt_template("{prompt}"), generate()], scorer=match()
        )

    # Run classification
    logs = eval(classify_features(), model="openai/gpt-4o-2024-08-06", max_connections=40)[0]

    # Update feature labels with classifications
    assert logs.samples is not None
    for sample in logs.samples:
        is_related = sample.output.completion.strip().lower().startswith("yes")
        feature_labels[sample.metadata["idx"]]["deception_related"] = is_related

    # Save updated feature labels
    with open(exp.cfg.folder / "sae_feature_labels.json", "w") as file:
        json.dump(feature_labels, file, indent=2)


if __name__ == "__main__":
    fire.Fire(
        {"labels": save_feature_labels, "classify": classify_feature_labels, "run": run_shift}
    )
