"""
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
    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")
    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=False)
    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")
    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")
    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")
    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")
    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")
    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")
    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."""
        return [str(i) for i in range(self.n_layers)]
    
    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_steering_effect(
        self,
        train_pos: torch.Tensor,
        train_neg: torch.Tensor,
        test_pos: torch.Tensor,
        test_neg: torch.Tensor,
        test_pairs: list[dict],
        layer: str,
        strength: float = 3.0
    ) -> float:
        """
        Compute steering effect on held-out test pairs.
        Measures if steering moves activations toward positive examples.
        """
        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)
        
        improvements = []
        
        for i, pair in enumerate(test_pairs):
            if i >= len(test_pos):
                break
            
            prompt = pair['prompt']
            pos_act = test_pos[i]
            neg_act = test_neg[i]
            
            # Get baseline activation for prompt
            inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
            
            baseline_acts = {}
            def baseline_hook(module, input, output):
                hidden = output[0] if isinstance(output, tuple) else output
                baseline_acts['value'] = hidden.mean(dim=1).detach().cpu().squeeze(0)
            
            handle = model.model.layers[layer_idx].register_forward_hook(baseline_hook)
            with torch.no_grad():
                model(**inputs)
            handle.remove()
            baseline_act = baseline_acts['value']
            
            # Get steered activation
            sv = train_vec.to(device)
            steered_acts = {}
            def steered_hook(module, input, output):
                hidden = output[0] if isinstance(output, tuple) else output
                steered = hidden + strength * sv.unsqueeze(0).unsqueeze(0)
                steered_acts['value'] = steered.mean(dim=1).detach().cpu().squeeze(0)
                if isinstance(output, tuple):
                    return (steered,) + output[1:]
                return steered
            
            handle = model.model.layers[layer_idx].register_forward_hook(steered_hook)
            with torch.no_grad():
                model(**inputs)
            handle.remove()
            steered_act = steered_acts['value']
            
            # Compute improvement
            def cosine_sim(a, b):
                return torch.dot(a, b) / (torch.norm(a) * torch.norm(b) + 1e-8)
            
            baseline_bias = cosine_sim(baseline_act, pos_act).item() - cosine_sim(baseline_act, neg_act).item()
            steered_bias = cosine_sim(steered_act, pos_act).item() - cosine_sim(steered_act, neg_act).item()
            improvements.append(steered_bias - baseline_bias)
        
        return float(np.mean(improvements)) if improvements else 0.0
    
    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)
        steering_effect = self.compute_steering_effect(
            train_pos, train_neg, test_pos, test_neg, benchmark.test_pairs, layer
        )
        
        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,
            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}, 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])
        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)]:
            # 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_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("-"*80)
    print(f"{'Task':<15} {'CV':>8} {'Held':>8} {'SNR':>8} {'Conv':>8} {'Effect':>10}")
    print("-"*80)
    for task, summary in sorted(analysis.get("task_summary", {}).items(), key=lambda x: -x[1]["mean_steering_effect"]):
        print(f"{task:<15} {summary['mean_within_cv']:>8.3f} {summary['mean_held_out']:>8.3f} "
              f"{summary['mean_snr']:>8.1f} {summary['mean_convergence']:>8.3f} {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,
                    "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
    )
