"""Small, optional learned reranker for the Astra/HSCM memory field.

The runtime implementation is deliberately NumPy-only.  Qiskit is used by the
experiment harness to train and validate the same two-qubit circuit, but normal
Astra recall only evaluates the eight learned angles stored in a JSON artifact.
"""
from __future__ import annotations

import hashlib
import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Mapping, Sequence

import numpy as np


FIELD_ARTIFACT_SCHEMA_VERSION = 1
FIELD_DEPLOYMENT_SCHEMA_VERSION = 1
FIELD_FEATURE_ORDER = (
    "semantic_support",
    "lexical_support",
    "coherence",
    "path_quality",
)


def _finite_vector(values: Sequence[float], size: int, label: str) -> np.ndarray:
    result = np.asarray(values, dtype=np.float64)
    if result.shape != (size,) or not np.all(np.isfinite(result)):
        raise ValueError(f"{label} must contain {size} finite values")
    return result


@dataclass(frozen=True)
class FieldRerankerArtifact:
    """Validated, secret-free representation of one learned field circuit."""

    artifact_id: str
    weights: tuple[float, ...]
    fusion_weight: float = 0.10
    score_calibration: str = "none"
    feature_min: tuple[float, ...] = (0.0, 0.0, 0.0, 0.0)
    feature_max: tuple[float, ...] = (1.0, 1.0, 1.0, 1.0)
    training_stage: str = "experimental"
    metadata: Mapping | None = None

    def __post_init__(self) -> None:
        if not str(self.artifact_id).strip():
            raise ValueError("field artifact_id must not be empty")
        _finite_vector(self.weights, 8, "field weights")
        lo = _finite_vector(self.feature_min, 4, "feature_min")
        hi = _finite_vector(self.feature_max, 4, "feature_max")
        if np.any(hi <= lo):
            raise ValueError("every feature_max must exceed feature_min")
        if not math.isfinite(float(self.fusion_weight)) or not 0.0 <= float(self.fusion_weight) <= 0.25:
            raise ValueError("fusion_weight must be finite and within [0, 0.25]")
        if self.score_calibration not in {"none", "batch_max"}:
            raise ValueError("unsupported field score calibration")
        try:
            json.dumps(dict(self.metadata or {}), allow_nan=False)
        except (TypeError, ValueError) as exc:
            raise ValueError("field metadata must be finite JSON data") from exc

    @classmethod
    def from_mapping(cls, payload: Mapping) -> "FieldRerankerArtifact":
        if int(payload.get("schema_version", -1)) != FIELD_ARTIFACT_SCHEMA_VERSION:
            raise ValueError("unsupported field artifact schema")
        if tuple(payload.get("feature_order", ())) != FIELD_FEATURE_ORDER:
            raise ValueError("field artifact feature order mismatch")
        circuit = payload.get("circuit") or {}
        if (int(circuit.get("qubits", -1)) != 2
                or int(circuit.get("trainable_weights", -1)) != 8
                or circuit.get("ansatz") != "astra-field-v1"
                or circuit.get("output") != "even-parity-probability"):
            raise ValueError("unsupported field circuit")
        scaler = payload.get("feature_scaler") or {}
        return cls(
            artifact_id=str(payload.get("artifact_id", "")),
            weights=tuple(float(value) for value in payload.get("weights", ())),
            fusion_weight=float(payload.get("fusion_weight", 0.10)),
            score_calibration=str(payload.get("score_calibration", "none")),
            feature_min=tuple(float(value) for value in scaler.get("min", ())),
            feature_max=tuple(float(value) for value in scaler.get("max", ())),
            training_stage=str(payload.get("training_stage", "experimental")),
            metadata=dict(payload.get("metadata") or {}),
        )

    def to_mapping(self) -> dict:
        return {
            "schema_version": FIELD_ARTIFACT_SCHEMA_VERSION,
            "artifact_id": self.artifact_id,
            "feature_order": list(FIELD_FEATURE_ORDER),
            "feature_scaler": {
                "min": [float(value) for value in self.feature_min],
                "max": [float(value) for value in self.feature_max],
            },
            "circuit": {
                "ansatz": "astra-field-v1",
                "qubits": 2,
                "trainable_weights": 8,
                "output": "even-parity-probability",
            },
            "weights": [float(value) for value in self.weights],
            "fusion_weight": float(self.fusion_weight),
            "score_calibration": self.score_calibration,
            "training_stage": self.training_stage,
            "metadata": dict(self.metadata or {}),
        }


def _ry(angle: float) -> np.ndarray:
    half = 0.5 * float(angle)
    return np.asarray([[math.cos(half), -math.sin(half)],
                       [math.sin(half), math.cos(half)]], dtype=np.complex128)


def _rz(angle: float) -> np.ndarray:
    half = 0.5 * float(angle)
    return np.asarray([[np.exp(-1j * half), 0.0],
                       [0.0, np.exp(1j * half)]], dtype=np.complex128)


_IDENTITY = np.eye(2, dtype=np.complex128)
_CZ = np.diag([1.0, 1.0, 1.0, -1.0]).astype(np.complex128)


def _single_qubit(gate: np.ndarray, qubit: int) -> np.ndarray:
    # Qiskit basis ordering is |q1 q0>; q0 is the least-significant qubit.
    return np.kron(_IDENTITY, gate) if int(qubit) == 0 else np.kron(gate, _IDENTITY)


def field_circuit_probability(features: Sequence[float], weights: Sequence[float]) -> float:
    """Evaluate the v1 circuit's even-parity probability exactly."""
    values = np.clip(_finite_vector(features, 4, "field features"), 0.0, 1.0)
    theta = _finite_vector(weights, 8, "field weights")
    state = np.asarray([1.0, 0.0, 0.0, 0.0], dtype=np.complex128)
    operations = (
        (_ry(math.pi * values[0]), 0), (_rz(math.pi * values[1]), 0),
        (_ry(math.pi * values[2]), 1), (_rz(math.pi * values[3]), 1),
    )
    for gate, qubit in operations:
        state = _single_qubit(gate, qubit) @ state
    state = _CZ @ state
    for gate, qubit in ((_ry(theta[0]), 0), (_rz(theta[1]), 0),
                        (_ry(theta[2]), 1), (_rz(theta[3]), 1)):
        state = _single_qubit(gate, qubit) @ state
    state = _CZ @ state
    for gate, qubit in ((_ry(theta[4]), 0), (_rz(theta[5]), 0),
                        (_ry(theta[6]), 1), (_rz(theta[7]), 1)):
        state = _single_qubit(gate, qubit) @ state
    probability = float(abs(state[0]) ** 2 + abs(state[3]) ** 2)
    return float(np.clip(probability, 0.0, 1.0))


class QuantumFieldReranker:
    """Runtime scorer backed by a validated two-qubit field artifact."""

    def __init__(self, artifact: FieldRerankerArtifact):
        self.artifact = artifact
        self.artifact_id = artifact.artifact_id
        self.fusion_weight = float(artifact.fusion_weight)

    def _normalize(self, features: np.ndarray) -> np.ndarray:
        values = np.asarray(features, dtype=np.float64)
        if values.ndim != 2 or values.shape[1] != 4:
            raise ValueError("field feature batch must have shape (n, 4)")
        if not np.all(np.isfinite(values)):
            raise ValueError("field feature batch contains non-finite values")
        lo = np.asarray(self.artifact.feature_min, dtype=np.float64)
        hi = np.asarray(self.artifact.feature_max, dtype=np.float64)
        return np.clip((values - lo) / (hi - lo), 0.0, 1.0)

    def score_batch(self, features: np.ndarray) -> np.ndarray:
        normalized = self._normalize(features)
        scores = np.asarray([
            field_circuit_probability(row, self.artifact.weights)
            for row in normalized
        ], dtype=np.float64)
        if self.artifact.score_calibration == "batch_max" and len(scores):
            scores = scores / max(float(np.max(scores)), 1e-12)
        return scores


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for block in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def verify_active_deployment(
        artifact_path: str | Path,
        deployment_path: str | Path | None = None) -> dict:
    """Verify the local allow-list record required for active reranking."""
    artifact = Path(artifact_path)
    deployment = (Path(deployment_path) if deployment_path is not None
                  else artifact.with_suffix(".deployment.json"))
    payload = json.loads(deployment.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError("field deployment root must be an object")
    if int(payload.get("schema_version", -1)) != FIELD_DEPLOYMENT_SCHEMA_VERSION:
        raise ValueError("unsupported field deployment schema")
    if payload.get("status") != "active":
        raise ValueError("field deployment is not active")
    if payload.get("rollback_mode") != "shadow":
        raise ValueError("active field deployment must declare shadow rollback")
    if payload.get("artifact_sha256") != _sha256(artifact):
        raise ValueError("field deployment artifact hash mismatch")
    gates = payload.get("activation_gates")
    if (not isinstance(gates, dict) or not gates
            or any(value is not True for value in gates.values())):
        raise ValueError("field deployment activation gates are not all passing")
    return payload


def load_field_reranker(path: str | Path, *, require_active: bool = False,
                        deployment_path: str | Path | None = None
                        ) -> QuantumFieldReranker:
    artifact_path = Path(path)
    payload = json.loads(artifact_path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError("field artifact root must be an object")
    reranker = QuantumFieldReranker(FieldRerankerArtifact.from_mapping(payload))
    if require_active:
        deployment = verify_active_deployment(artifact_path, deployment_path)
        if deployment.get("artifact_id") != reranker.artifact_id:
            raise ValueError("field deployment artifact id mismatch")
        feature_sha = (reranker.artifact.metadata or {}).get("feature_sha256")
        if deployment.get("feature_sha256") != feature_sha:
            raise ValueError("field deployment feature hash mismatch")
    return reranker