"""Ground truth types and collector.

GPU-intensive phase that produces training data mapping
(geometry_metrics) -> {method: accuracy} for each benchmark.
"""
from __future__ import annotations
import json
import tempfile
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any, Dict, List, Optional

# ── Data types ─────────────────────────────────────────────────


@dataclass(slots=True)
class MethodResult:
    """Accuracy of a single steering method on a single benchmark."""
    method: str
    accuracy: float
    layer: Optional[int] = None
    strength: Optional[float] = None


@dataclass(slots=True)
class BenchmarkGroundTruth:
    """One benchmark record: geometry metrics + per-method accuracy."""
    model: str
    benchmark: str
    metrics: Dict[str, Any]
    method_results: Dict[str, MethodResult]
    best_method: str
    best_accuracy: float

    def to_dict(self) -> dict:
        d = asdict(self)
        d["method_results"] = {
            k: asdict(v) for k, v in self.method_results.items()}
        return d

    @classmethod
    def from_dict(cls, d: dict) -> BenchmarkGroundTruth:
        mr = {k: MethodResult(**v)
              for k, v in d["method_results"].items()}
        return cls(
            model=d["model"], benchmark=d["benchmark"],
            metrics=d["metrics"], method_results=mr,
            best_method=d["best_method"],
            best_accuracy=d["best_accuracy"])


@dataclass(slots=True)
class GroundTruthDataset:
    """Collection of benchmark ground-truth records."""
    records: List[BenchmarkGroundTruth]

    def save(self, path: str | Path) -> None:
        p = Path(path)
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(json.dumps(
            [r.to_dict() for r in self.records], indent=2))

    @classmethod
    def load(cls, path: str | Path) -> GroundTruthDataset:
        raw = json.loads(Path(path).read_text())
        return cls(records=[
            BenchmarkGroundTruth.from_dict(d) for d in raw])

    def __len__(self) -> int:
        return len(self.records)


# ── Collector ──────────────────────────────────────────────────

PIPELINE_METHODS = (
    "CAA", "Hyperplane", "MLP", "PRISM", "PULSE", "TITAN", "Concept Flow")


def _load_repscan_metrics(
    model: str, benchmark: str, repscan_dir: str,
) -> Optional[Dict]:
    model_slug = model.replace("/", "_")
    path = Path(repscan_dir) / f"{model_slug}__{benchmark}.json"
    if not path.exists():
        print(f"  No repscan file: {path}")
        return None
    data = json.loads(path.read_text())
    return data.get("metrics", data)


def _make_default_config(method: str, num_layers: int):
    from wisent.core.cli.optimize_steering import (
        CAAConfig, HyperplaneConfig, MLPConfig,
        PRISMConfig, PULSEConfig, TITANConfig, ConceptFlowConfig,
    )
    mid = num_layers // 2
    upper = [l for l in [mid+1, mid+2, mid+3] if l < num_layers]
    return {
        "CAA": lambda: CAAConfig(method="CAA", layer=mid),
        "Hyperplane": lambda: HyperplaneConfig(
            method="Hyperplane", layer=mid),
        "MLP": lambda: MLPConfig(method="MLP", layer=mid),
        "PRISM": lambda: PRISMConfig(method="PRISM", layer=mid),
        "PULSE": lambda: PULSEConfig(
            method="PULSE", sensor_layer=mid,
            steering_layers=upper),
        "TITAN": lambda: TITANConfig(
            method="TITAN", sensor_layer=mid,
            steering_layers=upper),
        "Concept Flow": lambda: ConceptFlowConfig(
            method="concept_flow", layer=mid),
    }[method]()


def collect_benchmark_ground_truth(
    model: str, benchmark: str,
    repscan_dir: str = "repscan_results",
    limit: int = 100, device: Optional[str] = None,
    methods: Optional[List[str]] = None,
) -> Optional[BenchmarkGroundTruth]:
    """Run all methods on one benchmark, return ground truth."""
    from wisent.core.cli.optimize_steering import run_pipeline
    from transformers import AutoConfig

    metrics = _load_repscan_metrics(model, benchmark, repscan_dir)
    if metrics is None:
        return None
    ac = AutoConfig.from_pretrained(model, trust_remote_code=True)
    num_layers = ac.num_hidden_layers
    if methods is None:
        methods = list(PIPELINE_METHODS)
    method_results: Dict[str, MethodResult] = {}
    best_acc, best_method = -1.0, methods[0]
    for method in methods:
        if method not in PIPELINE_METHODS:
            print(f"  Skipping {method} (no pipeline support)")
            continue
        print(f"  Running {method} on {benchmark}...")
        config = _make_default_config(method, num_layers)
        work_dir = tempfile.mkdtemp(
            prefix=f"gt_{benchmark}_{method}_")
        try:
            result = run_pipeline(
                model=model, task=benchmark, config=config,
                work_dir=work_dir, limit=limit, device=device)
            acc = result.score
        except Exception as e:
            print(f"    {method} failed: {e}")
            acc = 0.0
        layer = getattr(config, 'layer',
                        getattr(config, 'sensor_layer', None))
        strength = getattr(config, 'max_alpha', 1.0)
        method_results[method] = MethodResult(
            method=method, accuracy=acc,
            layer=layer, strength=strength)
        if acc > best_acc:
            best_acc, best_method = acc, method
        print(f"    {method}: {acc:.4f}")
    return BenchmarkGroundTruth(
        model=model, benchmark=benchmark, metrics=metrics,
        method_results=method_results,
        best_method=best_method, best_accuracy=best_acc)


def collect_ground_truth(
    model: str, benchmarks: Optional[List[str]] = None,
    output_path: Optional[str] = None,
    repscan_dir: str = "repscan_results",
    limit: int = 100, device: Optional[str] = None,
    methods: Optional[List[str]] = None,
) -> GroundTruthDataset:
    """Collect ground truth across multiple benchmarks."""
    if benchmarks is None:
        slug = model.replace("/", "_")
        rp = Path(repscan_dir)
        benchmarks = sorted({
            f.stem.replace(f"{slug}__", "")
            for f in rp.glob(f"{slug}__*.json")})
    print(f"Collecting ground truth for {len(benchmarks)} benchmarks")
    records = []
    for i, bench in enumerate(benchmarks):
        print(f"\n[{i+1}/{len(benchmarks)}] {bench}")
        rec = collect_benchmark_ground_truth(
            model, bench, repscan_dir, limit, device, methods)
        if rec is not None:
            records.append(rec)
    dataset = GroundTruthDataset(records=records)
    if output_path:
        dataset.save(output_path)
        print(f"\nSaved {len(records)} records to {output_path}")
    return dataset
