"""Measurement-only continuation diagnostics. Never changes populations or commands."""

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

OFFSETS = np.array([1.5, 2.5, 3.5])
LINEAR = np.array([2.5, -1.5, 0.0])
QUADRATIC = np.array([4.375, -5.25, 1.875])


def sample_steps(c):
    """Same reference times at both scales; dense startup retained."""
    s2 = c["scale"] ** 2
    reference = sorted(set(range(0, 101)) | set(range(100, 3201, 4)))
    return np.array(reference, int) * s2


def factors(scale):
    """Lref=one s=1 root length, Tref=one s=1 root time, rho_ref=1."""
    return dict(
        length=1 / scale,
        time=1 / scale**2,
        velocity=scale,
        mass=1 / scale**2,
        momentum=1 / scale,
        impulse=1 / scale,
        force=scale,
        volume_rate=1.0,
        mass_rate=1.0,
        pressure=scale**2,
    )


def primitives(g, cx, cy):
    f = np.asarray(g, np.float64)
    rho = f.sum(axis=0)
    mom = np.stack([np.einsum("q,qyx->yx", v, f) for v in (cx, cy)])
    return rho, mom / rho


def contour_terms(rho, vel, origin, rectangle, nu, h=0.5):
    """Independent split of total outward momentum flux into advective,p,viscous."""
    x0, y0, x1, y1 = np.rint((np.array(rectangle) - np.tile(origin, 2)) / h).astype(int)
    duy, dux = np.gradient(vel[0], h)
    dvy, dvx = np.gradient(vel[1], h)
    adv = rho * vel[:, None] * vel[None, :]
    pressure = np.zeros_like(adv)
    pressure[0, 0] = pressure[1, 1] = (rho - 1) / 3
    viscous = -rho * nu * np.array([[2 * dux, duy + dvx], [duy + dvx, 2 * dvy]])
    out = []
    for field in (adv, pressure, viscous):
        flux = np.zeros(2)
        for i, sign in ((x0, -1), (x1, 1)):
            flux += (
                sign
                * 0.5
                * (field[:, 0, y0:y1, i - 1] + field[:, 0, y0:y1, i]).sum(axis=1)
                * h
            )
        for j, sign in ((y0, -1), (y1, 1)):
            flux += (
                sign
                * 0.5
                * (field[:, 1, j - 1, x0:x1] + field[:, 1, j, x0:x1]).sum(axis=1)
                * h
            )
        out.append(flux)
    return np.array(out)


def ray_geometry(fluid, origin, center, radius, h=0.5):
    z, w = np.polynomial.legendre.leggauss(32)
    theta = np.array([np.pi / 2, 3 * np.pi / 2])[:, None] + np.pi / 12 * z
    normal = np.stack((np.cos(theta), np.sin(theta)), axis=-1)
    positions = np.asarray(center) + normal[:, :, None, :] * (
        radius + h * OFFSETS[None, None, :, None]
    )
    local = (positions - origin) / h - 0.5
    ix = np.floor(local[..., 0]).astype(int)
    iy = np.floor(local[..., 1]).astype(int)
    a = local[..., 0] - ix
    b = local[..., 1] - iy
    if not (
        np.all(ix >= 0)
        and np.all(iy >= 0)
        and np.all(ix + 1 < fluid.shape[1])
        and np.all(iy + 1 < fluid.shape[0])
    ):
        raise ValueError("Ray interpolation leaves allocated grid")
    for dy, dx in ((0, 0), (0, 1), (1, 0), (1, 1)):
        if not fluid[iy + dy, ix + dx].all():
            raise ValueError("Ray stencil must be active fluid")
    return dict(
        ix=ix,
        iy=iy,
        a=a,
        b=b,
        normal=normal,
        ds=w * radius * np.pi / 12,
        positions=positions,
        origin=np.asarray(origin),
        h=h,
    )


def reconstruct_aperture(rho, vel, geo):
    field = np.concatenate((rho[None], vel), axis=0)
    x, y, a, b = geo["ix"], geo["iy"], geo["a"], geo["b"]
    ray = (
        field[:, y, x] * (1 - a) * (1 - b)
        + field[:, y, x + 1] * a * (1 - b)
        + field[:, y + 1, x] * (1 - a) * b
        + field[:, y + 1, x + 1] * a * b
    )
    estimates = []
    for coeff in (LINEAR, QUADRATIC):
        wall = np.einsum("cajk,k->caj", ray, coeff)
        un = np.einsum("caj,ajc->aj", wall[1:], geo["normal"])
        q = np.sum(un * geo["ds"], axis=1)
        mass = np.sum(wall[0] * un * geo["ds"], axis=1)
        estimates.append(np.stack((q, mass), axis=-1))
    return ray, np.array(estimates)


class MeasurementRecorder:
    """Retain reconstructible per-sample terms and sparse full primitive snapshots."""

    def __init__(self, c, output, fluid_u, leaves, origin_a, cx, cy):
        self.c = c
        self.output = Path(output)
        self.fluid_u = fluid_u
        self.leaves = leaves
        self.origin_a = np.asarray(origin_a)
        self.cx = cx
        self.cy = cy
        self.records = []
        self.rays = [
            ray_geometry(fluid_u, (0.0, 0.0), np.array(c["center"]), c["diameter"] / 2),
            ray_geometry(leaves[1], origin_a, np.array(c["center"]), c["diameter"] / 2),
        ]
        self.snapshots = []

    def capture(self, step, ug, levels, cv_data, nearest_flux, velocity_error):
        c = self.c
        s = c["scale"]
        record = dict(
            step=step,
            t_reference=step / s**2,
            t_star=step * c["u"] / c["diameter"],
            velocity_error=velocity_error,
        )
        arrays = [ug, *levels]
        masks = [self.fluid_u, *self.leaves]
        hs = [0.5, 1.0, 0.5]
        primitive = [primitives(g, self.cx, self.cy) for g in arrays]
        stats = []
        for (rho, vel), mask, h in zip(primitive, masks, hs):
            if (
                not np.isfinite(rho[mask]).all()
                or rho[mask].min() <= 0
                or not np.isfinite(vel[:, mask]).all()
            ):
                raise RuntimeError(
                    "Invalid density/velocity on an active composite leaf"
                )
            stats.append(
                [
                    rho[mask].min(),
                    rho[mask].max(),
                    np.max(abs(rho[mask] - 1)),
                    np.linalg.norm(vel[:, mask], axis=0).max() * s,
                    rho[mask].sum() * h * h / s**2,
                ]
            )
        record["level_stats"] = stats  # uniform, AMR root leaf, AMR fine leaf
        record["global_mass_reference"] = [stats[0][-1], stats[1][-1] + stats[2][-1]]
        record["cv_inventory"] = []
        record["cv_flux"] = []
        record["momentum_flux_terms"] = []
        record["arc_ray_values"] = []
        record["arc_estimates"] = []
        record["nearest_mass_flux"] = [
            nearest_flux[k] for k in ("uniform_adapter", "amr_api")
        ]
        for idx, (name, p, origin) in enumerate(
            [
                ("uniform_adapter", primitive[0], (0.0, 0.0)),
                ("amr_api", primitive[2], self.origin_a),
            ]
        ):
            rho, vel = p
            record["cv_inventory"].append([x[0] for x in cv_data[name]])
            record["cv_flux"].append([x[1] for x in cv_data[name]])
            terms = np.array(
                [contour_terms(rho, vel, origin, r, c["nu"]) for r in c["contours"]]
            )
            np.testing.assert_allclose(
                terms.sum(axis=1),
                np.array(record["cv_flux"][-1])[:, 1:],
                atol=2e-12,
                rtol=2e-10,
            )
            record["momentum_flux_terms"].append(terms)
            ray, estimates = reconstruct_aperture(rho, vel, self.rays[idx])
            record["arc_ray_values"].append(ray)
            record["arc_estimates"].append(estimates)
        self.records.append(record)
        if step // s**2 in (0, 100, 800, 1600, 3200) and step % s**2 == 0:
            fields = {}
            for name, (rho, vel), mask, h, origin in zip(
                ("uniform", "amr_root", "amr_fine"),
                primitive,
                masks,
                hs,
                ((0.0, 0.0), (0.0, 0.0), self.origin_a),
            ):
                fields[name + "_rho"] = np.where(mask, rho, np.nan).astype(np.float32)
                fields[name + "_velocity_reference"] = np.where(
                    mask[None], vel * s, np.nan
                ).astype(np.float32)
                fields[name + "_origin_reference"] = np.asarray(origin) / s
                fields[name + "_h_reference"] = np.array(h / s)
            file = self.output / f"fields_s{s}_tref{step//s**2}.npz"
            np.savez_compressed(file, **fields)
            self.snapshots.append(
                dict(
                    file=file.name,
                    sha256=hashlib.sha256(file.read_bytes()).hexdigest(),
                    bytes=file.stat().st_size,
                    t_reference=step / s**2,
                )
            )

    def flush(self):
        path = self.output / f"samples_s{self.c['scale']}.npz"
        arrays = {k: np.array([r[k] for r in self.records]) for k in self.records[0]}
        for i, g in enumerate(self.rays):
            for k in ("normal", "ds", "positions"):
                arrays[f"ray_{i}_{k}"] = g[k]
        arrays["ray_offsets_fine_cells"] = OFFSETS
        arrays["linear_coefficients"] = LINEAR
        arrays["quadratic_coefficients"] = QUADRATIC
        np.savez_compressed(path, **arrays)
        metadata = dict(
            schema=1,
            case=self.c,
            reference_conversion=factors(self.c["scale"]),
            sample_file=path.name,
            sample_sha256=hashlib.sha256(path.read_bytes()).hexdigest(),
            sample_count=len(self.records),
            snapshots=self.snapshots,
            level_stats_columns=[
                "rho_min",
                "rho_max",
                "max_abs_rho_minus_1",
                "max_speed_reference",
                "mass_reference",
            ],
            raw_partition_order=["wall", "aperture", "total", "blowing", "suction"],
            raw_columns=["Jx", "Jy", "angular_impulse", "mass_into_fluid"],
            backend_order=["uniform_adapter", "amr_api"],
            level_order=["uniform", "amr_root_leaf", "amr_fine_leaf"],
            momentum_term_order=["advection", "pressure", "negative_viscous_traction"],
            arc_estimate_axes=[
                "sample",
                "backend",
                "linear_or_quadratic",
                "blowing_or_suction",
                "Q_or_M",
            ],
            arc_ray_axes=[
                "sample",
                "backend",
                "rho_ux_uy",
                "aperture",
                "gauss_point",
                "normal_offset",
            ],
            warning="Extrapolation and nearest-node values are estimator-dependent, not exact wall values. Fields/ray arrays use root-lattice units except explicitly named reference arrays.",
        )
        (self.output / f"samples_s{self.c['scale']}.json").write_text(
            json.dumps(metadata, indent=2) + "\n"
        )