"""Audit whether a representation migration preserves declared behavior.

The implementation is intentionally dependency-light.  It makes representation
identity, coordinate dimensions, downstream probes, and tolerated drift explicit
instead of assuming that two equally shaped vectors mean the same thing.
"""

from __future__ import annotations

from dataclasses import dataclass
from itertools import combinations
import math
from typing import Iterable, Sequence


Vector = tuple[float, ...]
Matrix = tuple[Vector, ...]


def _finite_number(value: object, label: str) -> float:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError(f"{label} must be a real number")
    converted = float(value)
    if not math.isfinite(converted):
        raise ValueError(f"{label} must be finite")
    return converted


def _freeze_vector(values: Sequence[float], label: str) -> Vector:
    if isinstance(values, (str, bytes)) or not isinstance(values, Sequence):
        raise ValueError(f"{label} must be a finite numeric sequence")
    frozen = tuple(
        _finite_number(value, f"{label}[{index}]")
        for index, value in enumerate(values)
    )
    if not frozen:
        raise ValueError(f"{label} must not be empty")
    return frozen


def _freeze_matrix(rows: Sequence[Sequence[float]], label: str) -> Matrix:
    if isinstance(rows, (str, bytes)) or not isinstance(rows, Sequence):
        raise ValueError(f"{label} must be a sequence of rows")
    frozen = tuple(
        _freeze_vector(row, f"{label}[{index}]")
        for index, row in enumerate(rows)
    )
    if not frozen:
        raise ValueError(f"{label} must not be empty")
    width = len(frozen[0])
    if any(len(row) != width for row in frozen):
        raise ValueError(f"{label} must be rectangular")
    return frozen


def _require_identifier(value: object, label: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"{label} must be a non-empty string")
    return value


def dot(left: Vector, right: Vector) -> float:
    if len(left) != len(right):
        raise ValueError("dot product dimension mismatch")
    products = tuple(a * b for a, b in zip(left, right))
    if any(not math.isfinite(product) for product in products):
        raise ValueError("dot product produced a non-finite intermediate")
    return _finite_number(math.fsum(products), "dot product result")


def l2_distance(left: Vector, right: Vector) -> float:
    if len(left) != len(right):
        raise ValueError("distance dimension mismatch")
    differences = tuple(a - b for a, b in zip(left, right))
    if any(not math.isfinite(value) for value in differences):
        raise ValueError("distance produced a non-finite intermediate")
    return _finite_number(math.hypot(*differences), "distance result")


def matvec(matrix: Matrix, vector: Vector) -> Vector:
    if any(len(row) != len(vector) for row in matrix):
        raise ValueError("matrix/vector dimension mismatch")
    return tuple(dot(row, vector) for row in matrix)


@dataclass(frozen=True)
class RepresentationSet:
    representation_id: str
    entries: tuple[tuple[str, Vector], ...]
    dimension: int

    def __post_init__(self) -> None:
        _require_identifier(self.representation_id, "representation_id")
        if not isinstance(self.entries, tuple) or len(self.entries) < 2:
            raise ValueError("entries must be an immutable tuple with at least two items")
        seen: set[str] = set()
        for index, entry in enumerate(self.entries):
            if not isinstance(entry, tuple) or len(entry) != 2:
                raise ValueError(f"entries[{index}] must be an immutable (id, vector) tuple")
            item_id, vector = entry
            _require_identifier(item_id, f"entries[{index}].id")
            if item_id in seen:
                raise ValueError(f"duplicate item id: {item_id}")
            seen.add(item_id)
            if not isinstance(vector, tuple):
                raise ValueError(f"entries[{index}].vector must be immutable")
            checked = _freeze_vector(vector, f"entries[{index}].vector")
            if checked != vector:
                raise ValueError(f"entries[{index}].vector must contain normalized floats")
        if isinstance(self.dimension, bool) or not isinstance(self.dimension, int):
            raise ValueError("dimension must be an integer")
        if self.dimension != len(self.entries[0][1]) or any(
            len(vector) != self.dimension for _, vector in self.entries
        ):
            raise ValueError("dimension must match every representation vector")

    @classmethod
    def capture(
        cls,
        representation_id: str,
        entries: Iterable[tuple[str, Sequence[float]]],
    ) -> "RepresentationSet":
        frozen_entries: list[tuple[str, Vector]] = []
        seen: set[str] = set()
        for index, entry in enumerate(entries):
            if not isinstance(entry, (tuple, list)) or len(entry) != 2:
                raise ValueError(f"entries[{index}] must be an (id, vector) pair")
            item_id = _require_identifier(entry[0], f"entries[{index}].id")
            if item_id in seen:
                raise ValueError(f"duplicate item id: {item_id}")
            seen.add(item_id)
            frozen_entries.append(
                (item_id, _freeze_vector(entry[1], f"entries[{index}].vector"))
            )
        if len(frozen_entries) < 2:
            raise ValueError("at least two representation entries are required")
        dimension = len(frozen_entries[0][1])
        if any(len(vector) != dimension for _, vector in frozen_entries):
            raise ValueError("all representation vectors must have the same dimension")
        return cls(
            representation_id=_require_identifier(
                representation_id, "representation_id"
            ),
            entries=tuple(frozen_entries),
            dimension=dimension,
        )


@dataclass(frozen=True)
class LinearTransform:
    name: str
    source_representation_id: str
    target_representation_id: str
    rows: Matrix

    def __post_init__(self) -> None:
        _require_identifier(self.name, "transform name")
        _require_identifier(self.source_representation_id, "source_representation_id")
        _require_identifier(self.target_representation_id, "target_representation_id")
        if not isinstance(self.rows, tuple) or any(
            not isinstance(row, tuple) for row in self.rows
        ):
            raise ValueError("transform rows must be immutable tuples")
        checked = _freeze_matrix(self.rows, "transform rows")
        if checked != self.rows:
            raise ValueError("transform rows must contain normalized floats")

    @classmethod
    def capture(
        cls,
        *,
        name: str,
        source_representation_id: str,
        target_representation_id: str,
        rows: Sequence[Sequence[float]],
    ) -> "LinearTransform":
        return cls(
            name=_require_identifier(name, "transform name"),
            source_representation_id=_require_identifier(
                source_representation_id, "source_representation_id"
            ),
            target_representation_id=_require_identifier(
                target_representation_id, "target_representation_id"
            ),
            rows=_freeze_matrix(rows, "transform rows"),
        )


@dataclass(frozen=True)
class ProbeMigration:
    name: str
    source_weights: Vector
    target_weights: Vector

    def __post_init__(self) -> None:
        _require_identifier(self.name, "probe name")
        for label, weights in (
            ("source_weights", self.source_weights),
            ("target_weights", self.target_weights),
        ):
            if not isinstance(weights, tuple):
                raise ValueError(f"{label} must be immutable")
            checked = _freeze_vector(weights, label)
            if checked != weights:
                raise ValueError(f"{label} must contain normalized floats")

    @classmethod
    def capture(
        cls,
        name: str,
        source_weights: Sequence[float],
        target_weights: Sequence[float],
    ) -> "ProbeMigration":
        return cls(
            name=_require_identifier(name, "probe name"),
            source_weights=_freeze_vector(source_weights, "source_weights"),
            target_weights=_freeze_vector(target_weights, "target_weights"),
        )


@dataclass(frozen=True)
class InvarianceReport:
    source_representation_id: str
    target_representation_id: str
    source_dimension: int
    target_dimension: int
    item_ids: tuple[str, ...]
    max_score_drift: float
    max_distance_drift: float
    collisions: tuple[tuple[str, str], ...]
    tolerance: float

    @property
    def passed(self) -> bool:
        return (
            self.max_score_drift <= self.tolerance
            and self.max_distance_drift <= self.tolerance
            and not self.collisions
        )


def audit_invariance(
    representations: RepresentationSet,
    transform: LinearTransform,
    probes: Iterable[ProbeMigration],
    *,
    tolerance: float = 1e-9,
) -> InvarianceReport:
    """Compare downstream scores and pairwise geometry across coordinates.

    A probe migration declares how one downstream linear score changes basis.
    The audit does not guess that mapping: silently reusing old weights is one of
    the representation-migration failures this artifact is meant to expose.
    """

    checked_tolerance = _finite_number(tolerance, "tolerance")
    if checked_tolerance < 0:
        raise ValueError("tolerance must be non-negative")
    if transform.source_representation_id != representations.representation_id:
        raise ValueError("transform source does not match representation identity")
    if len(transform.rows[0]) != representations.dimension:
        raise ValueError("transform input dimension does not match representations")

    frozen_probes = tuple(probes)
    if not frozen_probes:
        raise ValueError("at least one downstream probe migration is required")
    if not all(isinstance(probe, ProbeMigration) for probe in frozen_probes):
        raise ValueError("probes must contain ProbeMigration values")
    target_dimension = len(transform.rows)
    probe_names: set[str] = set()
    for probe in frozen_probes:
        if probe.name in probe_names:
            raise ValueError(f"duplicate probe name: {probe.name}")
        probe_names.add(probe.name)
        if len(probe.source_weights) != representations.dimension:
            raise ValueError(f"source weight dimension mismatch for probe: {probe.name}")
        if len(probe.target_weights) != target_dimension:
            raise ValueError(f"target weight dimension mismatch for probe: {probe.name}")

    transformed = tuple(
        (item_id, matvec(transform.rows, vector))
        for item_id, vector in representations.entries
    )
    transformed_by_id = dict(transformed)

    score_drifts: list[float] = []
    for item_id, source_vector in representations.entries:
        target_vector = transformed_by_id[item_id]
        for probe in frozen_probes:
            source_score = dot(probe.source_weights, source_vector)
            target_score = dot(probe.target_weights, target_vector)
            score_drifts.append(abs(source_score - target_score))

    distance_drifts: list[float] = []
    collisions: list[tuple[str, str]] = []
    for (left_id, left), (right_id, right) in combinations(
        representations.entries, 2
    ):
        source_distance = l2_distance(left, right)
        target_distance = l2_distance(
            transformed_by_id[left_id], transformed_by_id[right_id]
        )
        distance_drifts.append(abs(source_distance - target_distance))
        if source_distance > checked_tolerance and target_distance <= checked_tolerance:
            collisions.append((left_id, right_id))

    return InvarianceReport(
        source_representation_id=representations.representation_id,
        target_representation_id=transform.target_representation_id,
        source_dimension=representations.dimension,
        target_dimension=target_dimension,
        item_ids=tuple(item_id for item_id, _ in representations.entries),
        max_score_drift=max(score_drifts),
        max_distance_drift=max(distance_drifts),
        collisions=tuple(collisions),
        tolerance=checked_tolerance,
    )


EXAMPLE_REPRESENTATIONS = RepresentationSet.capture(
    "support-intent-v3",
    [
        ("ticket-a", [1.0, 2.0]),
        ("ticket-b", [3.0, 1.0]),
        ("ticket-c", [-1.0, 1.0]),
    ],
)
EXAMPLE_ROTATION = LinearTransform.capture(
    name="quarter-turn-basis-change",
    source_representation_id="support-intent-v3",
    target_representation_id="support-intent-v4",
    rows=[[0.0, -1.0], [1.0, 0.0]],
)
EXAMPLE_PROBES = (
    ProbeMigration.capture("escalation", [2.0, -1.0], [1.0, 2.0]),
)


if __name__ == "__main__":
    report = audit_invariance(
        EXAMPLE_REPRESENTATIONS, EXAMPLE_ROTATION, EXAMPLE_PROBES
    )
    print(
        f"representation={report.source_representation_id}"
        f"->{report.target_representation_id}"
    )
    print(f"dimensions={report.source_dimension}->{report.target_dimension}")
    print(f"max_score_drift={report.max_score_drift:.6f}")
    print(f"max_distance_drift={report.max_distance_drift:.6f}")
    print(f"collisions={len(report.collisions)}")