#!/usr/bin/env python3
"""
φ-Coherence vs Industry Standard Hallucination Detection Benchmark
Abhishek Srivastava | 137-Resonance Logic

Compares φ-Coherence against:
- HHEM (Vectara's Hallucination Evaluation Model)
- SelfCheckGPT-NLI
- Baseline methods

Datasets:
- TruthfulQA (817 questions)
- HaluEval (35,000 samples)

"Truth has structure. Lies are noise."
"""

import json
import time
import argparse
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass, asdict
from collections import defaultdict

# φ-Coherence
from phi_coherence import PhiCoherence, PHI, ALPHA

# Will be imported conditionally
datasets = None
torch = None
transformers = None


@dataclass
class BenchmarkResult:
    method: str
    dataset: str
    subset: str
    accuracy: float
    precision: float
    recall: float
    f1: float
    avg_time_ms: float
    total_samples: int
    true_positives: int
    false_positives: int
    true_negatives: int
    false_negatives: int


def install_dependencies():
    """Check and install required packages."""
    import subprocess
    import sys

    packages = {
        'datasets': 'datasets',
        'torch': 'torch',
        'transformers': 'transformers',
        'numpy': 'numpy',
        'tqdm': 'tqdm',
    }

    for module, package in packages.items():
        try:
            __import__(module)
        except ImportError:
            print(f"[*] Installing {package}...")
            subprocess.check_call([sys.executable, '-m', 'pip', 'install', package, '-q'])


def load_truthfulqa(max_samples: Optional[int] = None) -> List[Dict]:
    """Load TruthfulQA dataset."""
    from datasets import load_dataset

    print("[*] Loading TruthfulQA dataset...")
    ds = load_dataset("truthfulqa/truthful_qa", "multiple_choice", split="validation")

    samples = []
    for i, item in enumerate(ds):
        if max_samples and i >= max_samples:
            break

        # Get question and choices
        question = item['question']
        mc1_targets = item['mc1_targets']

        # mc1_targets has 'choices' and 'labels' (1 for correct, 0 for incorrect)
        choices = mc1_targets['choices']
        labels = mc1_targets['labels']

        # Create samples: correct answers (label=1) are NOT hallucinations
        # incorrect answers (label=0) ARE hallucinations
        for choice, label in zip(choices, labels):
            full_text = f"Question: {question}\nAnswer: {choice}"
            samples.append({
                'text': full_text,
                'is_hallucination': label == 0,  # 0 = incorrect = hallucination
                'source': 'truthfulqa',
                'question': question,
                'answer': choice,
            })

    print(f"[*] Loaded {len(samples)} samples from TruthfulQA")
    return samples


def load_halueval(subset: str = "qa", max_samples: Optional[int] = None) -> List[Dict]:
    """Load HaluEval dataset."""
    from datasets import load_dataset

    print(f"[*] Loading HaluEval dataset (subset: {subset})...")
    ds = load_dataset("pminervini/HaluEval", subset, split="data")

    samples = []
    for i, item in enumerate(ds):
        if max_samples and i >= max_samples:
            break

        if subset == "qa":
            # QA subset has knowledge, question, right_answer, hallucinated_answer
            knowledge = item.get('knowledge', '')
            question = item.get('question', '')
            right_answer = item.get('right_answer', '')
            halluc_answer = item.get('hallucinated_answer', '')

            # Right answer - NOT hallucination
            if right_answer:
                samples.append({
                    'text': f"Context: {knowledge}\nQuestion: {question}\nAnswer: {right_answer}",
                    'is_hallucination': False,
                    'source': 'halueval_qa',
                })

            # Hallucinated answer - IS hallucination
            if halluc_answer:
                samples.append({
                    'text': f"Context: {knowledge}\nQuestion: {question}\nAnswer: {halluc_answer}",
                    'is_hallucination': True,
                    'source': 'halueval_qa',
                })

        elif subset == "summarization":
            document = item.get('document', '')
            right_summary = item.get('right_summary', '')
            halluc_summary = item.get('hallucinated_summary', '')

            if right_summary:
                samples.append({
                    'text': f"Document: {document[:500]}...\nSummary: {right_summary}",
                    'is_hallucination': False,
                    'source': 'halueval_summarization',
                })

            if halluc_summary:
                samples.append({
                    'text': f"Document: {document[:500]}...\nSummary: {halluc_summary}",
                    'is_hallucination': True,
                    'source': 'halueval_summarization',
                })

        elif subset == "dialogue":
            dialogue_history = item.get('dialogue_history', '')
            right_response = item.get('right_response', '')
            halluc_response = item.get('hallucinated_response', '')

            if right_response:
                samples.append({
                    'text': f"Dialogue: {dialogue_history}\nResponse: {right_response}",
                    'is_hallucination': False,
                    'source': 'halueval_dialogue',
                })

            if halluc_response:
                samples.append({
                    'text': f"Dialogue: {dialogue_history}\nResponse: {halluc_response}",
                    'is_hallucination': True,
                    'source': 'halueval_dialogue',
                })

    print(f"[*] Loaded {len(samples)} samples from HaluEval ({subset})")
    return samples


class PhiCoherenceDetector:
    """φ-Coherence hallucination detector."""

    def __init__(self, threshold: float = 0.55):
        self.coherence = PhiCoherence()
        self.threshold = threshold
        self.name = f"φ-Coherence (t={threshold})"

    def predict(self, text: str) -> Tuple[bool, float]:
        """
        Predict if text is hallucination.
        Returns: (is_hallucination, confidence_score)
        """
        score = self.coherence.calculate(text)
        # Lower score = more likely hallucination
        is_hallucination = score < self.threshold
        return is_hallucination, score


class HHEMDetector:
    """Vectara HHEM hallucination detector."""

    def __init__(self, threshold: float = 0.5):
        from transformers import AutoModelForSequenceClassification, AutoTokenizer
        import torch

        self.threshold = threshold
        self.name = f"HHEM-2.1 (t={threshold})"

        print("[*] Loading HHEM model...")
        self.tokenizer = AutoTokenizer.from_pretrained(
            "vectara/hallucination_evaluation_model"
        )
        self.model = AutoModelForSequenceClassification.from_pretrained(
            "vectara/hallucination_evaluation_model",
            trust_remote_code=True
        )
        self.model.eval()

        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model.to(self.device)
        print(f"[*] HHEM loaded on {self.device}")

    def predict(self, text: str) -> Tuple[bool, float]:
        """
        Predict if text is hallucination.
        HHEM outputs: 0 = hallucination, 1 = factual
        """
        import torch

        # HHEM expects premise-hypothesis format for NLI
        # For standalone text, we use the text as both
        inputs = self.tokenizer(
            text, text,
            return_tensors="pt",
            truncation=True,
            max_length=512,
            padding=True
        ).to(self.device)

        with torch.no_grad():
            outputs = self.model(**inputs)
            probs = torch.softmax(outputs.logits, dim=-1)
            # Score closer to 1 = factual, closer to 0 = hallucination
            factual_score = probs[0][1].item()

        is_hallucination = factual_score < self.threshold
        return is_hallucination, factual_score


class LengthBaselineDetector:
    """Simple baseline: shorter texts are more likely hallucinations."""

    def __init__(self, threshold: int = 100):
        self.threshold = threshold
        self.name = f"Length Baseline (t={threshold})"

    def predict(self, text: str) -> Tuple[bool, float]:
        length = len(text)
        score = min(1.0, length / 200)  # Normalize to 0-1
        is_hallucination = length < self.threshold
        return is_hallucination, score


class RandomBaselineDetector:
    """Random baseline for comparison."""

    def __init__(self):
        import random
        self.name = "Random Baseline"
        self.random = random

    def predict(self, text: str) -> Tuple[bool, float]:
        score = self.random.random()
        return score < 0.5, score


def evaluate_detector(
    detector,
    samples: List[Dict],
    verbose: bool = False
) -> BenchmarkResult:
    """Evaluate a detector on samples."""
    from tqdm import tqdm

    tp = fp = tn = fn = 0
    total_time = 0

    iterator = tqdm(samples, desc=detector.name, disable=not verbose)

    for sample in iterator:
        text = sample['text']
        actual_halluc = sample['is_hallucination']

        start = time.time()
        predicted_halluc, score = detector.predict(text)
        elapsed = (time.time() - start) * 1000  # ms
        total_time += elapsed

        if predicted_halluc and actual_halluc:
            tp += 1
        elif predicted_halluc and not actual_halluc:
            fp += 1
        elif not predicted_halluc and not actual_halluc:
            tn += 1
        else:
            fn += 1

    total = len(samples)
    accuracy = (tp + tn) / total if total > 0 else 0
    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
    avg_time = total_time / total if total > 0 else 0

    return BenchmarkResult(
        method=detector.name,
        dataset=samples[0]['source'] if samples else "unknown",
        subset="",
        accuracy=round(accuracy, 4),
        precision=round(precision, 4),
        recall=round(recall, 4),
        f1=round(f1, 4),
        avg_time_ms=round(avg_time, 2),
        total_samples=total,
        true_positives=tp,
        false_positives=fp,
        true_negatives=tn,
        false_negatives=fn,
    )


def find_optimal_threshold(
    detector_class,
    samples: List[Dict],
    thresholds: List[float]
) -> Tuple[float, float]:
    """Find optimal threshold for a detector."""
    best_threshold = 0.5
    best_f1 = 0

    for t in thresholds:
        detector = detector_class(threshold=t)
        result = evaluate_detector(detector, samples, verbose=False)
        if result.f1 > best_f1:
            best_f1 = result.f1
            best_threshold = t

    return best_threshold, best_f1


def print_results_table(results: List[BenchmarkResult]):
    """Print results in a nice table."""
    print("\n" + "=" * 100)
    print(f"{'Method':<30} {'Dataset':<20} {'Accuracy':<10} {'Precision':<10} {'Recall':<10} {'F1':<10} {'Time(ms)':<10}")
    print("=" * 100)

    for r in sorted(results, key=lambda x: x.f1, reverse=True):
        print(f"{r.method:<30} {r.dataset:<20} {r.accuracy:<10.4f} {r.precision:<10.4f} {r.recall:<10.4f} {r.f1:<10.4f} {r.avg_time_ms:<10.2f}")

    print("=" * 100)


def run_benchmark(
    max_samples: int = 500,
    include_hhem: bool = True,
    datasets_to_test: List[str] = ["truthfulqa", "halueval_qa"],
    optimize_thresholds: bool = True,
):
    """Run the full benchmark."""

    print("\n" + "=" * 70)
    print("  φ-COHERENCE HALLUCINATION DETECTION BENCHMARK")
    print("  Comparing against industry standard methods")
    print("=" * 70)
    print(f"\n  Constants: φ = {PHI:.6f}  |  α = {ALPHA}")
    print(f"  Max samples per dataset: {max_samples}")
    print()

    # Load datasets
    all_samples = {}

    if "truthfulqa" in datasets_to_test:
        all_samples["truthfulqa"] = load_truthfulqa(max_samples)

    if "halueval_qa" in datasets_to_test:
        all_samples["halueval_qa"] = load_halueval("qa", max_samples)

    if "halueval_summarization" in datasets_to_test:
        all_samples["halueval_summarization"] = load_halueval("summarization", max_samples)

    if "halueval_dialogue" in datasets_to_test:
        all_samples["halueval_dialogue"] = load_halueval("dialogue", max_samples)

    # Initialize detectors
    detectors = []

    # φ-Coherence with different thresholds
    if optimize_thresholds:
        print("\n[*] Finding optimal threshold for φ-Coherence...")
        test_samples = list(all_samples.values())[0][:200]  # Use first 200 for tuning
        thresholds = [0.40, 0.45, 0.50, 0.55, 0.60, 0.65, 0.70]
        best_t, best_f1 = find_optimal_threshold(PhiCoherenceDetector, test_samples, thresholds)
        print(f"[*] Optimal threshold: {best_t} (F1={best_f1:.4f})")
        detectors.append(PhiCoherenceDetector(threshold=best_t))
    else:
        detectors.append(PhiCoherenceDetector(threshold=0.55))

    # Also test fixed thresholds for comparison
    detectors.append(PhiCoherenceDetector(threshold=0.50))
    detectors.append(PhiCoherenceDetector(threshold=0.60))

    # HHEM
    if include_hhem:
        try:
            detectors.append(HHEMDetector(threshold=0.5))
        except Exception as e:
            print(f"[!] Could not load HHEM: {e}")

    # Baselines
    detectors.append(LengthBaselineDetector(threshold=100))
    detectors.append(RandomBaselineDetector())

    # Run evaluation
    all_results = []

    for dataset_name, samples in all_samples.items():
        print(f"\n[*] Evaluating on {dataset_name} ({len(samples)} samples)...")

        for detector in detectors:
            try:
                result = evaluate_detector(detector, samples, verbose=True)
                result.dataset = dataset_name
                all_results.append(result)
            except Exception as e:
                print(f"[!] Error with {detector.name}: {e}")

    # Print results
    print_results_table(all_results)

    # Summary by method (averaged across datasets)
    print("\n" + "-" * 70)
    print("  SUMMARY BY METHOD (averaged across datasets)")
    print("-" * 70)

    method_scores = defaultdict(list)
    for r in all_results:
        method_scores[r.method].append(r.f1)

    for method, scores in sorted(method_scores.items(), key=lambda x: sum(x[1])/len(x[1]), reverse=True):
        avg_f1 = sum(scores) / len(scores)
        print(f"  {method:<35} Avg F1: {avg_f1:.4f}")

    print("-" * 70)

    # Save results
    results_dict = {
        "benchmark": "phi-coherence-comparison",
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "max_samples": max_samples,
        "constants": {"phi": PHI, "alpha": ALPHA},
        "results": [asdict(r) for r in all_results],
    }

    with open("benchmark_comparison_results.json", "w") as f:
        json.dump(results_dict, f, indent=2)

    print("\n[*] Results saved to benchmark_comparison_results.json")

    return all_results


def main():
    parser = argparse.ArgumentParser(description="φ-Coherence Benchmark Comparison")
    parser.add_argument("--max-samples", type=int, default=500, help="Max samples per dataset")
    parser.add_argument("--no-hhem", action="store_true", help="Skip HHEM (faster)")
    parser.add_argument("--quick", action="store_true", help="Quick test with 100 samples")
    parser.add_argument("--datasets", nargs="+", default=["truthfulqa", "halueval_qa"],
                        help="Datasets to test")

    args = parser.parse_args()

    if args.quick:
        args.max_samples = 100

    # Install dependencies
    install_dependencies()

    # Run benchmark
    run_benchmark(
        max_samples=args.max_samples,
        include_hhem=not args.no_hhem,
        datasets_to_test=args.datasets,
    )


if __name__ == "__main__":
    main()