import numpy as np
import torch
from huggingface_hub import hf_hub_download

from liar_detector_v5 import (
    LiarDetectorModel,
    FeatureScaler,
    synth_true,
    apply_lie,
    ID_FAMILIES,
)

REPO = "zeechimp/liar-detector-v5"


def main():
    # 1. Load model and feature scaler
    print(f"Loading {REPO} ...")
    model = LiarDetectorModel.from_pretrained(REPO).eval()
    scaler = FeatureScaler.load(hf_hub_download(REPO, "feature_extractor.json"))
    print(f"  parameters: {sum(p.numel() for p in model.parameters())}")

    # 2. Predict on a batch of synthetic pairs (one per family)
    n_trials = 20
    print(f"\nRunning {n_trials} trials per family ...\n")

    print(f"  {'family':<14s}  {'binary vote':<14s}  {'family vote':<14s}  {'presence vote'}")
    print("  " + "-" * 62)

    for family in ID_FAMILIES:
        bin_hits = 0
        pres_hits = 0
        fam_hits = 0

        for i in range(n_trials):
            rng = np.random.default_rng(1000 + i)
            A_clean = synth_true(rng)
            sx = float(np.std(A_clean)) + 1e-9
            noise_std = rng.uniform(0.05, 0.15) * sx

            A = A_clean + noise_std * rng.standard_normal(len(A_clean))
            B = (apply_lie(A_clean, rng, family)
                 + noise_std * rng.standard_normal(len(A_clean)))

            x = scaler(A, B)
            with torch.no_grad():
                out = model(features=torch.from_numpy(x).unsqueeze(0))

            bin_pred = int(out.binary_probs[0].argmax())
            fam_pred = int(out.family_probs[0].argmax())
            pres_pred = int(out.presence_probs[0].argmax())

            bin_hits += int(bin_pred == 1)            # ground truth: B lies
            pres_hits += int(pres_pred == 1)          # ground truth: somebody lies
            fam_hits += int(fam_pred == ID_FAMILIES.index(family))

        print(f"  {family:<14s}  "
              f"B lies {bin_hits:>2d}/{n_trials}   "
              f"correct {fam_hits:>2d}/{n_trials}   "
              f"somebody lies {pres_hits:>2d}/{n_trials}")


if __name__ == "__main__":
    main()