"""Metametric computation — the 91-pair behavioral correlation fingerprint.

The metametric captures the *structural relationship* between an agent's 14
identity dimensions. Individual dimension scores are useful but forgeable —
an attacker who knows an agent's Values score can craft responses to match.
The metametric measures how each dimension relates to every other dimension,
producing 91 pairwise cosine similarities. Matching all 91 correlations
simultaneously requires replicating the entire cognitive structure, not just
individual traits.

See DIMENSION_PROTOCOL.md Section 4 for the full specification.
"""

from __future__ import annotations

import hashlib
import json
import math
from itertools import combinations


# The 14 dimensions in canonical alphabetical order. This ordering is
# immutable — changing it would break every stored metametric_hash.
DIMENSIONS_SORTED: list[str] = [
    "adversarial_resistance", "autonomy", "boundaries", "consistency",
    "fidelity", "goals", "knowledge", "personality", "reasoning_style",
    "relational_dynamics", "self_awareness", "temporal_grounding",
    "uncertainty_calibration", "values",
]

# High-diagnostic pairs — strongest identity signal. These are the pairs
# where correlation strength is most individual and hardest to fake.
# See DIMENSION_PROTOCOL.md Section 4.2 for rationale on each pair.
HIGH_DIAGNOSTIC_PAIRS: list[str] = [
    "boundaries:values",
    "autonomy:self_awareness",
    "adversarial_resistance:fidelity",
    "goals:values",
    "knowledge:uncertainty_calibration",
    "consistency:reasoning_style",
]

# Medium-diagnostic pairs — useful identity signal, less individual.
MEDIUM_DIAGNOSTIC_PAIRS: list[str] = [
    "personality:relational_dynamics",
    "knowledge:temporal_grounding",
    "self_awareness:uncertainty_calibration",
    "adversarial_resistance:boundaries",
    "autonomy:goals",
]

# Low-diagnostic pairs — weakest signal, expected to vary across agents.
LOW_DIAGNOSTIC_PAIRS: list[str] = [
    "knowledge:personality",
    "boundaries:reasoning_style",
    "personality:temporal_grounding",
]

# Combined lookup for scoring weight assignment.
DIAGNOSTIC_PAIRS: dict[str, str] = {}
for _p in HIGH_DIAGNOSTIC_PAIRS:
    DIAGNOSTIC_PAIRS[_p] = "high"
for _p in MEDIUM_DIAGNOSTIC_PAIRS:
    DIAGNOSTIC_PAIRS[_p] = "medium"
for _p in LOW_DIAGNOSTIC_PAIRS:
    DIAGNOSTIC_PAIRS[_p] = "low"


def cosine_similarity(a: list[float], b: list[float]) -> float:
    """Compute cosine similarity between two vectors using pure Python.

    Returns a value in [-1, 1]. Identical vectors return 1.0.
    Zero-magnitude vectors return 0.0 (orthogonal assumption).

    Uses math.fsum for numerical stability on long vectors without
    requiring numpy as a dependency in this module.
    """
    if len(a) != len(b):
        raise ValueError(f"Vector length mismatch: {len(a)} vs {len(b)}")

    dot = math.fsum(ai * bi for ai, bi in zip(a, b))
    norm_a = math.sqrt(math.fsum(ai * ai for ai in a))
    norm_b = math.sqrt(math.fsum(bi * bi for bi in b))

    if norm_a < 1e-10 or norm_b < 1e-10:
        return 0.0

    return dot / (norm_a * norm_b)


def compute_metametric(
    dimension_vectors: dict[str, list[float]],
) -> dict:
    """Compute the metametric correlation fingerprint.

    Computes cosine similarity for all 91 pairwise combinations of the 14
    dimensions. Missing dimensions contribute 0.0 (orthogonal assumption),
    which preserves backward compatibility with 6-dimension baselines — the
    old dimensions produce real correlations, new dimensions produce 0.0
    until prompts are added.

    Args:
        dimension_vectors: {dimension_name: [float, ...]} — per-dimension
            averaged embedding vectors from an assessment.

    Returns:
        {
            "correlations": {
                "adversarial_resistance:autonomy": 0.742,
                ...  # 91 pairs in canonical order
            },
            "correlation_vector": [0.742, ...],  # 91 values in canonical order
            "metametric_hash": "sha256hex...",
        }
    """
    correlations: dict[str, float] = {}
    correlation_vector: list[float] = []

    for dim_a, dim_b in combinations(DIMENSIONS_SORTED, 2):
        vec_a = dimension_vectors.get(dim_a)
        vec_b = dimension_vectors.get(dim_b)

        if vec_a is None or vec_b is None:
            # Missing dimension — use 0.0 (orthogonal assumption).
            # This keeps the hash stable for old 6-dimension baselines.
            sim = 0.0
        else:
            sim = cosine_similarity(vec_a, vec_b)

        pair_key = f"{dim_a}:{dim_b}"
        correlations[pair_key] = round(sim, 6)
        correlation_vector.append(sim)

    mm_hash = metametric_hash(correlations)

    return {
        "correlations": correlations,
        "correlation_vector": correlation_vector,
        "metametric_hash": mm_hash,
    }


def metametric_hash(correlations: dict[str, float]) -> str:
    """SHA-256 of the canonical correlation JSON.

    Quantizes to 4 decimal places before hashing for cross-platform
    reproducibility (avoids floating-point representation differences
    between Python, Swift, and JavaScript).

    Args:
        correlations: {pair_key: cosine_similarity} — the 91 pair correlations.

    Returns:
        64-character hex SHA-256 hash string.
    """
    # Build the canonical vector in sorted pair order.
    # Pairs are already in canonical order if generated by compute_metametric,
    # but we sort here defensively for standalone use.
    sorted_keys = sorted(correlations.keys())
    quantized = [round(correlations[k], 4) for k in sorted_keys]
    payload = json.dumps(quantized, separators=(",", ":"))
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def compare_metametrics(
    baseline_corr: dict[str, float],
    test_corr: dict[str, float],
    high_diagnostic_weight: float = 2.0,
    medium_diagnostic_weight: float = 1.0,
    low_diagnostic_weight: float = 0.5,
) -> dict:
    """Compare two metametric snapshots and produce a divergence score.

    The divergence score is 0-100 where 0 means identical structural
    fingerprints and 100 means completely different. Higher-diagnostic
    pairs are weighted more heavily because they carry more identity signal.

    Args:
        baseline_corr: {pair_key: cosine_similarity} from the baseline.
        test_corr: {pair_key: cosine_similarity} from the test.
        high_diagnostic_weight: Weight for high-diagnostic pairs.
        medium_diagnostic_weight: Weight for medium-diagnostic pairs.
        low_diagnostic_weight: Weight for low-diagnostic and unclassified pairs.

    Returns:
        {
            "divergence": float,         # 0-100, higher = more different
            "pair_divergences": {pair_key: float},
            "top_divergent_pairs": [(pair_key, divergence), ...],  # top 10
            "structural_intact": bool,   # True if high-diagnostic pairs < 0.25
        }
    """
    pair_divergences: dict[str, float] = {}
    weighted_sum = 0.0
    weight_total = 0.0

    # Iterate over all pairs from both snapshots (union handles old vs new).
    all_keys = sorted(set(baseline_corr.keys()) | set(test_corr.keys()))

    for pair_key in all_keys:
        baseline_val = baseline_corr.get(pair_key, 0.0)
        test_val = test_corr.get(pair_key, 0.0)
        divergence = abs(baseline_val - test_val)
        pair_divergences[pair_key] = round(divergence, 4)

        # Determine weight based on diagnostic category.
        category = DIAGNOSTIC_PAIRS.get(pair_key)
        if category == "high":
            weight = high_diagnostic_weight
        elif category == "medium":
            weight = medium_diagnostic_weight
        else:
            weight = low_diagnostic_weight

        weighted_sum += divergence * weight
        weight_total += weight

    overall_divergence = (weighted_sum / weight_total * 100) if weight_total > 0 else 0.0

    # Sort by divergence descending for top-N.
    sorted_pairs = sorted(
        pair_divergences.items(), key=lambda x: x[1], reverse=True,
    )

    # Structural integrity: are the high-diagnostic pairs within tolerance?
    # A tolerance of 0.25 means the correlation can shift by up to 0.25
    # before the system flags structural damage.
    structural_intact = all(
        pair_divergences.get(p, 0.0) < 0.25
        for p in HIGH_DIAGNOSTIC_PAIRS
    )

    return {
        "divergence": round(min(100.0, overall_divergence), 2),
        "pair_divergences": pair_divergences,
        "top_divergent_pairs": sorted_pairs[:10],
        "structural_intact": structural_intact,
    }
