"""Faithful sparse-jaw diffusion experiment for Claim 5.

The paper does not identify or release the jaw volume shown in Figure 1.  This
module therefore uses the independently published OpenMandible cortical-bone
model.  The source is commit- and hash-pinned, and the substitution is recorded
as a limitation instead of being presented as the authors' original scan.
"""

from __future__ import annotations

import hashlib
import math
import time
import urllib.request
from collections import deque
from typing import Callable

import numpy as np
from scipy.ndimage import gaussian_filter
from scipy.sparse.linalg import LinearOperator, eigsh

from armadillo import USER_AGENT, _normalize_to_unit_ball, _sample_surface


def _download(spec: dict) -> tuple[bytes, dict]:
    request = urllib.request.Request(
        spec["url"], headers={"User-Agent": USER_AGENT}
    )
    with urllib.request.urlopen(request, timeout=180) as response:
        payload = response.read()
    observed = hashlib.sha256(payload).hexdigest()
    if observed != spec["sha256"]:
        raise RuntimeError(
            "OpenMandible hash mismatch: "
            f"expected {spec['sha256']}, got {observed}"
        )
    return payload, {
        "dataset": spec["dataset"],
        "dataset_paper_doi": spec["dataset_paper_doi"],
        "repository": spec["repository"],
        "commit": spec["commit"],
        "url": spec["url"],
        "sha256": observed,
        "bytes": len(payload),
        "retrieval_user_agent": USER_AGENT,
    }


def _parse_ascii_stl(payload: bytes) -> tuple[np.ndarray, np.ndarray, dict]:
    coordinates: list[list[float]] = []
    for raw_line in payload.splitlines():
        fields = raw_line.split()
        if fields and fields[0] == b"vertex":
            if len(fields) != 4:
                raise RuntimeError("malformed OpenMandible STL vertex")
            coordinates.append(
                [float(fields[1]), float(fields[2]), float(fields[3])]
            )
    vertices = np.asarray(coordinates, dtype=np.float64)
    if vertices.shape[0] == 0 or vertices.shape[0] % 3:
        raise RuntimeError("OpenMandible STL is not an all-triangle mesh")
    faces = np.arange(vertices.shape[0], dtype=np.int64).reshape(-1, 3)
    triangles = vertices[faces]
    doubled_area = np.linalg.norm(
        np.cross(
            triangles[:, 1] - triangles[:, 0],
            triangles[:, 2] - triangles[:, 0],
        ),
        axis=1,
    )
    if np.any(doubled_area <= 0.0):
        raise RuntimeError("OpenMandible STL contains degenerate triangles")
    return vertices, faces, {
        "format": "ASCII STL",
        "triangle_count": int(faces.shape[0]),
        "vertex_records": int(vertices.shape[0]),
        "all_triangles": True,
        "degenerate_triangles": 0,
    }


def _sparse_surface_voxels(
    vertices: np.ndarray,
    faces: np.ndarray,
    edge: float,
    sample_count: int,
    seed: int,
) -> tuple[np.ndarray, np.ndarray, dict]:
    samples, sampling = _sample_surface(
        vertices, faces, sample_count, seed
    )
    grid_size = int(round(2.0 / edge))
    indices = np.floor((samples + 1.0) / edge).astype(np.int64)
    indices = np.clip(indices, 0, grid_size - 1)
    indices = np.unique(indices, axis=0)
    centers = -1.0 + edge * (indices.astype(np.float64) + 0.5)
    return indices, centers, {
        **sampling,
        "representation": "sparse regular-grid surface voxels",
        "grid_shape": [grid_size, grid_size, grid_size],
        "voxel_edge": float(edge),
        "nonempty_voxels": int(indices.shape[0]),
        "occupancy_fraction": float(indices.shape[0] / grid_size**3),
        "index_extent": (
            indices.max(axis=0) - indices.min(axis=0) + 1
        ).tolist(),
    }


def _largest_component_fraction(indices: np.ndarray) -> float:
    lookup = {tuple(int(value) for value in row) for row in indices}
    remaining = set(lookup)
    largest = 0
    offsets = (
        (1, 0, 0),
        (-1, 0, 0),
        (0, 1, 0),
        (0, -1, 0),
        (0, 0, 1),
        (0, 0, -1),
    )
    while remaining:
        root = remaining.pop()
        queue: deque[tuple[int, int, int]] = deque([root])
        size = 0
        while queue:
            current = queue.popleft()
            size += 1
            for offset in offsets:
                neighbor = (
                    current[0] + offset[0],
                    current[1] + offset[1],
                    current[2] + offset[2],
                )
                if neighbor in remaining:
                    remaining.remove(neighbor)
                    queue.append(neighbor)
        largest = max(largest, size)
    return float(largest / max(indices.shape[0], 1))


def _voxel_gaussian_operator(
    indices: np.ndarray,
    grid_size: int,
    sigma_grid: float,
    truncate: float,
) -> Callable[[np.ndarray], np.ndarray]:
    workspace = np.zeros(
        (grid_size, grid_size, grid_size), dtype=np.float64
    )

    def matvec(vector: np.ndarray) -> np.ndarray:
        workspace.fill(0.0)
        workspace[indices[:, 0], indices[:, 1], indices[:, 2]] = vector
        convolved = gaussian_filter(
            workspace,
            sigma=sigma_grid,
            mode="constant",
            cval=0.0,
            truncate=truncate,
        )
        return convolved[
            indices[:, 0], indices[:, 1], indices[:, 2]
        ]

    return matvec


def _sinkhorn(
    kernel_matvec: Callable[[np.ndarray], np.ndarray],
    weights: np.ndarray,
) -> tuple[np.ndarray, list[float], int, float]:
    scaling = np.ones(weights.shape[0], dtype=np.float64)
    curve: list[float] = []
    threshold_iteration = -1
    residual_max = math.inf
    for iteration in range(1, 301):
        kernel_scaled = kernel_matvec(weights * scaling)
        row_values = scaling * kernel_scaled
        mean_error = float(np.sum(weights * np.abs(row_values - 1.0)))
        curve.append(mean_error)
        if threshold_iteration < 0 and mean_error < 1e-3:
            threshold_iteration = iteration
        residual_max = float(np.max(np.abs(row_values - 1.0)))
        if residual_max < 1e-12:
            return scaling, curve, threshold_iteration, residual_max
        scaling = np.sqrt(
            scaling / np.maximum(kernel_scaled, 1e-300)
        )
    raise RuntimeError(
        f"OpenMandible Sinkhorn did not converge: {residual_max}"
    )


def _record_diffusion(
    indices: np.ndarray,
    centers: np.ndarray,
    weights: np.ndarray,
    kernel_matvec: Callable[[np.ndarray], np.ndarray],
    scaling: np.ndarray,
    steps: list[int],
    normalization: str,
) -> dict:
    def apply(signal: np.ndarray) -> np.ndarray:
        return scaling * kernel_matvec(weights * scaling * signal)

    row_values = apply(np.ones(weights.shape[0], dtype=np.float64))
    source_index = int(np.argmin(centers[:, 0]))
    source = centers[source_index]
    signal = np.zeros(weights.shape[0], dtype=np.float64)
    signal[source_index] = 1.0 / weights[source_index]
    snapshots: list[dict] = []
    maximum_step = max(steps)
    for step in range(maximum_step + 1):
        if step in steps:
            next_signal = apply(signal)
            constant = float(np.sum(weights * signal))
            centered_signal = signal - constant
            q_roughness = float(
                np.sum(weights * signal * (signal - next_signal))
            )
            weighted_l2_from_constant = float(
                np.sum(weights * centered_signal**2)
            )
            spatial_second_moment = float(
                np.sum(
                    weights
                    * np.maximum(signal, 0.0)
                    * np.sum((centers - source) ** 2, axis=1)
                )
            )
            snapshots.append(
                {
                    "step": step,
                    "mass": constant,
                    "minimum": float(signal.min()),
                    "maximum": float(signal.max()),
                    "q_roughness": q_roughness,
                    "weighted_l2_from_constant": (
                        weighted_l2_from_constant
                    ),
                    "spatial_second_moment": spatial_second_moment,
                    "signal": signal.tolist(),
                }
            )
        if step < maximum_step:
            signal = apply(signal)

    diagonal = np.sqrt(weights) * scaling
    symmetric_operator = LinearOperator(
        (weights.shape[0], weights.shape[0]),
        matvec=lambda vector: diagonal
        * kernel_matvec(diagonal * vector),
        rmatvec=lambda vector: diagonal
        * kernel_matvec(diagonal * vector),
        dtype=np.float64,
    )
    largest = eigsh(
        symmetric_operator,
        k=6,
        which="LA",
        return_eigenvectors=False,
        tol=2e-9,
        maxiter=1_000,
    )
    smallest = eigsh(
        symmetric_operator,
        k=3,
        which="SA",
        return_eigenvectors=False,
        tol=2e-9,
        maxiter=1_000,
    )
    return {
        "normalization": normalization,
        "source_index": source_index,
        "source_voxel_index": indices[source_index].tolist(),
        "row_residual_max": float(np.max(np.abs(row_values - 1.0))),
        "constant_preservation_max_error": float(
            np.max(np.abs(row_values - 1.0))
        ),
        "largest_symmetric_eigenvalues": np.sort(largest)[::-1].tolist(),
        "smallest_symmetric_eigenvalues": np.sort(smallest).tolist(),
        "snapshots": snapshots,
    }


def run_claim5_jaw(config: dict, spectral_result: dict) -> tuple[dict, dict]:
    spec = config["claim5_jaw"]
    start = time.perf_counter()
    payload, source = _download(spec)
    vertices, faces, mesh = _parse_ascii_stl(payload)
    vertices, normalization = _normalize_to_unit_ball(vertices)
    indices, centers, voxelization = _sparse_surface_voxels(
        vertices,
        faces,
        float(spec["voxel_edge"]),
        int(spec["surface_sample_count"]),
        int(spec["seed"]),
    )
    voxelization["largest_6_connected_component_fraction"] = (
        _largest_component_fraction(indices)
    )
    grid_size = int(round(2.0 / float(spec["voxel_edge"])))
    sigma_grid = float(spec["kernel_sigma"]) / float(spec["voxel_edge"])
    truncate = float(spec["gaussian_truncate_sigma"])
    kernel_matvec = _voxel_gaussian_operator(
        indices, grid_size, sigma_grid, truncate
    )
    weights = np.full(indices.shape[0], 1.0 / indices.shape[0])
    scaling, curve, threshold_iteration, residual = _sinkhorn(
        kernel_matvec, weights
    )
    steps = [int(value) for value in spec["diffusion_steps"]]
    sinkhorn_record = _record_diffusion(
        indices,
        centers,
        weights,
        kernel_matvec,
        scaling,
        steps,
        "symmetric Sinkhorn",
    )
    sinkhorn_record.update(
        {
            "sinkhorn_curve": curve,
            "sinkhorn_iterations": len(curve),
            "sinkhorn_iterations_to_1e-3": threshold_iteration,
            "sinkhorn_residual_max": residual,
        }
    )
    raw_record = _record_diffusion(
        indices,
        centers,
        weights,
        kernel_matvec,
        np.ones_like(weights),
        steps,
        "raw unnormalized Gaussian",
    )
    modalities = spectral_result["modalities"]
    cross_modalities = {
        "point_cloud": modalities["point_5000"],
        "covariance_aware_gmm": modalities["gmm_500"],
        "sparse_armadillo_voxels": modalities["surface_voxels"],
    }
    common = {
        "claim_id": 5,
        "source_statement": (
            "The method is demonstrated on point clouds, sparse voxel "
            "grids (jaw bone geometry), and Gaussian mixture models with "
            "covariance-aware kernels, showing Laplacian-like smoothing."
        ),
        "paper_source_anchor": "Figure 1, Figure 3, Sections 5-6, Eq.6",
        "jaw_source": source,
        "jaw_source_substitution": (
            "OpenMandible cortical bone replaces the paper's unidentified "
            "and unreleased jaw scan; it is not claimed to be the same scan."
        ),
        "mesh": mesh,
        "normalization": normalization,
        "voxelization": voxelization,
        "voxel_indices": indices.tolist(),
        "weights": weights.tolist(),
        "kernel": {
            "type": "matrix-free separable Gaussian convolution",
            "physical_sigma": float(spec["kernel_sigma"]),
            "sigma_grid_cells": sigma_grid,
            "truncate_sigma": truncate,
            "maximum_omitted_axis_weight": float(
                math.exp(-0.5 * truncate**2)
            ),
        },
        "cross_modalities": cross_modalities,
        "runtime_seconds": time.perf_counter() - start,
        "seed": int(spec["seed"]),
    }
    actual = {
        **common,
        "scaling": scaling.tolist(),
        "diffusion": sinkhorn_record,
    }
    negative = {
        **common,
        "scaling": np.ones_like(weights).tolist(),
        "diffusion": raw_record,
        "negative_control": (
            "omit Sinkhorn scaling while retaining the same jaw, voxels, "
            "kernel, weights, source signal, and evaluation checks"
        ),
    }
    return actual, negative