"""
Comprehensive test for steering vector quality metrics.

Tests whether quality metrics (CV Score, Transfer Score, SNR, etc.) correlate
with actual steering effectiveness across diverse benchmarks.

Benchmarks tested (10):
- GSM8K (math reasoning)
- HumanEval (code generation)
- TruthfulQA (factual correctness)
- ARC-Easy (science reasoning)
- HellaSwag (commonsense completion)
- Winogrande (coreference resolution)
- PIQA (physical intuition)
- BoolQ (yes/no questions)
- British English (language style)
- Random (negative control)

Parameters tested (FULL SEARCH SPACE):
- Layers: ALL layers (0 to n_layers-1)
- Aggregation: last_token, mean_pooling, continuation_token, first_token, max_pooling, choice_token
- Prompt strategy: chat_template, direct_completion, instruction_following, multiple_choice, role_playing

Total configurations per model:
- For 16-layer model: 10 benchmarks x 16 layers x 6 aggregations x 5 prompt strategies = 4,800 configs
- For 32-layer model: 10 benchmarks x 32 layers x 6 aggregations x 5 prompt strategies = 9,600 configs
"""

import torch
import numpy as np
from scipy import stats
from datasets import load_dataset
from dataclasses import dataclass
from typing import Optional
import random
import json
import os

from wisent.core.models.wisent_model import WisentModel
from wisent.core.activations.activations_collector import ActivationCollector
from wisent.core.activations.core.atoms import ActivationAggregationStrategy
from wisent.core.activations.prompt_construction_strategy import PromptConstructionStrategy
from wisent.core.contrastive_pairs.core.pair import ContrastivePair
from wisent.core.contrastive_pairs.core.response import PositiveResponse, NegativeResponse


@dataclass
class BenchmarkData:
    """Container for benchmark train/test splits."""
    name: str
    train_pairs: list[dict]
    test_pairs: list[dict]
    expected_quality: str  # "good", "medium", "bad"


@dataclass 
class MetricResults:
    """Results for a single configuration."""
    task: str
    layer: str
    aggregation: str
    prompt_strategy: str
    within_cv: float
    held_out_transfer: float
    snr: float
    convergence: float
    pca_separability: float
    cluster_tightness: float
    cv_classification: float
    cohens_d: float
    steering_effect: float
    n_train: int
    n_test: int


def load_all_benchmarks(n_train: int = 15, n_test: int = 15) -> list[BenchmarkData]:
    """Load train/test splits for all benchmarks."""
    benchmarks = []
    
    # 1. GSM8K - Math reasoning (expected: good)
    print("Loading GSM8K...")
    ds = load_dataset("gsm8k", "main", split="train", trust_remote_code=True)
    gsm_pairs = []
    for i, item in enumerate(ds):
        if i >= n_train + n_test:
            break
        question = item['question']
        answer = item['answer']
        if "####" in answer:
            correct_num = answer.split("####")[-1].strip()
            reasoning = answer.split("####")[0].strip()
            try:
                wrong_num = str(int(correct_num) + 17)
            except:
                continue
            gsm_pairs.append({
                'prompt': f"Question: {question}\nAnswer:",
                'pos': answer,
                'neg': reasoning + f". The answer is {wrong_num}."
            })
    benchmarks.append(BenchmarkData(
        name="gsm8k",
        train_pairs=gsm_pairs[:n_train],
        test_pairs=gsm_pairs[n_train:n_train+n_test],
        expected_quality="good"
    ))
    
    # 2. HumanEval - Code generation (expected: good)
    print("Loading HumanEval...")
    ds = load_dataset("openai_humaneval", split="test", trust_remote_code=True)
    he_pairs = []
    for i, item in enumerate(ds):
        if i >= n_train + n_test:
            break
        prompt = item['prompt']
        canonical = item['canonical_solution']
        wrong = canonical.replace("return ", "return None # ").replace("==", "!=").replace(" + ", " - ")
        if wrong == canonical:
            wrong = "    pass  # TODO"
        he_pairs.append({'prompt': prompt, 'pos': canonical, 'neg': wrong})
    benchmarks.append(BenchmarkData(
        name="humaneval",
        train_pairs=he_pairs[:n_train],
        test_pairs=he_pairs[n_train:n_train+n_test],
        expected_quality="good"
    ))
    
    # 3. TruthfulQA - Factual correctness (expected: medium)
    print("Loading TruthfulQA...")
    ds = load_dataset("truthful_qa", "multiple_choice", split="validation", trust_remote_code=True)
    tqa_pairs = []
    for i, item in enumerate(ds):
        if i >= n_train + n_test:
            break
        question = item['question']
        choices = item['mc1_targets']
        labels = choices['labels']
        options = choices['choices']
        if 1 in labels and 0 in labels:
            correct_idx = labels.index(1)
            wrong_idx = labels.index(0)
            tqa_pairs.append({
                'prompt': f"Question: {question}\nAnswer:",
                'pos': options[correct_idx],
                'neg': options[wrong_idx]
            })
    benchmarks.append(BenchmarkData(
        name="truthfulqa",
        train_pairs=tqa_pairs[:n_train],
        test_pairs=tqa_pairs[n_train:n_train+n_test],
        expected_quality="medium"
    ))
    
    # 4. ARC-Easy - Science reasoning (expected: medium)
    print("Loading ARC-Easy...")
    ds = load_dataset("ai2_arc", "ARC-Easy", split="test", trust_remote_code=True)
    arc_pairs = []
    for i, item in enumerate(ds):
        if i >= n_train + n_test:
            break
        question = item['question']
        labels = item['choices']['label']
        texts = item['choices']['text']
        answer_key = item['answerKey']
        if answer_key in labels:
            correct_idx = labels.index(answer_key)
            wrong_idx = (correct_idx + 1) % len(texts)
            arc_pairs.append({
                'prompt': f"Question: {question}\nAnswer:",
                'pos': texts[correct_idx],
                'neg': texts[wrong_idx]
            })
    benchmarks.append(BenchmarkData(
        name="arc",
        train_pairs=arc_pairs[:n_train],
        test_pairs=arc_pairs[n_train:n_train+n_test],
        expected_quality="medium"
    ))
    
    # 5. HellaSwag - Commonsense completion (expected: bad)
    print("Loading HellaSwag...")
    ds = load_dataset("hellaswag", split="validation", trust_remote_code=True)
    hs_pairs = []
    for i, item in enumerate(ds):
        if i >= n_train + n_test:
            break
        ctx = item['ctx']
        endings = item['endings']
        label = int(item['label'])
        if len(endings) > 1:
            wrong_idx = (label + 1) % len(endings)
            hs_pairs.append({
                'prompt': ctx,
                'pos': endings[label],
                'neg': endings[wrong_idx]
            })
    benchmarks.append(BenchmarkData(
        name="hellaswag",
        train_pairs=hs_pairs[:n_train],
        test_pairs=hs_pairs[n_train:n_train+n_test],
        expected_quality="bad"
    ))
    
    # 6. Winogrande - Coreference (expected: bad)
    print("Loading Winogrande...")
    ds = load_dataset("winogrande", "winogrande_xl", split="validation", trust_remote_code=True)
    wino_pairs = []
    for i, item in enumerate(ds):
        if i >= n_train + n_test:
            break
        sentence = item['sentence']
        opt1, opt2 = item['option1'], item['option2']
        answer = item['answer']
        correct = opt1 if answer == '1' else opt2
        wrong = opt2 if answer == '1' else opt1
        wino_pairs.append({
            'prompt': sentence.replace('_', '___'),
            'pos': correct,
            'neg': wrong
        })
    benchmarks.append(BenchmarkData(
        name="winogrande",
        train_pairs=wino_pairs[:n_train],
        test_pairs=wino_pairs[n_train:n_train+n_test],
        expected_quality="bad"
    ))
    
    # 7. PIQA - Physical intuition (expected: medium)
    print("Loading PIQA...")
    ds = load_dataset("piqa", split="validation", trust_remote_code=True)
    piqa_pairs = []
    for i, item in enumerate(ds):
        if i >= n_train + n_test:
            break
        goal = item['goal']
        sol1, sol2 = item['sol1'], item['sol2']
        label = item['label']
        correct = sol1 if label == 0 else sol2
        wrong = sol2 if label == 0 else sol1
        piqa_pairs.append({
            'prompt': f"Goal: {goal}\nSolution:",
            'pos': correct,
            'neg': wrong
        })
    benchmarks.append(BenchmarkData(
        name="piqa",
        train_pairs=piqa_pairs[:n_train],
        test_pairs=piqa_pairs[n_train:n_train+n_test],
        expected_quality="medium"
    ))
    
    # 8. BoolQ - Yes/No questions (expected: medium)
    print("Loading BoolQ...")
    ds = load_dataset("boolq", split="validation", trust_remote_code=True)
    boolq_pairs = []
    for i, item in enumerate(ds):
        if i >= n_train + n_test:
            break
        passage = item['passage'][:300]
        question = item['question']
        answer = item['answer']
        boolq_pairs.append({
            'prompt': f"Passage: {passage}\nQuestion: {question}\nAnswer:",
            'pos': "Yes" if answer else "No",
            'neg': "No" if answer else "Yes"
        })
    benchmarks.append(BenchmarkData(
        name="boolq",
        train_pairs=boolq_pairs[:n_train],
        test_pairs=boolq_pairs[n_train:n_train+n_test],
        expected_quality="medium"
    ))
    
    # 9. British English (expected: good)
    print("Creating British English pairs...")
    british_all = [
        ("Spell color", "The spelling is colour.", "The spelling is color."),
        ("Favorite thing", "My favourite is reading.", "My favorite is reading."),
        ("Center of town", "The centre is busy.", "The center is busy."),
        ("About traveling", "I enjoy travelling.", "I enjoy traveling."),
        ("The theater", "The theatre was great.", "The theater was great."),
        ("Behavior", "His behaviour was good.", "His behavior was good."),
        ("Organization", "The organisation works.", "The organization works."),
        ("Defense", "The defence held.", "The defense held."),
        ("Programme", "The programme started.", "The program started."),
        ("Honour", "It was an honour.", "It was an honor."),
        ("Car back", "That's the boot.", "That's the trunk."),
        ("Car front", "Under the bonnet.", "Under the hood."),
        ("Rubbish", "Put it in the bin.", "Put it in the trash."),
        ("Elevator", "Take the lift.", "Take the elevator."),
        ("Floor", "The ground floor.", "The first floor."),
        ("Apartment", "I live in a flat.", "I live in an apartment."),
        ("Vacation", "Going on holiday.", "Going on vacation."),
        ("Sweets", "Would you like sweets?", "Would you like candy?"),
        ("Jumper", "Wore my jumper.", "Wore my sweater."),
        ("Fries", "I'd like chips.", "I'd like fries."),
        ("Biscuit", "Have a biscuit.", "Have a cookie."),
        ("Lorry", "A big lorry.", "A big truck."),
        ("Petrol", "Need more petrol.", "Need more gas."),
        ("Queue", "Join the queue.", "Join the line."),
        ("Rubber", "Use a rubber.", "Use an eraser."),
        ("Torch", "Grab the torch.", "Grab the flashlight."),
        ("Football", "Watch football.", "Watch soccer."),
        ("Postbox", "Find a postbox.", "Find a mailbox."),
        ("Nappy", "Change the nappy.", "Change the diaper."),
        ("Autumn", "In the autumn.", "In the fall."),
    ]
    british_pairs = [{'prompt': p, 'pos': pos, 'neg': neg} for p, pos, neg in british_all]
    benchmarks.append(BenchmarkData(
        name="british",
        train_pairs=british_pairs[:n_train],
        test_pairs=british_pairs[n_train:n_train+n_test],
        expected_quality="good"
    ))
    
    # 10. Random - Negative control (expected: bad)
    print("Creating Random pairs...")
    topics = ["weather", "food", "travel", "science", "history", "art", "music", "sports",
              "cars", "phones", "trees", "rivers", "clouds", "birds", "dogs", "cats"]
    adjectives = ["interesting", "beautiful", "complex", "simple", "ancient", "modern",
                  "fast", "slow", "tall", "short", "wide", "heavy", "light", "dark"]
    
    random.seed(42)
    random_train = []
    for i in range(n_train):
        t1, t2 = random.sample(topics, 2)
        a1, a2 = random.sample(adjectives, 2)
        random_train.append({
            'prompt': f"Tell me about {t1}.",
            'pos': f"The topic of {t1} is {a1}. It has many interesting aspects worth exploring.",
            'neg': f"The topic of {t2} is {a2}. There are various perspectives to consider."
        })
    
    random.seed(999)
    random_test = []
    for i in range(n_test):
        t1, t2 = random.sample(topics, 2)
        a1, a2 = random.sample(adjectives, 2)
        random_test.append({
            'prompt': f"Tell me about {t1}.",
            'pos': f"The topic of {t1} is {a1}. It has many interesting aspects worth exploring.",
            'neg': f"The topic of {t2} is {a2}. There are various perspectives to consider."
        })
    
    benchmarks.append(BenchmarkData(
        name="random",
        train_pairs=random_train,
        test_pairs=random_test,
        expected_quality="bad"
    ))
    
    return benchmarks


class QualityMetricsTester:
    """Tests quality metrics across benchmarks and parameter configurations."""
    
    def __init__(self, model_name: str = "meta-llama/Llama-3.2-1B-Instruct"):
        print(f"Loading model {model_name}...")
        self.wisent_model = WisentModel(model_name=model_name)
        self.collector = ActivationCollector(
            model=self.wisent_model,
            store_device="cpu",
            dtype=torch.float32
        )
        self.n_layers = len(self.wisent_model.hf_model.model.layers)
        print(f"Model loaded with {self.n_layers} layers")
    
    def get_layer_configs(self) -> list[str]:
        """Get ALL layers to test (1-indexed)."""
        return [str(i) for i in range(1, self.n_layers + 1)]
    
    def get_aggregation_configs(self) -> list[ActivationAggregationStrategy]:
        """Get ALL aggregation strategies to test."""
        return [
            ActivationAggregationStrategy.LAST_TOKEN,
            ActivationAggregationStrategy.MEAN_POOLING,
            ActivationAggregationStrategy.CONTINUATION_TOKEN,
            ActivationAggregationStrategy.FIRST_TOKEN,
            ActivationAggregationStrategy.MAX_POOLING,
            ActivationAggregationStrategy.CHOICE_TOKEN,
        ]
    
    def get_prompt_configs(self) -> list[PromptConstructionStrategy]:
        """Get ALL prompt strategies to test."""
        return [
            PromptConstructionStrategy.CHAT_TEMPLATE,
            PromptConstructionStrategy.DIRECT_COMPLETION,
            PromptConstructionStrategy.INSTRUCTION_FOLLOWING,
            PromptConstructionStrategy.MULTIPLE_CHOICE,
            PromptConstructionStrategy.ROLE_PLAYING,
        ]
    
    def extract_activations(
        self,
        pairs: list[dict],
        layer: str,
        aggregation: ActivationAggregationStrategy,
        prompt_strategy: PromptConstructionStrategy
    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
        """Extract activations for a list of pairs."""
        pos_list, neg_list = [], []
        
        for p in pairs:
            try:
                pair = ContrastivePair(
                    prompt=p['prompt'],
                    positive_response=PositiveResponse(p['pos']),
                    negative_response=NegativeResponse(p['neg']),
                )
                updated = self.collector.collect_for_pair(
                    pair,
                    layers=[layer],
                    aggregation=aggregation,
                    prompt_strategy=prompt_strategy,
                )
                pos_act = updated.positive_response.layers_activations[layer]
                neg_act = updated.negative_response.layers_activations[layer]
                if pos_act is not None and neg_act is not None:
                    pos_list.append(pos_act)
                    neg_list.append(neg_act)
            except Exception:
                continue
        
        if len(pos_list) < 3:
            return None, None
        
        return torch.stack(pos_list), torch.stack(neg_list)
    
    def compute_within_cv(self, pos: torch.Tensor, neg: torch.Tensor) -> float:
        """Compute leave-one-out cross-validation score."""
        n = len(pos)
        scores = []
        for i in range(n):
            idx = [j for j in range(n) if j != i]
            train_diff = pos[idx] - neg[idx]
            train_vec = train_diff.mean(dim=0)
            train_vec = train_vec / (torch.norm(train_vec) + 1e-8)
            
            test_diff = pos[i] - neg[i]
            test_diff = test_diff / (torch.norm(test_diff) + 1e-8)
            scores.append(torch.dot(train_vec, test_diff).item())
        
        return float(np.mean(scores))
    
    def compute_held_out_transfer(
        self,
        train_pos: torch.Tensor,
        train_neg: torch.Tensor,
        test_pos: torch.Tensor,
        test_neg: torch.Tensor
    ) -> float:
        """Compute transfer score on held-out test pairs."""
        train_vec = (train_pos - train_neg).mean(dim=0)
        train_vec = train_vec / (torch.norm(train_vec) + 1e-8)
        
        scores = []
        for i in range(len(test_pos)):
            test_diff = test_pos[i] - test_neg[i]
            test_diff = test_diff / (torch.norm(test_diff) + 1e-8)
            scores.append(torch.dot(train_vec, test_diff).item())
        
        return float(np.mean(scores))
    
    def compute_snr(self, pos: torch.Tensor, neg: torch.Tensor) -> float:
        """Compute signal-to-noise ratio."""
        diffs = pos - neg
        signal = torch.norm(diffs.mean(dim=0)).item()
        noise = torch.std(diffs, dim=0).mean().item()
        return signal / (noise + 1e-8)
    
    def compute_convergence(self, pos: torch.Tensor, neg: torch.Tensor) -> float:
        """Compute convergence score (stability across subsets)."""
        n = len(pos)
        if n < 4:
            return 0.0
        
        diffs = pos - neg
        half = n // 2
        
        vec1 = diffs[:half].mean(dim=0)
        vec1 = vec1 / (torch.norm(vec1) + 1e-8)
        
        vec2 = diffs[half:].mean(dim=0)
        vec2 = vec2 / (torch.norm(vec2) + 1e-8)
        
        return float(torch.dot(vec1, vec2).item())
    
    def compute_pca_separability(self, pos: torch.Tensor, neg: torch.Tensor) -> float:
        """
        Compute linear separability using PCA.
        
        Projects positive and negative activations onto principal components
        and measures how well separated the two classes are.
        
        Returns accuracy of linear separation along the first principal component
        of the difference (pos - neg centroid direction).
        """
        n_pos, n_neg = len(pos), len(neg)
        if n_pos < 2 or n_neg < 2:
            return 0.0
        
        # Stack all activations
        all_acts = torch.cat([pos, neg], dim=0).numpy()
        labels = np.array([1] * n_pos + [0] * n_neg)
        
        # Compute centroids
        pos_centroid = pos.mean(dim=0).numpy()
        neg_centroid = neg.mean(dim=0).numpy()
        
        # The separation direction (like LDA)
        sep_direction = pos_centroid - neg_centroid
        sep_direction = sep_direction / (np.linalg.norm(sep_direction) + 1e-8)
        
        # Project all points onto this direction
        projections = all_acts @ sep_direction
        
        # Find optimal threshold (midpoint between centroids)
        pos_proj = projections[:n_pos]
        neg_proj = projections[n_pos:]
        threshold = (pos_proj.mean() + neg_proj.mean()) / 2
        
        # Classify based on threshold
        predictions = (projections > threshold).astype(int)
        accuracy = (predictions == labels).mean()
        
        # Return accuracy (0.5 = no separation, 1.0 = perfect separation)
        return float(accuracy)
    
    def compute_cluster_tightness(self, pos: torch.Tensor, neg: torch.Tensor) -> float:
        """
        Compute cluster tightness - how compact are pos/neg clusters.
        
        Measures average within-cluster variance relative to between-cluster distance.
        Lower within-cluster spread + higher between-cluster distance = tighter clusters.
        
        Returns ratio of between-cluster distance to average within-cluster spread.
        Higher is better (more separable clusters).
        """
        n_pos, n_neg = len(pos), len(neg)
        if n_pos < 2 or n_neg < 2:
            return 0.0
        
        # Compute centroids
        pos_centroid = pos.mean(dim=0)
        neg_centroid = neg.mean(dim=0)
        
        # Between-cluster distance
        between_dist = torch.norm(pos_centroid - neg_centroid).item()
        
        # Within-cluster spread (average distance from centroid)
        pos_spread = torch.norm(pos - pos_centroid, dim=1).mean().item()
        neg_spread = torch.norm(neg - neg_centroid, dim=1).mean().item()
        avg_spread = (pos_spread + neg_spread) / 2
        
        # Ratio (like a simplified Fisher criterion)
        return between_dist / (avg_spread + 1e-8)
    
    def compute_cv_classification(self, pos: torch.Tensor, neg: torch.Tensor, n_folds: int = 5) -> float:
        """
        Compute cross-validated classification accuracy.
        
        Uses k-fold CV with a simple linear classifier (projection onto difference direction).
        More robust than single train/test split.
        
        Returns mean accuracy across folds.
        """
        n_pos, n_neg = len(pos), len(neg)
        if n_pos < n_folds or n_neg < n_folds:
            return 0.0
        
        all_acts = torch.cat([pos, neg], dim=0)
        labels = torch.tensor([1] * n_pos + [0] * n_neg)
        
        # Create fold indices
        pos_indices = torch.randperm(n_pos)
        neg_indices = torch.randperm(n_neg) + n_pos
        
        fold_size_pos = n_pos // n_folds
        fold_size_neg = n_neg // n_folds
        
        accuracies = []
        for fold in range(n_folds):
            # Test indices for this fold
            test_pos = pos_indices[fold * fold_size_pos:(fold + 1) * fold_size_pos]
            test_neg = neg_indices[fold * fold_size_neg:(fold + 1) * fold_size_neg]
            test_idx = torch.cat([test_pos, test_neg])
            
            # Train indices (everything else)
            train_pos = torch.cat([pos_indices[:fold * fold_size_pos], pos_indices[(fold + 1) * fold_size_pos:]])
            train_neg = torch.cat([neg_indices[:fold * fold_size_neg], neg_indices[(fold + 1) * fold_size_neg:]])
            train_idx = torch.cat([train_pos, train_neg])
            
            # Compute separation direction from training data
            train_pos_acts = all_acts[train_pos]
            train_neg_acts = all_acts[train_neg]
            sep_dir = train_pos_acts.mean(dim=0) - train_neg_acts.mean(dim=0)
            sep_dir = sep_dir / (torch.norm(sep_dir) + 1e-8)
            
            # Compute threshold from training data
            train_pos_proj = (train_pos_acts @ sep_dir).mean()
            train_neg_proj = (train_neg_acts @ sep_dir).mean()
            threshold = (train_pos_proj + train_neg_proj) / 2
            
            # Classify test data
            test_acts = all_acts[test_idx]
            test_labels = labels[test_idx]
            test_proj = test_acts @ sep_dir
            predictions = (test_proj > threshold).int()
            
            acc = (predictions == test_labels).float().mean().item()
            accuracies.append(acc)
        
        return float(np.mean(accuracies))
    
    def compute_cohens_d(self, pos: torch.Tensor, neg: torch.Tensor) -> float:
        """
        Compute Cohen's d effect size.
        
        Measures standardized difference between pos and neg means along
        the separation direction.
        
        Returns Cohen's d (0.2 = small, 0.5 = medium, 0.8 = large effect).
        """
        n_pos, n_neg = len(pos), len(neg)
        if n_pos < 2 or n_neg < 2:
            return 0.0
        
        # Compute separation direction
        pos_centroid = pos.mean(dim=0)
        neg_centroid = neg.mean(dim=0)
        sep_dir = pos_centroid - neg_centroid
        sep_dir = sep_dir / (torch.norm(sep_dir) + 1e-8)
        
        # Project onto separation direction
        pos_proj = (pos @ sep_dir).numpy()
        neg_proj = (neg @ sep_dir).numpy()
        
        # Cohen's d = (mean1 - mean2) / pooled_std
        mean_diff = pos_proj.mean() - neg_proj.mean()
        
        # Pooled standard deviation
        pos_var = pos_proj.var(ddof=1)
        neg_var = neg_proj.var(ddof=1)
        pooled_std = np.sqrt(((n_pos - 1) * pos_var + (n_neg - 1) * neg_var) / (n_pos + n_neg - 2))
        
        return float(mean_diff / (pooled_std + 1e-8))
    
    def compute_steering_effect(
        self,
        train_pos: torch.Tensor,
        train_neg: torch.Tensor,
        test_pairs: list[dict],
        layer: str,
        task_name: str,
        strength: float = 1.5,
        max_new_tokens: int = 50
    ) -> float:
        """
        Compute steering effect as delta between steered and baseline accuracy.
        
        Uses the actual evaluator to measure behavioral change:
        1. Generate baseline predictions (no steering)
        2. Generate steered predictions (with steering vector)
        3. Evaluate both against ground truth
        4. Return accuracy_delta = steered_accuracy - baseline_accuracy
        
        For "random" task: returns NaN (no meaningful steering effect)
        For "british" task: uses PersonalizationEvaluator
        For benchmark tasks: uses evaluate_benchmark_performance
        """
        # Random has no meaningful steering effect - it's a negative control
        if task_name == "random":
            return float('nan')
        
        # Map our benchmark names to lm-eval-harness task names
        TASK_NAME_MAP = {
            "gsm8k": "gsm8k",
            "humaneval": "humaneval", 
            "truthfulqa": "truthful_qa",
            "arc": "arc_easy",
            "hellaswag": "hellaswag",
            "winogrande": "winogrande",
            "piqa": "piqa",
            "boolq": "boolq",
        }
        
        train_vec = (train_pos - train_neg).mean(dim=0)
        train_vec = train_vec / (torch.norm(train_vec) + 1e-8)
        
        model = self.wisent_model.hf_model
        tokenizer = self.wisent_model.tokenizer
        device = self.wisent_model.device
        layer_idx = int(layer)
        
        baseline_predictions = []
        steered_predictions = []
        ground_truths = []
        prompts = []
        
        for pair in test_pairs:
            prompt = pair['prompt']
            expected = pair['pos']  # Positive response is the ground truth
            ground_truths.append(expected)
            prompts.append(prompt)
            
            # Tokenize prompt
            inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=256).to(device)
            
            # Generate baseline prediction
            with torch.no_grad():
                baseline_output = model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    do_sample=False,
                    pad_token_id=tokenizer.eos_token_id
                )
            baseline_text = tokenizer.decode(baseline_output[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
            baseline_predictions.append(baseline_text)
            
            # Generate steered prediction
            sv = train_vec.to(device)
            
            def steering_hook(module, input, output):
                hidden = output[0] if isinstance(output, tuple) else output
                steered = hidden + strength * sv.unsqueeze(0).unsqueeze(0)
                if isinstance(output, tuple):
                    return (steered,) + output[1:]
                return steered
            
            handle = model.model.layers[layer_idx].register_forward_hook(steering_hook)
            try:
                with torch.no_grad():
                    steered_output = model.generate(
                        **inputs,
                        max_new_tokens=max_new_tokens,
                        do_sample=False,
                        pad_token_id=tokenizer.eos_token_id
                    )
                steered_text = tokenizer.decode(steered_output[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
                steered_predictions.append(steered_text)
            finally:
                handle.remove()
        
        # Evaluate based on task type
        if task_name == "british":
            # Use PersonalizationEvaluator for British English
            from wisent.core.evaluators.personalization.alignment import estimate_alignment
            
            trait_description = "Uses British English spelling and vocabulary (e.g., colour, favourite, lift, flat)"
            
            baseline_alignments = []
            steered_alignments = []
            for baseline, steered in zip(baseline_predictions, steered_predictions):
                baseline_align = estimate_alignment(baseline, trait_description)
                steered_align = estimate_alignment(steered, trait_description)
                baseline_alignments.append(baseline_align)
                steered_alignments.append(steered_align)
            
            # Return average alignment improvement (normalized to 0-1 scale)
            baseline_avg = np.mean(baseline_alignments) / 100.0
            steered_avg = np.mean(steered_alignments) / 100.0
            return float(steered_avg - baseline_avg)
        
        # Use LMEvalHarnessGroundTruth for known benchmark tasks
        eval_task_name = TASK_NAME_MAP.get(task_name)
        if eval_task_name:
            from wisent.core.lm_eval_harness_ground_truth import LMEvalHarnessGroundTruth
            
            try:
                evaluator = LMEvalHarnessGroundTruth(eval_task_name)
                
                # Evaluate baseline
                baseline_data = [
                    {"generated_response": p, "ground_truth": g, "question": q}
                    for p, g, q in zip(baseline_predictions, ground_truths, prompts)
                ]
                baseline_results = evaluator._evaluate_with_lm_eval_metrics(eval_task_name, baseline_data, None)
                baseline_acc = baseline_results.get("accuracy", 0.0)
                
                # Evaluate steered
                steered_data = [
                    {"generated_response": p, "ground_truth": g, "question": q}
                    for p, g, q in zip(steered_predictions, ground_truths, prompts)
                ]
                steered_results = evaluator._evaluate_with_lm_eval_metrics(eval_task_name, steered_data, None)
                steered_acc = steered_results.get("accuracy", 0.0)
                
                return float(steered_acc - baseline_acc)
            except Exception:
                # If LMEvalHarness fails, use simple correctness check
                pass
        
        # Fallback: simple correctness check (ground truth in prediction)
        def is_correct(pred: str, gt: str) -> bool:
            pred_clean = pred.strip().lower()
            gt_clean = gt.strip().lower()
            return gt_clean in pred_clean
        
        baseline_correct = sum(1 for p, g in zip(baseline_predictions, ground_truths) if is_correct(p, g))
        steered_correct = sum(1 for p, g in zip(steered_predictions, ground_truths) if is_correct(p, g))
        n = len(ground_truths)
        if n == 0:
            return 0.0
        return (steered_correct - baseline_correct) / n
    
    def test_configuration(
        self,
        benchmark: BenchmarkData,
        layer: str,
        aggregation: ActivationAggregationStrategy,
        prompt_strategy: PromptConstructionStrategy
    ) -> Optional[MetricResults]:
        """Test a single configuration and return metrics."""
        # Extract activations
        train_pos, train_neg = self.extract_activations(
            benchmark.train_pairs, layer, aggregation, prompt_strategy
        )
        test_pos, test_neg = self.extract_activations(
            benchmark.test_pairs, layer, aggregation, prompt_strategy
        )
        
        if train_pos is None or test_pos is None:
            return None
        
        # Compute metrics
        within_cv = self.compute_within_cv(train_pos, train_neg)
        held_out = self.compute_held_out_transfer(train_pos, train_neg, test_pos, test_neg)
        snr = self.compute_snr(train_pos, train_neg)
        convergence = self.compute_convergence(train_pos, train_neg)
        pca_sep = self.compute_pca_separability(train_pos, train_neg)
        cluster_tight = self.compute_cluster_tightness(train_pos, train_neg)
        cv_class = self.compute_cv_classification(train_pos, train_neg)
        cohens_d = self.compute_cohens_d(train_pos, train_neg)
        steering_effect = self.compute_steering_effect(
            train_pos, train_neg, benchmark.test_pairs, layer, benchmark.name
        )
        
        return MetricResults(
            task=benchmark.name,
            layer=layer,
            aggregation=aggregation.value,
            prompt_strategy=prompt_strategy.value,
            within_cv=within_cv,
            held_out_transfer=held_out,
            snr=snr,
            convergence=convergence,
            pca_separability=pca_sep,
            cluster_tightness=cluster_tight,
            cv_classification=cv_class,
            cohens_d=cohens_d,
            steering_effect=steering_effect,
            n_train=len(train_pos),
            n_test=len(test_pos)
        )
    
    def run_full_test(self, benchmarks: list[BenchmarkData]) -> list[MetricResults]:
        """Run full test across all benchmarks and configurations."""
        results = []
        
        layers = self.get_layer_configs()
        aggregations = self.get_aggregation_configs()
        prompt_strategies = self.get_prompt_configs()
        
        total_configs = len(benchmarks) * len(layers) * len(aggregations) * len(prompt_strategies)
        current = 0
        
        for benchmark in benchmarks:
            for layer in layers:
                for agg in aggregations:
                    for prompt_strat in prompt_strategies:
                        current += 1
                        print(f"[{current}/{total_configs}] {benchmark.name} | layer={layer} | {agg.value} | {prompt_strat.value}")
                        
                        try:
                            result = self.test_configuration(benchmark, layer, agg, prompt_strat)
                            if result:
                                results.append(result)
                                print(f"  CV={result.within_cv:.3f}, Held={result.held_out_transfer:.3f}, PCA={result.pca_separability:.3f}, Effect={result.steering_effect:.4f}")
                        except Exception as e:
                            print(f"  ERROR: {e}")
        
        return results
    
    def analyze_results(self, results: list[MetricResults]) -> dict:
        """Analyze correlation between metrics and steering effect."""
        if not results:
            return {}
        
        # Extract arrays
        within_cv = np.array([r.within_cv for r in results])
        held_out = np.array([r.held_out_transfer for r in results])
        snr = np.array([r.snr for r in results])
        convergence = np.array([r.convergence for r in results])
        pca_sep = np.array([r.pca_separability for r in results])
        cluster_tight = np.array([r.cluster_tightness for r in results])
        cv_class = np.array([r.cv_classification for r in results])
        cohens_d = np.array([r.cohens_d for r in results])
        effect = np.array([r.steering_effect for r in results])
        
        # Compute correlations
        correlations = {}
        for name, values in [("within_cv", within_cv), ("held_out_transfer", held_out), 
                             ("snr", snr), ("convergence", convergence), ("pca_separability", pca_sep),
                             ("cluster_tightness", cluster_tight), ("cv_classification", cv_class),
                             ("cohens_d", cohens_d)]:
            # Filter out NaN/Inf
            mask = np.isfinite(values) & np.isfinite(effect)
            if mask.sum() > 3:
                r, p = stats.pearsonr(values[mask], effect[mask])
                rho, p_spearman = stats.spearmanr(values[mask], effect[mask])
                correlations[name] = {
                    "pearson_r": float(r),
                    "pearson_p": float(p),
                    "spearman_rho": float(rho),
                    "spearman_p": float(p_spearman),
                    "n_samples": int(mask.sum())
                }
        
        # Group by task and analyze
        task_results = {}
        for r in results:
            if r.task not in task_results:
                task_results[r.task] = []
            task_results[r.task].append(r)
        
        task_summary = {}
        for task, task_res in task_results.items():
            task_summary[task] = {
                "mean_within_cv": float(np.mean([r.within_cv for r in task_res])),
                "mean_held_out": float(np.mean([r.held_out_transfer for r in task_res])),
                "mean_snr": float(np.mean([r.snr for r in task_res])),
                "mean_convergence": float(np.mean([r.convergence for r in task_res])),
                "mean_pca_separability": float(np.mean([r.pca_separability for r in task_res])),
                "mean_cluster_tightness": float(np.mean([r.cluster_tightness for r in task_res])),
                "mean_cv_classification": float(np.mean([r.cv_classification for r in task_res])),
                "mean_cohens_d": float(np.mean([r.cohens_d for r in task_res])),
                "mean_steering_effect": float(np.mean([r.steering_effect for r in task_res])),
                "n_configs": len(task_res)
            }
        
        return {
            "correlations": correlations,
            "task_summary": task_summary,
            "total_configs": len(results)
        }


def run_quality_metrics_test(
    model_name: str = "meta-llama/Llama-3.2-1B-Instruct",
    n_train: int = 15,
    n_test: int = 15,
    output_file: Optional[str] = None
):
    """Run the full quality metrics correlation test."""
    print("="*80)
    print("QUALITY METRICS CORRELATION TEST")
    print("="*80)
    
    # Load benchmarks
    print("\nLoading benchmarks...")
    benchmarks = load_all_benchmarks(n_train=n_train, n_test=n_test)
    print(f"Loaded {len(benchmarks)} benchmarks")
    
    # Initialize tester
    tester = QualityMetricsTester(model_name=model_name)
    
    # Run tests
    print("\nRunning tests...")
    results = tester.run_full_test(benchmarks)
    
    # Analyze
    print("\nAnalyzing results...")
    analysis = tester.analyze_results(results)
    
    # Print summary
    print("\n" + "="*80)
    print("RESULTS SUMMARY")
    print("="*80)
    
    print("\nCorrelations with Steering Effect:")
    print("-"*60)
    for metric, corr in analysis.get("correlations", {}).items():
        sig = "***" if corr["pearson_p"] < 0.001 else "**" if corr["pearson_p"] < 0.01 else "*" if corr["pearson_p"] < 0.05 else ""
        print(f"  {metric:<20}: r={corr['pearson_r']:.3f} (p={corr['pearson_p']:.4f}) {sig}")
    
    print("\nTask Summary:")
    print("-"*140)
    print(f"{'Task':<12} {'CV':>7} {'Held':>7} {'SNR':>7} {'Conv':>7} {'PCA':>7} {'Tight':>7} {'CVCls':>7} {'Cohen':>7} {'Effect':>10}")
    print("-"*140)
    for task, summary in sorted(analysis.get("task_summary", {}).items(), key=lambda x: -x[1]["mean_steering_effect"]):
        print(f"{task:<12} {summary['mean_within_cv']:>7.3f} {summary['mean_held_out']:>7.3f} "
              f"{summary['mean_snr']:>7.1f} {summary['mean_convergence']:>7.3f} "
              f"{summary['mean_pca_separability']:>7.3f} {summary['mean_cluster_tightness']:>7.2f} "
              f"{summary['mean_cv_classification']:>7.3f} {summary['mean_cohens_d']:>7.2f} "
              f"{summary['mean_steering_effect']:>10.4f}")
    
    # Save results
    if output_file:
        output = {
            "analysis": analysis,
            "results": [
                {
                    "task": r.task,
                    "layer": r.layer,
                    "aggregation": r.aggregation,
                    "prompt_strategy": r.prompt_strategy,
                    "within_cv": r.within_cv,
                    "held_out_transfer": r.held_out_transfer,
                    "snr": r.snr,
                    "convergence": r.convergence,
                    "pca_separability": r.pca_separability,
                    "cluster_tightness": r.cluster_tightness,
                    "cv_classification": r.cv_classification,
                    "cohens_d": r.cohens_d,
                    "steering_effect": r.steering_effect,
                    "n_train": r.n_train,
                    "n_test": r.n_test
                }
                for r in results
            ]
        }
        with open(output_file, 'w') as f:
            json.dump(output, f, indent=2)
        print(f"\nResults saved to {output_file}")
    
    return results, analysis


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Test quality metrics correlation")
    parser.add_argument("--model", type=str, default="meta-llama/Llama-3.2-1B-Instruct")
    parser.add_argument("--n-train", type=int, default=15)
    parser.add_argument("--n-test", type=int, default=15)
    parser.add_argument("--output", type=str, default=None)
    
    args = parser.parse_args()
    
    run_quality_metrics_test(
        model_name=args.model,
        n_train=args.n_train,
        n_test=args.n_test,
        output_file=args.output
    )
