#!/usr/bin/env python3
"""Offline common-unit/sampling sensitivity comparison; never grants physics acceptance."""

import argparse
import hashlib
import json
from pathlib import Path
import numpy as np


def integrate(t, y):
    return np.sum(
        np.diff(t).reshape((-1,) + (1,) * (y.ndim - 1)) * (y[1:] + y[:-1]) * 0.5, axis=0
    )


def analyze(directory):
    meta_path = next(Path(directory).glob("samples_s*.json"))
    meta = json.loads(meta_path.read_text())
    c = meta["case"]
    s = c["scale"]
    sample_path = meta_path.parent / meta["sample_file"]
    checksum_verified = False
    if "sample_sha256" in meta:
        if (
            hashlib.sha256(sample_path.read_bytes()).hexdigest()
            != meta["sample_sha256"]
        ):
            raise ValueError("Sample file hash mismatch")
        checksum_verified = True
    a = np.load(sample_path)
    t = a["t_reference"]
    if not len(t) or t[0] != 0 or np.any(np.diff(t) <= 0):
        raise ValueError("Sample reference times must increase strictly from zero")
    from jet_cv_measurements import sample_steps, old_sample_steps

    complete_schedule = np.array_equal(t, sample_steps(c) / s**2)
    records = []
    # Input retains root units; convert before all budgets/comparisons.
    inv = a["cv_inventory"] * np.array([1 / s**2, 1 / s, 1 / s])
    flux = a["cv_flux"] * np.array([1.0, s, s])
    raw = a["raw_cumulative"] * np.array([1 / s, 1 / s, 1.0, 1 / s**2])
    raw[:, :, :, 2] /= (
        s**2
    )  # angular impulse reference = root/s^2 (not used in force tests)
    q0 = 2 * np.pi / 12 * (c["diameter"] / 2) * c["jet_speed"]
    f0 = 0.5 * (c["u"] * s) ** 2 * (c["diameter"] / s)
    dense_root = c.get("dense_root", False)
    old_times = old_sample_steps(c) / s**2
    ray_indices = (
        a["ray_sample_indices"] if "ray_sample_indices" in a else np.arange(len(t))
    )
    ray_times = t[ray_indices]
    if dense_root and complete_schedule and not np.array_equal(ray_times, old_times):
        raise ValueError("Sparse ray records do not match the declared old cadence")
    for lo in range(0, c.get("steps", 3200 * s * s) // (s * s), 100):
        hi = lo + 100
        ids = np.flatnonzero((t >= lo) & (t <= hi))
        if len(ids) < 2 or t[ids[0]] != lo or t[ids[-1]] != hi:
            continue
        first, last = ids[0], ids[-1]
        delta = inv[last] - inv[first]
        ledger = raw[last] - raw[first]
        cv_variants = []
        selections = [np.unique(np.r_[ids[::stride], ids[-1]]) for stride in (1, 2, 4)]
        if dense_root:
            selections.append(ids[np.isin(t[ids], old_times)])
        for selected in selections:
            if t[selected[0]] != lo or t[selected[-1]] != hi:
                raise ValueError(
                    "Every quadrature selection must retain both window endpoints"
                )
            total = integrate(t[selected], flux[selected])
            m = delta[:, :, 0] + total[:, :, 0]
            j = -delta[:, :, 1:] - total[:, :, 1:]
            cv_variants.append(
                np.concatenate(
                    (
                        (m - ledger[:, 2, 3, None])[:, :, None] / (100 * q0),
                        (j - ledger[:, 2, None, :2]) / (100 * f0),
                    ),
                    axis=-1,
                )
            )
        ray_ids = np.flatnonzero((ray_times >= lo) & (ray_times <= hi))
        arc = (
            integrate(ray_times[ray_ids], a["arc_estimates"][ray_ids]) / 100
        )  # Q/M rate conversion is unity
        # Each estimator, aperture, Q/M; compare against signed prescribed Q (rho_ref=1 for M only).
        arc_ratio = arc / np.array([q0, -q0])[None, None, :, None]
        global_delta = (
            a["global_mass_reference"][last] - a["global_mass_reference"][first]
        )
        dense_details = {}
        if dense_root:
            absolute = (
                a["raw_absolute_aperture_mass_cumulative"][last]
                - a["raw_absolute_aperture_mass_cumulative"][first]
            ) / s**2
            stock_raw = np.diff(a["global_mass_reference"][ids], axis=0) - np.diff(
                raw[ids, :, 2, 3], axis=0
            )
            partition = ledger[:, 2] - ledger[:, 0] - ledger[:, 1]
            dense_details = dict(
                quadrature_names=[
                    "root_stride1",
                    "root_stride2",
                    "root_stride4",
                    "exact_old_cadence",
                ],
                root_dt_reference=1 / s**2,
                old_cadence_change_from_dense=(
                    cv_variants[3] - cv_variants[0]
                ).tolist(),
                root_temporal_sensitivity_max=float(
                    np.max(abs(np.array(cv_variants)[1:3] - cv_variants[0]))
                ),
                raw_absolute_aperture_net_interval_mass_reference=absolute.tolist(),
                prescribed_unsigned_nominal_mass_reference=float(2 * q0 * 100),
                legacy_raw_gross_field_definition="raw_gross_mass_over_nominal retains abs(window net blowing)+abs(window net suction), divided by ONE nominal aperture throughput; it is not root-interval absolute throughput.",
                raw_absolute_to_prescribed_unsigned_nominal=(
                    absolute / (2 * q0 * 100)
                ).tolist(),
                global_stock_raw_balance_max_abs_root_interval_reference=np.max(
                    abs(stock_raw), axis=0
                ).tolist(),
                global_stock_raw_balance_sum_reference=stock_raw.sum(axis=0).tolist(),
                raw_total_minus_wall_aperture_reference=partition.tolist(),
                raw_aperture_minus_blowing_suction_reference=(
                    ledger[:, 1] - ledger[:, 3] - ledger[:, 4]
                ).tolist(),
                measured_ray_absolute_volume_mass_reference=integrate(
                    ray_times[ray_ids], a["arc_absolute_estimates"][ray_ids]
                ).tolist(),
                measured_ray_cadence="Sparse old cadence; estimator-dependent, not dense-root or exact wall throughput",
            )
        records.append(
            dict(
                **dense_details,
                t_reference=[lo, hi],
                startup=hi <= 800,
                normalized_cv_residual=np.array(cv_variants).tolist(),
                temporal_sensitivity_max=float(
                    np.max(abs(np.array(cv_variants)[1:] - cv_variants[0]))
                ),
                arc_ratio_to_command_reference_density=arc_ratio.tolist(),
                global_mass_residual_over_nominal=(
                    (global_delta - ledger[:, 2, 3]) / (100 * q0)
                ).tolist(),
                raw_reaction_coefficient=(ledger[:, 2, :2] / (100 * f0)).tolist(),
                raw_gross_mass_over_nominal=(
                    (abs(ledger[:, 3, 3]) + abs(ledger[:, 4, 3])) / (100 * q0)
                ).tolist(),
                raw_wall_mass_over_nominal=(ledger[:, 0, 3] / (100 * q0)).tolist(),
            )
        )
    return dict(
        scale=s,
        dense_root=dense_root,
        previous_temporal_failure_retained=(
            "Prior continuation stride2 changes exceeded baseline residual in all48 audited comparisons; this diagnostic does not retroactively change those gates."
            if dense_root
            else None
        ),
        complete_schedule=bool(complete_schedule),
        sample_checksum_verified=checksum_verified,
        windows=records,
        sampled_density_range=[
            float(a["level_stats"][..., 0].min()),
            float(a["level_stats"][..., 1].max()),
        ],
        sampled_max_speed_reference=float(a["level_stats"][..., 3].max()),
        global_mass_initial_final=a["global_mass_reference"][[0, -1]].tolist(),
        max_sampled_fine_velocity_error=float(a["velocity_error"].max()),
        physics_validation_pass=False,
    )


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("directories", nargs="+", type=Path)
    p.add_argument("--output", type=Path, required=True)
    a = p.parse_args()
    results = [analyze(d) for d in a.directories]
    payload = dict(
        status="TRANSIENT_DIAGNOSTIC_COMPARISON",
        results=results,
        physics_validation_pass=False,
    )
    if (
        not any(x.get("dense_root", False) for x in results)
        and sorted(x["scale"] for x in results) == [1, 2]
        and all(x["complete_schedule"] for x in results)
    ):
        ordered = sorted(results, key=lambda x: x["scale"])
        summary = []
        for startup in (True, False):
            r = [
                np.array(
                    [
                        w["normalized_cv_residual"][0]
                        for w in item["windows"]
                        if w["startup"] == startup
                    ]
                )
                for item in ordered
            ]
            e = [
                np.array(
                    [
                        w["temporal_sensitivity_max"]
                        for w in item["windows"]
                        if w["startup"] == startup
                    ]
                )
                for item in ordered
            ]
            rms = [np.sqrt(np.mean(x * x, axis=(0, 1, 2))) for x in r]
            unc = max(float(x.max()) for x in e)
            summary.append(
                dict(
                    startup=startup,
                    rms_normalized_residual_mass_Cx_Cy=[v.tolist() for v in rms],
                    max_temporal_sensitivity=unc,
                    reduction_exceeds_twice_temporal_proxy=(
                        rms[0] - rms[1] > 2 * unc
                    ).tolist(),
                )
            )
        payload["refinement_trend"] = summary
    a.output.write_text(json.dumps(payload, indent=2) + "\n")


if __name__ == "__main__":
    main()