# %%
import os
from pathlib import Path
from typing import cast

import pandas as pd
import plotly.express as px
import torch
from dotenv import load_dotenv
from goodfire import Client

from deception_detection.activations import PairedActivations
from deception_detection.detectors import LogisticRegressionDetector
from deception_detection.experiment import Experiment, ExperimentConfig

# %%

path = Path("/data/nix_goldowsky_dill/deception-detection/results/33/sae_rp")
exp = Experiment(ExperimentConfig.from_path(path))
detector = exp.get_detector()
assert isinstance(detector, LogisticRegressionDetector)

# %%


# %%

load_dotenv()
api_key = os.getenv("GOODFIRE_API_KEY")
assert api_key is not None
client = Client(api_key)

# %%
assert detector.directions is not None
weights = detector.directions[0]
topk = weights.abs().topk(k=200)
feature_idxs = cast(list[int], topk.indices.tolist())
result = client.features.lookup(indices=feature_idxs, model="meta-llama/Llama-3.3-70B-Instruct")
# %%

for idx in feature_idxs:
    print(f"{weights[idx]:+.2f} {result[idx].label}")

# %%

split_paired_acts = PairedActivations.from_dataset(
    dataset=exp.train_dataset,
    tokenizer=exp.tokenizer,
    model=exp.model,
    detect_only_start_of_turn=exp.cfg.detect_only_start_of_turn,
    detect_only_last_token=exp.cfg.detect_only_last_token,
    layers=exp.cfg.detect_layers,
    use_goodfire_sae_acts=exp.cfg.use_goodfire_sae_acts,
)
honest_fit_acts, deceptive_fit_acts = split_paired_acts.get_all_acts()

# %%

honest_fit_acts.to_dense()
deceptive_fit_acts.to_dense()

# %%


def get_freq_active_tokens(acts: torch.Tensor, idxs: list[int]) -> list[float]:
    ndim = acts.ndim
    return (acts[..., idxs] > 0).float().mean(dim=tuple(range(0, ndim - 1))).tolist()


# %%

honest_freqs = get_freq_active_tokens(honest_fit_acts.get_masked_activations(), feature_idxs)
deceptive_freqs = get_freq_active_tokens(deceptive_fit_acts.get_masked_activations(), feature_idxs)

# %%

# Create a DataFrame for plotting
plot_data = pd.DataFrame(
    {
        "honest_freq": honest_freqs,
        "deceptive_freq": deceptive_freqs,
        "weight": weights[feature_idxs].tolist(),
        "feature": [result[idx].label for idx in feature_idxs],
    }
)

# Create interactive scatter plot
fig = px.scatter(
    plot_data,
    x="honest_freq",
    y="deceptive_freq",
    color="weight",
    hover_data={"feature": True},
    log_x=True,
    log_y=True,
    labels={
        "honest_freq": "Honest Activation Frequency",
        "deceptive_freq": "Deceptive Activation Frequency",
        "weight": "Feature Weight",
    },
    title="Feature Activation Frequencies in Honest vs Deceptive Text",
    color_continuous_scale="RdBu",
    color_continuous_midpoint=0,
)

# Customize the layout
fig.update_layout(showlegend=True, width=1000, height=800)

fig.show()

# %%
