#!/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

    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)
    for lo in range(0, 3200, 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 = []
        for stride in (1, 2, 4):
            selected = np.unique(np.r_[ids[::stride], ids[-1]])
            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,
                )
            )
        arc = (
            integrate(t[ids], a["arc_estimates"][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]
        )
        records.append(
            dict(
                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,
        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 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()