"""Private frozen-source prior-weighting probe; no production edits or timing claim."""
from __future__ import annotations

import dataclasses
import hashlib
import inspect
import json
import subprocess
import types
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import jsonschema
import numpy as np
from pydantic import BaseModel, ConfigDict, TypeAdapter
from scipy.linalg import cholesky

from calibrex.core import continuous_time_sparse as sparse
from calibrex.core.continuous_time_marginalization import KnotMarginalizationPrior
from calibrex.core.geometry import SE3
from calibrex.core.se3_manifold import se3_exp, se3_log

ROOT = Path(__file__).resolve().parents[2]
OUT = Path(__file__).resolve().parent
SOURCE = "8861ed65ebbcffb13a182ac4cd549db88ef763f2"


def sha(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def serializable(value: Any) -> Any:
    if isinstance(value, np.ndarray):
        return value.tolist()
    raise TypeError(type(value).__name__)


def write(name: str, data: Any) -> None:
    path = OUT / name
    assert not path.exists(), path
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes((json.dumps(data, indent=2, allow_nan=False, default=serializable) + "\n").encode())


def candidate_factor(information: np.ndarray) -> np.ndarray:
    # Only finite symmetric PSD fixtures are admitted to this diagnostic.
    # This does not choose the eventual public validation/tolerance policy.
    assert information.shape == (6, 6) and np.isfinite(information).all()
    assert np.array_equal(information, information.T)
    try:
        return cholesky(information, lower=False)
    except np.linalg.LinAlgError:
        values, vectors = np.linalg.eigh(information)
        roundoff = 60 * np.finfo(float).eps * max(1.0, np.max(np.abs(values)))
        assert values.min() >= -roundoff
        return np.diag(np.sqrt(np.maximum(values, 0))) @ vectors.T


def candidate_blocks(prior: KnotMarginalizationPrior, knot: SE3) -> tuple[np.ndarray, np.ndarray]:
    delta = se3_log(prior.anchor_pose, knot)
    weight = candidate_factor(prior.information)
    jacobian = np.zeros((6, 6))
    for axis in range(6):
        step = np.eye(6)[axis] * 1e-7
        jacobian[:, axis] = (se3_log(prior.anchor_pose, se3_exp(knot, step)) - delta) / 1e-7
    return weight @ delta, weight @ jacobian


def clone(function: Any, **overrides: Any) -> Any:
    return types.FunctionType(function.__code__, dict(vars(sparse), **overrides),
                              argdefs=function.__defaults__)


def score(problem: Any, knots: list[SE3]) -> float:
    return sparse._objective(problem,
        [(0, item.point_body_m) for item in problem.point_measurements], [], [], [], [],
        knots, np.zeros(3), np.zeros(3), 0.0, np.zeros(3), np.zeros(3), np.ones(3), np.ones(3))


def finite_gradient(problem: Any, knots: list[SE3], step: float) -> np.ndarray:
    gradient = np.zeros(6 * len(knots))
    for index in range(len(knots)):
        for axis in range(6):
            plus, minus = list(knots), list(knots)
            delta = np.eye(6)[axis] * step
            plus[index], minus[index] = se3_exp(knots[index], delta), se3_exp(knots[index], -delta)
            gradient[6 * index + axis] = (score(problem, plus) - score(problem, minus)) / (2 * step)
    return gradient


def relative_error(a: np.ndarray | float, b: np.ndarray | float) -> float:
    return float(np.max(np.abs(np.asarray(a) - np.asarray(b))) / max(1.0, np.max(np.abs(b))))


class Captured(Exception):
    pass


class Envelope(BaseModel):
    model_config = ConfigDict(extra="forbid")
    payload: dict[str, Any]
    provenance: dict[str, Any]


class Outcome(BaseModel):
    model_config = ConfigDict(extra="forbid")
    case: str
    implementation: str
    initial_objective: float
    initial_residual_squared: float
    initial_cost_relative_error: float
    initial_gradient_relative_error: float
    independent_fd_halfstep_disagreement: float
    sparse_dense_gradient_relative_error: float
    sparse_dense_normal_relative_error: float
    fit_status: str
    iterations: int
    final_objective: float
    reported_gradient_norm: float
    final_declared_gradient_inf: float
    provenance: dict[str, Any]


def main() -> None:
    assert subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT, text=True).strip() == SOURCE
    original = {p.relative_to(ROOT).as_posix(): sha(p) for p in (ROOT / "src/calibrex").rglob("*.py")}
    upper = np.eye(6)
    upper[0, 1], upper[1, 2], upper[3, 4], upper[2, 5] = .6, -.3, .2, .4
    low_rank = np.array([[1., 2., 0., 1., 0., 0.], [0., 1., -1., 0., 1., 0.]])
    information = {
        "diagonal_spd_control": np.diag([2., 3., 5., 7., 11., 13.]),
        "correlated_spd": upper.T @ upper,
        "correlated_rank_two_psd": low_rank.T @ low_rank,
        "zero_information": np.zeros((6, 6)),
        "small_correlated_spd": 1e-12 * (upper.T @ upper),
        "large_correlated_spd": 1e12 * (upper.T @ upper),
    }
    origin = {"source_commit": SOURCE, "created_at": datetime.now(timezone.utc).isoformat(),
        "generator": "outputs/prior-production-screen-v1/probe.py", "launcher_sha256": sha(Path(__file__)),
        "source_files_raw_sha256": original,
        "scope": "Private in-memory prior whitening only; all attempts retained; no accuracy/timing/adoption claim"}
    write("protocol.json", {"case_names": list(information), "max_iterations": 30,
        "central_fd_steps": [1e-6, 5e-7], "candidate_gradient_relative_gate": 1e-6,
        "nullspace": "No positive floor; tiny negative eigensolver roundoff projected only in PSD fixture",
        "public_contract_unresolved": "Symmetry, finite/shape checks and negative tolerance still require a separate decision",
        "provenance": origin})
    write("envelope.schema.json", Envelope.model_json_schema())
    write("outcome.schema.json", Outcome.model_json_schema())
    anchor0 = se3_exp(SE3.identity(), np.array([.1, .2, -.1, .1, -.05, .2]))
    anchor1 = se3_exp(SE3.identity(), np.array([.4, -.1, .2, -.1, .15, .05]))
    initial = (se3_exp(anchor0, np.array([.3, -.2, .1, .04, -.02, .03])),
               se3_exp(anchor1, np.array([-.1, .15, -.05, -.03, .04, -.02])))
    targets = (se3_exp(anchor0, np.array([-.2, .1, .15, -.05, .03, .02])), anchor1)
    body_points = ((0., 0., 0.), (.3, 0., 0.), (0., .4, 0.), (0., 0., .5), (-.2, .1, .3))
    points = tuple(sparse.TrajectoryPointMeasurement(
        f"{knot}-{axis}", float(knot), point, targets[knot].transform_point(point), 1.0 + .1 * axis)
        for knot in range(2) for axis, point in enumerate(body_points))
    settings = sparse.ContinuousTimeTrajectoryFitOptions(max_iterations=30)
    outcomes = []
    for name, matrix in information.items():
        problem = sparse.ContinuousTimeTrajectoryFitProblem((0., 1.), initial,
            point_measurements=points, knot_marginalization_priors=(KnotMarginalizationPrior(0, anchor0, matrix),))
        envelope = Envelope(payload={"problem": dataclasses.asdict(problem), "options": dataclasses.asdict(settings)},
                            provenance=dict(origin, case=name))
        write(name + "/input.json", envelope.model_dump())
        numeric = finite_gradient(problem, list(initial), 1e-6)
        half = finite_gradient(problem, list(initial), 5e-7)
        for implementation in ("production", "candidate"):
            block_fn = sparse.marginalization_prior_blocks if implementation == "production" else candidate_blocks
            captured = []

            def assembly(blocks: Any, *, residual_count: int, parameter_count: int) -> Any:
                jacobian = sparse._assemble_sparse_jacobian(blocks, residual_count=residual_count,
                                                            parameter_count=parameter_count)
                frame = inspect.currentframe().f_back
                state = frame.f_locals
                weighted = jacobian.multiply(state["weight_array"][:, None])
                residual = state["weight_array"] * state["residual_array"]
                captured.append(((weighted.T @ weighted).toarray(), np.asarray(weighted.T @ residual).ravel(),
                                 float(residual @ residual)))
                raise Captured

            fit = clone(sparse.fit_continuous_trajectory, marginalization_prior_blocks=block_fn)
            capture_fit = clone(sparse.fit_continuous_trajectory, marginalization_prior_blocks=block_fn,
                                _assemble_sparse_jacobian=assembly)
            try:
                capture_fit(problem, settings)
            except Captured:
                pass
            assert len(captured) == 1
            normal, gradient, residual_squared = captured[0]
            dense = clone(sparse.build_dense_knot_normal_equations, marginalization_prior_blocks=block_fn)
            dense_normal, dense_gradient = dense(problem, list(initial))
            result = fit(problem, settings)
            folder = name + "/" + implementation
            provenance = dict(origin, case=name, implementation=implementation)
            write(folder + "/fit-result.json", {"payload": dataclasses.asdict(result), "provenance": provenance})
            result_schema = TypeAdapter(type(result)).json_schema()
            fit_schema = {"type": "object", "additionalProperties": False,
                "required": ["payload", "provenance"], "properties": {
                    "payload": result_schema, "provenance": {"type": "object"}}}
            # Keep TypeAdapter's definitions at the schema root for internal refs.
            if "$defs" in result_schema:
                fit_schema["$defs"] = result_schema.pop("$defs")
            write(folder + "/fit-result.schema.json", fit_schema)
            jsonschema.validate(json.loads((OUT / folder / "fit-result.json").read_text()), fit_schema)
            declared = score(problem, list(initial))
            row = Outcome(case=name, implementation=implementation,
                initial_objective=declared, initial_residual_squared=residual_squared,
                initial_cost_relative_error=relative_error(residual_squared, declared),
                initial_gradient_relative_error=relative_error(2 * gradient, numeric),
                independent_fd_halfstep_disagreement=relative_error(numeric, half),
                sparse_dense_gradient_relative_error=relative_error(gradient, dense_gradient),
                sparse_dense_normal_relative_error=relative_error(normal, dense_normal),
                fit_status=result.status, iterations=result.iterations, final_objective=result.final_objective,
                reported_gradient_norm=result.gradient_norm,
                final_declared_gradient_inf=float(np.max(np.abs(finite_gradient(problem, list(result.knot_poses), 1e-6)))),
                provenance=provenance)
            write(folder + "/outcome.json", row.model_dump())
            outcomes.append(row.model_dump())
            print(json.dumps({k: row.model_dump()[k] for k in ("case", "implementation", "initial_cost_relative_error",
                "initial_gradient_relative_error", "fit_status", "iterations", "final_objective", "final_declared_gradient_inf")}), flush=True)
    assert original == {p.relative_to(ROOT).as_posix(): sha(p) for p in (ROOT / "src/calibrex").rglob("*.py")}
    write("outcomes.json", {"payload": {"outcomes": outcomes, "source_unchanged": True}, "provenance": origin})
    for row in outcomes:
        jsonschema.validate(row, Outcome.model_json_schema())
    print("All twelve private attempts retained; production source unchanged", flush=True)


if __name__ == "__main__":
    main()
