#!/usr/bin/env python3
"""
GroundingProjector LLM Experiment
=================================

Validates G = A(AᵀA)⁻¹Aᵀ as a hallucination detector using real LLM hidden states.

Previous approach (inference-time toroidal logit bias on raw token IDs) was disproven --
null result on replication (v7 paper). This approach operates on *embedding vectors*,
not token IDs: fundamentally different mechanism.

Methodology:
  1. Load a small model (Qwen2.5-0.5B) via HuggingFace transformers
  2. Build evidence matrix A from hidden states of known-correct completions
  3. Construct orthogonal projection G = A(AᵀA)⁻¹Aᵀ
  4. For test questions, compute hallucination_score = ||(I-G)x|| / ||x||
  5. Check correlation: do correct answers have lower hallucination scores?

Math mirrors the Rust GroundingProjector in topological-coherence/src/lib.rs.
"""

import sys
import time
import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

MODEL_NAME = "Qwen/Qwen2.5-0.5B"
MAX_NEW_TOKENS = 10  # greedy decode length for answer generation

# ---------------------------------------------------------------------------
# Evidence facts -- known-correct statements from which we extract embeddings.
# We run the model on these and grab the hidden state at the key factual token.
# Format: (full sentence, answer_token_text)
# We extract the hidden state at the *first token* of the answer.
# ---------------------------------------------------------------------------

EVIDENCE_FACTS = [
    ("The capital of France is Paris.", "Paris"),
    ("Water boils at 100 degrees Celsius at sea level.", "100"),
    ("The Earth orbits the Sun.", "Sun"),
    ("Light travels at approximately 300000 kilometers per second.", "300"),
    ("The chemical symbol for gold is Au.", "Au"),
    ("DNA stands for deoxyribonucleic acid.", "de"),
    ("The speed of sound in air is approximately 343 meters per second.", "343"),
    ("Oxygen has the atomic number 8.", "8"),
    ("The Great Wall of China is located in China.", "China"),
    ("Pi is approximately 3.14159.", "3"),
]

# ---------------------------------------------------------------------------
# Test questions -- factual prompts with known answers.
# Format: (prompt, expected_answer_substring, domain)
# ---------------------------------------------------------------------------

TEST_QUESTIONS = [
    # Geography
    ("The capital of Japan is", "Tokyo", "geography"),
    ("The capital of Germany is", "Berlin", "geography"),
    ("The capital of Australia is", "Canberra", "geography"),
    ("The capital of Brazil is", "Brasilia", "geography"),
    ("The largest ocean on Earth is the", "Pacific", "geography"),
    # Science
    ("The chemical symbol for water is", "H2O", "science"),
    ("The boiling point of water in Fahrenheit is", "212", "science"),
    ("The speed of light in a vacuum is approximately", "300", "science"),
    ("The chemical symbol for iron is", "Fe", "science"),
    ("Gravity on Earth accelerates objects at approximately", "9.8", "science"),
    # History
    ("The first president of the United States was", "Washington", "history"),
    ("World War II ended in the year", "1945", "history"),
    ("The Berlin Wall fell in", "1989", "history"),
    ("The Declaration of Independence was signed in", "1776", "history"),
    # Math
    ("The square root of 144 is", "12", "math"),
    ("Two plus two equals", "four", "math"),
    ("The value of pi to two decimal places is", "3.14", "math"),
    # Science (more)
    ("Absolute zero is approximately minus", "273", "science"),
    ("The closest star to Earth is the", "Sun", "science"),
    ("Photosynthesis converts carbon dioxide and water into glucose and", "oxygen", "science"),
    # Geography (more)
    ("The longest river in the world is the", "Nile", "geography"),
    ("Mount Everest is located in the", "Himalayas", "geography"),
    # History (more)
    ("The French Revolution began in", "1789", "history"),
    ("The Titanic sank in the year", "1912", "history"),
]


# ---------------------------------------------------------------------------
# Core: GroundingProjector (mirrors Rust implementation)
# ---------------------------------------------------------------------------

class GroundingProjector:
    """
    Orthogonal projection onto the subspace spanned by evidence embeddings.

    G = A(AᵀA)⁻¹Aᵀ

    where A is (dim x m) with columns = evidence vectors.
    Hallucination score = ||(I - G)x|| / ||x|| in [0, 1].
    """

    def __init__(self, evidence_matrix: np.ndarray):
        """
        Args:
            evidence_matrix: shape (m, dim) -- each row is an evidence embedding.
        """
        m, dim = evidence_matrix.shape
        self.dim = dim
        self.rank = m

        # A is dim x m (columns are evidence vectors)
        A = evidence_matrix.T  # (dim, m)

        # AᵀA is (m, m)
        AtA = A.T @ A

        # Invert AᵀA -- use pseudo-inverse for numerical stability
        # (matches Rust Gauss-Jordan but handles near-singular gracefully)
        try:
            AtA_inv = np.linalg.inv(AtA)
        except np.linalg.LinAlgError:
            print("WARNING: AᵀA is singular, falling back to pseudo-inverse")
            AtA_inv = np.linalg.pinv(AtA)

        # G = A (AᵀA)⁻¹ Aᵀ  -- (dim, dim)
        self.G = A @ AtA_inv @ A.T

        # Verify projection properties
        self._verify()

    def _verify(self):
        """Verify G is idempotent (G² = G) and symmetric (Gᵀ = G)."""
        G2 = self.G @ self.G
        idempotent_err = np.max(np.abs(G2 - self.G))
        symmetric_err = np.max(np.abs(self.G - self.G.T))
        if idempotent_err > 1e-4:
            print(f"WARNING: idempotency error = {idempotent_err:.6f}")
        if symmetric_err > 1e-4:
            print(f"WARNING: symmetry error = {symmetric_err:.6f}")

    def hallucination_score(self, x: np.ndarray) -> float:
        """
        ||(I - G)x|| / ||x|| in [0, 1].
        0 = fully grounded, 1 = fully hallucinated.
        """
        x_norm_sq = np.dot(x, x)
        if x_norm_sq < 1e-12:
            return 0.0
        hallucinated = x - self.G @ x
        h_norm_sq = np.dot(hallucinated, hallucinated)
        return float(np.sqrt(h_norm_sq / x_norm_sq))

    def project_grounded(self, x: np.ndarray) -> np.ndarray:
        """Return Gx -- the component supported by evidence."""
        return self.G @ x

    def project_hallucinated(self, x: np.ndarray) -> np.ndarray:
        """Return (I-G)x -- the component NOT supported by evidence."""
        return x - self.G @ x


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def get_device():
    """Select best available device."""
    if torch.backends.mps.is_available():
        return torch.device("mps")
    elif torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def extract_hidden_state(model, tokenizer, text: str, device: torch.device,
                         target_token: str = None) -> np.ndarray:
    """
    Run model on `text` and extract the last hidden state.

    If target_token is specified, find the first occurrence in the tokenized
    output and return that position's hidden state. Otherwise, return the
    last token's hidden state.

    Returns: numpy array of shape (hidden_dim,)
    """
    inputs = tokenizer(text, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs, output_hidden_states=True)

    # Last layer hidden states: shape (1, seq_len, hidden_dim)
    last_hidden = outputs.hidden_states[-1][0]  # (seq_len, hidden_dim)

    if target_token is not None:
        # Find position of target token in the sequence
        input_ids = inputs["input_ids"][0].tolist()
        # Tokenize the target to get its token ID(s)
        target_ids = tokenizer.encode(target_token, add_special_tokens=False)
        if len(target_ids) > 0:
            target_id = target_ids[0]
            # Find first occurrence
            for pos, tid in enumerate(input_ids):
                if tid == target_id:
                    return last_hidden[pos].cpu().float().numpy()
        # Fallback: also try with leading space (common in BPE tokenizers)
        target_ids_space = tokenizer.encode(" " + target_token, add_special_tokens=False)
        if len(target_ids_space) > 0:
            target_id = target_ids_space[0]
            for pos, tid in enumerate(input_ids):
                if tid == target_id:
                    return last_hidden[pos].cpu().float().numpy()
        # If target not found, fall back to last token
        print(f"  WARNING: target token '{target_token}' not found, using last token")

    # Return last token's hidden state
    return last_hidden[-1].cpu().float().numpy()


def generate_answer(model, tokenizer, prompt: str, device: torch.device,
                    max_new_tokens: int = MAX_NEW_TOKENS) -> tuple:
    """
    Generate answer via greedy decoding. Returns (answer_text, hidden_state).

    hidden_state is the last-layer hidden state at the LAST prompt token
    (i.e., the representation used to predict the first generated token).
    """
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    prompt_len = inputs["input_ids"].shape[1]

    # Get hidden state at last prompt token
    with torch.no_grad():
        outputs = model(**inputs, output_hidden_states=True)
    last_hidden = outputs.hidden_states[-1][0]  # (seq_len, hidden_dim)
    hidden_state = last_hidden[-1].cpu().float().numpy()  # last token

    # Generate answer (greedy)
    with torch.no_grad():
        gen_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            temperature=None,
            top_p=None,
        )

    answer_ids = gen_ids[0][prompt_len:]
    answer_text = tokenizer.decode(answer_ids, skip_special_tokens=True).strip()

    return answer_text, hidden_state


def check_answer(generated: str, expected: str) -> bool:
    """Case-insensitive substring match."""
    return expected.lower() in generated.lower()


# ---------------------------------------------------------------------------
# Main experiment
# ---------------------------------------------------------------------------

def main():
    print("=" * 72)
    print("GroundingProjector Experiment")
    print("Validating G = A(AᵀA)⁻¹Aᵀ on real LLM hidden states")
    print("=" * 72)
    print()

    device = get_device()
    print(f"Device: {device}")
    print(f"Model:  {MODEL_NAME}")
    print()

    # ------------------------------------------------------------------
    # Step 1: Load model and tokenizer
    # ------------------------------------------------------------------
    print("[1/5] Loading model and tokenizer...")
    t0 = time.time()
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True,
        torch_dtype=torch.float32,  # full precision for projection math
    ).to(device)
    model.eval()
    print(f"  Loaded in {time.time() - t0:.1f}s")
    hidden_dim = model.config.hidden_size
    print(f"  Hidden dimension: {hidden_dim}")
    print()

    # ------------------------------------------------------------------
    # Step 2: Extract evidence embeddings
    # ------------------------------------------------------------------
    print("[2/5] Extracting evidence embeddings from known-correct facts...")
    evidence_vectors = []
    for i, (sentence, target) in enumerate(EVIDENCE_FACTS):
        vec = extract_hidden_state(model, tokenizer, sentence, device, target_token=target)
        evidence_vectors.append(vec)
        print(f"  [{i+1}/{len(EVIDENCE_FACTS)}] \"{sentence}\" -> ||h|| = {np.linalg.norm(vec):.2f}")

    evidence_matrix = np.array(evidence_vectors, dtype=np.float64)
    print(f"  Evidence matrix shape: {evidence_matrix.shape}")
    print()

    # ------------------------------------------------------------------
    # Step 3: Build GroundingProjector
    # ------------------------------------------------------------------
    print("[3/5] Building GroundingProjector G = A(AᵀA)⁻¹Aᵀ ...")
    projector = GroundingProjector(evidence_matrix)
    print(f"  Projection matrix shape: {projector.G.shape}")
    print(f"  Rank (evidence vectors): {projector.rank}")

    # Quick sanity: evidence vectors should be (nearly) fully grounded
    print("  Sanity check -- hallucination scores for evidence vectors:")
    for i, (sentence, target) in enumerate(EVIDENCE_FACTS):
        score = projector.hallucination_score(evidence_vectors[i].astype(np.float64))
        print(f"    Evidence[{i}]: score = {score:.6f}  (expected ~0)")
    print()

    # ------------------------------------------------------------------
    # Step 4: Test on questions
    # ------------------------------------------------------------------
    print("[4/5] Running test questions with greedy decoding...")
    print("-" * 72)
    results = []
    for i, (prompt, expected, domain) in enumerate(TEST_QUESTIONS):
        answer, hidden = generate_answer(model, tokenizer, prompt, device)
        correct = check_answer(answer, expected)
        h_score = projector.hallucination_score(hidden.astype(np.float64))
        results.append({
            "prompt": prompt,
            "expected": expected,
            "answer": answer[:40],
            "correct": correct,
            "h_score": h_score,
            "domain": domain,
        })
        status = "CORRECT" if correct else "WRONG"
        print(f"  [{i+1:2d}/{len(TEST_QUESTIONS)}] {status}  h_score={h_score:.4f}  "
              f"| \"{prompt}\" -> \"{answer[:30]}\"")

    print("-" * 72)
    print()

    # ------------------------------------------------------------------
    # Step 5: Analysis
    # ------------------------------------------------------------------
    print("[5/5] Analysis")
    print("=" * 72)

    correct_scores = [r["h_score"] for r in results if r["correct"]]
    wrong_scores = [r["h_score"] for r in results if not r["correct"]]

    n_correct = len(correct_scores)
    n_wrong = len(wrong_scores)
    n_total = len(results)

    print(f"  Total questions:  {n_total}")
    print(f"  Correct answers:  {n_correct} ({100*n_correct/n_total:.1f}%)")
    print(f"  Wrong answers:    {n_wrong} ({100*n_wrong/n_total:.1f}%)")
    print()

    if n_correct > 0:
        mean_correct = np.mean(correct_scores)
        std_correct = np.std(correct_scores)
        print(f"  Hallucination score (correct): {mean_correct:.4f} +/- {std_correct:.4f}")
    else:
        mean_correct = float("nan")
        print("  Hallucination score (correct): N/A (no correct answers)")

    if n_wrong > 0:
        mean_wrong = np.mean(wrong_scores)
        std_wrong = np.std(wrong_scores)
        print(f"  Hallucination score (wrong):   {mean_wrong:.4f} +/- {std_wrong:.4f}")
    else:
        mean_wrong = float("nan")
        print("  Hallucination score (wrong):   N/A (all answers correct)")

    print()

    # Effect size and significance
    if n_correct > 0 and n_wrong > 0:
        delta = mean_wrong - mean_correct
        print(f"  Delta (wrong - correct): {delta:+.4f}")
        if delta > 0:
            print("  Direction: EXPECTED (wrong answers have higher hallucination scores)")
        else:
            print("  Direction: UNEXPECTED (wrong answers have LOWER hallucination scores)")

        # Welch's t-test (does not assume equal variance)
        from scipy import stats
        t_stat, p_value = stats.ttest_ind(wrong_scores, correct_scores, equal_var=False)
        print(f"  Welch's t-test: t = {t_stat:.4f}, p = {p_value:.4f}")
        if p_value < 0.05:
            print("  Result: SIGNIFICANT at p < 0.05")
        else:
            print(f"  Result: NOT significant (p = {p_value:.4f})")

        # Cohen's d
        pooled_std = np.sqrt(
            ((n_correct - 1) * np.var(correct_scores) +
             (n_wrong - 1) * np.var(wrong_scores)) /
            (n_correct + n_wrong - 2)
        )
        if pooled_std > 0:
            cohens_d = delta / pooled_std
            print(f"  Cohen's d: {cohens_d:.4f}")
        print()

        # Mann-Whitney U (non-parametric)
        u_stat, u_p = stats.mannwhitneyu(wrong_scores, correct_scores, alternative="greater")
        print(f"  Mann-Whitney U (one-sided, wrong > correct): U = {u_stat:.1f}, p = {u_p:.4f}")
        print()

    elif n_wrong == 0:
        print("  All answers correct -- cannot compare groups.")
        print("  Individual hallucination scores still measure projection distance.")
        print()

    # Per-domain breakdown
    print("  Per-domain breakdown:")
    domains = sorted(set(r["domain"] for r in results))
    for domain in domains:
        domain_results = [r for r in results if r["domain"] == domain]
        domain_correct = [r["h_score"] for r in domain_results if r["correct"]]
        domain_wrong = [r["h_score"] for r in domain_results if not r["correct"]]
        n_dc = len(domain_correct)
        n_dw = len(domain_wrong)
        mc = np.mean(domain_correct) if domain_correct else float("nan")
        mw = np.mean(domain_wrong) if domain_wrong else float("nan")
        print(f"    {domain:12s}: {n_dc} correct (mean h={mc:.4f}), "
              f"{n_dw} wrong (mean h={mw:.4f})" if not (np.isnan(mc) and np.isnan(mw))
              else f"    {domain:12s}: {n_dc} correct, {n_dw} wrong")

    print()

    # Full results table
    print("  Full results table:")
    print(f"  {'#':>3s}  {'Status':>7s}  {'H-Score':>8s}  {'Domain':>10s}  "
          f"{'Prompt':<35s}  {'Answer':<30s}")
    print("  " + "-" * 100)
    for i, r in enumerate(results):
        status = "OK" if r["correct"] else "WRONG"
        print(f"  {i+1:3d}  {status:>7s}  {r['h_score']:8.4f}  {r['domain']:>10s}  "
              f"{r['prompt'][:35]:<35s}  {r['answer'][:30]:<30s}")

    print()
    print("=" * 72)
    if n_correct > 0 and n_wrong > 0:
        if delta > 0 and p_value < 0.05:
            print("CONCLUSION: GroundingProjector hallucination score correlates with")
            print("factual accuracy. Correct answers have significantly lower scores.")
            print("This provides empirical support for the orthogonal projection approach.")
        elif delta > 0:
            print("CONCLUSION: Direction is correct (wrong > correct) but not yet")
            print(f"statistically significant (p = {p_value:.4f}). Consider more questions")
            print("or more evidence vectors to strengthen the signal.")
        else:
            print("CONCLUSION: Unexpected direction -- further investigation needed.")
            print("The evidence subspace may not capture the relevant factual geometry.")
    elif n_wrong == 0:
        print("CONCLUSION: Model answered all questions correctly. The hallucination")
        print("scores show the projection distance from evidence subspace but no")
        print("wrong/correct comparison is possible. Consider harder questions.")
    print("=" * 72)


if __name__ == "__main__":
    main()
