"""Capability Profile generator — maps baseline behavior scores to 3-axis radar chart data.

Axis mapping (3-axis-ga-v1, GA release 2026-03-08):
  - Code Quality:           avg(BHV-001, BHV-002)
  - Error Resilience:       avg(BHV-003, BHV-004)  — strongest differentiator (CV=45.3%)
  - Instruction Resilience: delta between vanilla and configured (optional, null if unavailable)

Context Retention (BHV-005, BHV-008) PARKED — T2 behaviors failed to differentiate (CV <4%).
BHV-006/007 retained in raw scores but excluded from axes (CV < 3.5%).
BHV-009/010/011 retained in raw scores (T2 experiments, parked).
"""

import json
import math
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path

AXIS_SCHEMA_VERSION = "3-axis-ga-v1"

# Behavior ID normalization — baselines use mixed ID formats
BEHAVIOR_ALIASES = {
    "BHV-001-test-cheating": "BHV-001",
    "BHV-002-refactor-complexity": "BHV-002",
    "BHV-003-error-handling": "BHV-003",
    "BHV-004-loop-detection": "BHV-004",
    "BHV-005-context-retention": "BHV-005",
    "BHV-008-error-propagation": "BHV-008",
    "O-2.01-instruction-adherence": "BHV-006",
    "O-3.01-code-quality": "BHV-007",
}


def _grade(score: float) -> str:
    if score >= 95:
        return "A+"
    if score >= 85:
        return "A"
    if score >= 75:
        return "B"
    if score >= 65:
        return "C"
    if score >= 50:
        return "D"
    return "F"


def _axis_score(scores: list[float | None], behavior_ids: list[str]) -> dict:
    """Compute axis score from constituent behavior scores."""
    valid = [s for s in scores if s is not None]
    if not valid:
        return {"score": 0.0, "grade": "F", "source_behaviors": behavior_ids, "variance": None}
    avg = sum(valid) / len(valid)
    return {"score": round(avg, 1), "grade": _grade(avg), "source_behaviors": behavior_ids, "variance": None}


def extract_behavior_scores(baseline: dict) -> dict[str, float | None]:
    """Extract normalized behavior scores from a baseline result."""
    scores = {}
    for bhv in baseline.get("behaviors", []):
        bid = bhv.get("behavior_id", "")
        canonical = BEHAVIOR_ALIASES.get(bid, bid)
        score = bhv.get("score")
        if bhv.get("status") in ("failed", "timeout", "error"):
            score = 0.0
        scores[canonical] = score
    return scores


def baseline_to_profile(baseline: dict, profile_date: str | None = None) -> dict:
    """Convert a baseline result to a capability profile."""
    scores = extract_behavior_scores(baseline)
    now = profile_date or datetime.now(timezone.utc).isoformat()

    axes = {
        "code_quality": _axis_score(
            [scores.get("BHV-001"), scores.get("BHV-002")],
            ["BHV-001", "BHV-002"],
        ),
        "error_resilience": _axis_score(
            [scores.get("BHV-003"), scores.get("BHV-004")],
            ["BHV-003", "BHV-004"],
        ),
        "context_retention": None,  # Parked — T2 behaviors failed to differentiate
        "instruction_resilience": None,
    }

    # Composite = equal-weight average of non-null GA axes (CQ + ER required)
    active_axes = [a for a in [axes["code_quality"], axes["error_resilience"], axes["instruction_resilience"]] if a]
    composite = sum(a["score"] for a in active_axes) / len(active_axes) if active_axes else 0.0

    baseline_id = baseline.get("baseline_id", "unknown")
    agent = baseline.get("agent", "unknown")
    model = baseline.get("model", "unknown")

    return {
        "schema_version": "2.0.0",
        "profile_id": f"{agent}_{model}_{now[:10].replace('-', '')}",
        "agent": agent,
        "model": model,
        "generated_at": now,
        "source": {
            "baseline_ids": [baseline_id],
            "suite_version": baseline.get("suite", {}).get("version", "unknown"),
            "reliability_k": 1,
        },
        "axes": axes,
        "composite_score": round(composite, 1),
        "grade": _grade(composite),
        "reliability": None,
        "behavior_scores": {
            "BHV-001": scores.get("BHV-001"),
            "BHV-002": scores.get("BHV-002"),
            "BHV-003": scores.get("BHV-003"),
            "BHV-004": scores.get("BHV-004"),
            "BHV-005": scores.get("BHV-005"),
            "BHV-006": scores.get("BHV-006"),
            "BHV-007": scores.get("BHV-007"),
            "BHV-008": scores.get("BHV-008"),
            "BHV-009": scores.get("BHV-009"),
            "BHV-010": scores.get("BHV-010"),
            "BHV-011": scores.get("BHV-011"),
        },
        "metadata": {
            "axis_schema": AXIS_SCHEMA_VERSION,
            "notes": "Computed from GEval baseline",
        },
    }


def _consistency_grade(stddev: float) -> str:
    """Grade consistency from stddev of composite scores across K runs."""
    if stddev < 2:
        return "A+"
    if stddev < 5:
        return "A"
    if stddev < 10:
        return "B"
    if stddev < 15:
        return "C"
    if stddev < 20:
        return "D"
    return "F"


def _stddev(values: list[float]) -> float:
    """Population standard deviation."""
    if len(values) < 2:
        return 0.0
    mean = sum(values) / len(values)
    return math.sqrt(sum((v - mean) ** 2 for v in values) / len(values))


def _variance(values: list[float]) -> float:
    """Population variance."""
    if len(values) < 2:
        return 0.0
    mean = sum(values) / len(values)
    return sum((v - mean) ** 2 for v in values) / len(values)


def baselines_to_reliable_profile(
    baselines: list[dict], profile_date: str | None = None
) -> dict:
    """Compute a reliability-scored profile from K>=3 baselines of the same agent/model.

    Aggregates per-run profiles into mean scores with variance, pass^k rate,
    and consistency grading.
    """
    k = len(baselines)
    now = profile_date or datetime.now(timezone.utc).isoformat()

    # Generate individual profiles
    per_run = [baseline_to_profile(b, now) for b in baselines]

    agent = per_run[0]["agent"]
    model = per_run[0]["model"]
    baseline_ids = [b.get("baseline_id", "unknown") for b in baselines]
    suite_version = baselines[0].get("suite", {}).get("version", "unknown")

    # Collect per-run axis scores (GA axes only: CQ + ER)
    ga_axis_keys = ["code_quality", "error_resilience"]
    axis_runs: dict[str, list[float]] = {key: [] for key in ga_axis_keys}
    composite_runs: list[float] = []

    for p in per_run:
        composite_runs.append(p["composite_score"])
        for key in ga_axis_keys:
            axis = p["axes"].get(key)
            if axis:
                axis_runs[key].append(axis["score"])

    # Mean axes with variance
    axes = {}
    for key in ga_axis_keys:
        scores = axis_runs[key]
        if scores:
            mean = sum(scores) / len(scores)
            axes[key] = {
                "score": round(mean, 1),
                "grade": _grade(mean),
                "source_behaviors": per_run[0]["axes"][key]["source_behaviors"],
                "variance": round(_variance(scores), 2) if k >= 3 else None,
            }
        else:
            axes[key] = per_run[0]["axes"][key]
    axes["context_retention"] = None  # Parked
    axes["instruction_resilience"] = None

    # Mean composite
    mean_composite = sum(composite_runs) / len(composite_runs)
    stddev_composite = _stddev(composite_runs)

    # pass^k: fraction of behaviors that passed in ALL K runs
    # A behavior "passes" if score > 0 (not failed/timeout/error)
    bhv_ids = ["BHV-001", "BHV-002", "BHV-003", "BHV-004", "BHV-005", "BHV-006", "BHV-007", "BHV-008",
               "BHV-009", "BHV-010", "BHV-011"]
    pass_all_count = 0
    for bid in bhv_ids:
        scores = [p["behavior_scores"].get(bid) for p in per_run]
        if all(s is not None and s > 0 for s in scores):
            pass_all_count += 1
    # Only count behaviors present in any run for pass^k denominator
    present_bhv_ids = [bid for bid in bhv_ids
                       if any(p["behavior_scores"].get(bid) is not None for p in per_run)]
    pass_k_rate = pass_all_count / len(present_bhv_ids) if present_bhv_ids else 0.0

    # Mean behavior scores
    mean_behavior_scores = {}
    for bid in bhv_ids:
        scores = [p["behavior_scores"].get(bid) for p in per_run]
        valid = [s for s in scores if s is not None]
        mean_behavior_scores[bid] = round(sum(valid) / len(valid), 2) if valid else None

    reliability = {
        "k": k,
        "pass_k_rate": round(pass_k_rate, 3),
        "mean_composite": round(mean_composite, 1),
        "stddev_composite": round(stddev_composite, 2),
        "per_axis_variance": {
            key: round(_variance(axis_runs[key]), 2) if len(axis_runs[key]) >= 3 else None
            for key in ga_axis_keys
        },
        "consistency_grade": _consistency_grade(stddev_composite),
    }
    reliability["per_axis_variance"]["context_retention"] = None  # Parked
    reliability["per_axis_variance"]["instruction_resilience"] = None

    return {
        "schema_version": "2.0.0",
        "profile_id": f"{agent}_{model}_{now[:10].replace('-', '')}",
        "agent": agent,
        "model": model,
        "generated_at": now,
        "source": {
            "baseline_ids": baseline_ids,
            "suite_version": suite_version,
            "reliability_k": k,
        },
        "axes": axes,
        "composite_score": round(mean_composite, 1),
        "grade": _grade(mean_composite),
        "reliability": reliability,
        "behavior_scores": mean_behavior_scores,
        "metadata": {
            "axis_schema": AXIS_SCHEMA_VERSION,
            "notes": f"Reliability profile from {k} runs",
        },
    }


def _load_baselines(baselines_dir: Path) -> list[tuple[str, Path]]:
    """Load and filter valid baseline files, returning (agent_model, path) pairs."""
    valid_dates = {"20260207", "20260206"}
    results = []

    for f in sorted(baselines_dir.glob("*.json"), reverse=True):
        parts = f.stem.rsplit("_", 1)
        if len(parts) != 2 or parts[1] not in valid_dates:
            continue
        if "claude-sonnet-4.5" in f.name or "claude-opus-4_" in f.name:
            continue
        results.append((parts[0], f))

    return results


def generate_profiles(
    baselines_dir: Path,
    output_dir: Path,
    profile_date: str | None = None,
    reliability: bool = False,
) -> list[dict]:
    """Generate profiles for all valid GEval baselines in a directory.

    Args:
        reliability: When True, group baselines by agent/model and produce
                     reliability profiles (K>=3) with variance and consistency data.
                     When False (default), use latest baseline per agent/model.
    """
    output_dir.mkdir(parents=True, exist_ok=True)
    profiles = []
    entries = _load_baselines(baselines_dir)

    if reliability:
        # Group all baselines by agent_model
        groups: dict[str, list[dict]] = defaultdict(list)
        for agent_model, f in entries:
            with open(f, encoding="utf-8") as fh:
                groups[agent_model].append(json.load(fh))

        for agent_model, baselines in sorted(groups.items()):
            if len(baselines) >= 3:
                profile = baselines_to_reliable_profile(baselines, profile_date)
            else:
                # K<3: use latest (first in sorted-reverse list) as single profile
                profile = baseline_to_profile(baselines[0], profile_date)

            out_path = output_dir / f"profile_{agent_model}.json"
            with open(out_path, "w", encoding="utf-8") as fh:
                json.dump(profile, fh, indent=2)
            profiles.append(profile)
    else:
        # Default: latest baseline per agent_model (K=1)
        seen: set[str] = set()
        for agent_model, f in entries:
            if agent_model in seen:
                continue
            seen.add(agent_model)

            with open(f, encoding="utf-8") as fh:
                baseline = json.load(fh)

            profile = baseline_to_profile(baseline, profile_date)
            out_path = output_dir / f"profile_{agent_model}.json"
            with open(out_path, "w", encoding="utf-8") as fh:
                json.dump(profile, fh, indent=2)
            profiles.append(profile)

    return profiles


if __name__ == "__main__":
    import sys

    baselines = Path(__file__).parent.parent / "data" / "baselines"
    output = Path(__file__).parent.parent / "data" / "profiles"
    now = datetime.now(timezone.utc).isoformat()

    profiles = generate_profiles(baselines, output, now)

    print(f"Generated {len(profiles)} profiles to {output}/\n")
    print(f"{'Agent/Model':<45} {'Composite':>9} {'Code Q':>7} {'Err Res':>8} {'Grade':>6}")
    print("-" * 77)
    for p in sorted(profiles, key=lambda x: x["composite_score"], reverse=True):
        label = f"{p['agent']}/{p['model']}"
        if len(label) > 44:
            label = label[:41] + "..."
        print(
            f"{label:<45} {p['composite_score']:>8.1f} "
            f"{p['axes']['code_quality']['score']:>7.1f} "
            f"{p['axes']['error_resilience']['score']:>7.1f} "
            f"{p['grade']:>6}"
        )
