from pathlib import Path

import fire

from deception_detection.detectors import DirectionDetector
from deception_detection.experiment import Experiment, ExperimentConfig


def print_detector_weights(exp_path: Path | str, k: int = 25):
    cfg = ExperimentConfig.from_path(Path(exp_path))
    experiment = Experiment(cfg)
    detector = experiment.get_detector()
    assert isinstance(detector, DirectionDetector)
    assert cfg.use_local_sae_acts
    assert detector.directions is not None
    weights = detector.directions[0]

    k = min(k, len(weights))
    top_k = weights.abs().topk(k)
    assert cfg.sae_latent_whitelist is not None
    orig_sae_idxs = [cfg.sae_latent_whitelist[i] for i in top_k.indices]

    for sae_idx, weight_idx in zip(orig_sae_idxs, top_k.indices, strict=True):
        print(f"SAE {str(sae_idx).ljust(6)} weight: {+weights[weight_idx].item():.2f}")


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