"""Frozen matched SRT/MRT DFG comparison. GPU execution requires explicit root admission."""

import argparse
import csv
import hashlib
import json
import math
import os
from pathlib import Path
import sys
import time
from tempfile import TemporaryDirectory

import numpy as np

HERE = Path(__file__).resolve().parent
DEFAULT_SOURCE = HERE.parents[1]
BOUNDS = (40, 24, 320, 140)
BASE_BLOCK = (2, 881)  # Exact 166 x 881 domain; no geometric padding.
FINE_BLOCK = (8, 8)  # 240 x 568 including four fine halo cells.
SPEED = 0.01


def helpers(source):
    sys.path.insert(0, str(source / "src"))
    sys.path.insert(0, str(source / "tests/validation"))
    import CelerisLab

    if (
        Path(CelerisLab.__file__).resolve()
        != source.resolve() / "src/CelerisLab/__init__.py"
    ):
        raise RuntimeError("Wrong CelerisLab import")
    import re100_helpers as h
    import periodic

    if h.ROOT.resolve() != source.resolve():
        raise RuntimeError("Helper module loaded from wrong source")
    return h, periodic


def config(h, diameter, collision="SRT"):
    cfg, body, p = h.make_config(diameter, SPEED, collision, "neq_extrap", "zou_he_local")
    cfg["method"]["outlet"]["srt_neq_damp"] = 0.5
    return cfg, body, p


def plan(h, collision="SRT"):
    _, _, c = config(h, 40, collision)
    _, _, f = config(h, 80, collision)
    x0, y0, x1, y1 = BOUNDS
    fine_active = 4 * (x1 - x0) * (y1 - y0)
    coarse_uncovered = c["nx"] * c["ny"] - fine_active // 4
    return dict(
        status="prepared_not_run",
        collision=collision,
        root_gpu_gates_required=True,
        root=c,
        uniform=f,
        bounds_root_edge=list(BOUNDS),
        physical_patch_m=[
            (x0 - 0.5) * c["dx"],
            (y0 - 1) * c["dx"],
            (x1 - 0.5) * c["dx"],
            (y1 - 1) * c["dx"],
        ],
        body_root_edge=[80.5, 81.0],
        radius_root=20.0,
        root_block=list(BASE_BLOCK),
        fine_block=list(FINE_BLOCK),
        root_wall_edges=[1.0, 165.0],
        root_inlet_outlet_edges=[0.5, 880.5],
        physical_domain_m=[2.2, 0.41],
        physical_center_m=[0.2, 0.2],
        physical_diameter_m=0.1,
        physical_mean_speed_m_s=1.0,
        physical_viscosity_m2_s=0.001,
        finest_dx_m=f["dx"],
        finest_dt_s=f["dt"],
        approximate_leaf_cell_ratio=(coarse_uncovered + fine_active)
        / (f["nx"] * f["ny"]),
        approximate_collision_work_ratio=(coarse_uncovered + 2 * fine_active)
        / (2 * f["nx"] * f["ny"]),
        cost_caveat="Geometry-only estimates include ghosts/solids; exclude block halos, transfers, observation and launch overhead. Not measured speedup.",
        sampling_offset="For origin (38,22), fine physical x=(j+75.5)*dx80 versus uniform i*dx80: half-fine-cell x offset. Fine physical y=(k+42.5)*dx80 equals uniform (l-.5)*dx80 at l=k+43: y samples align. Body, wall and source planes coincide.",
        outlet=(
            "SRT pressure NEQ, all nine fneq damped by .5, no backflow clamp; noncanonical DFG outlet"
            if collision == "SRT" else
            "MRT pressure NEQ, unknowns q2/6/8 extrapolated; no backflow clamp; noncanonical DFG outlet"
        ),
        modes={"pilot_tstar": 5.0, "full_tstar": 300.0},
        observation_dense_tstar=0.02,
    )


def macros(g):
    rho = np.sum(g, axis=0, dtype=np.float64)
    ux = (g[1] - g[2] + g[5] - g[6] + g[7] - g[8]) / rho
    uy = (g[3] - g[4] + g[5] - g[6] - g[7] + g[8]) / rho
    return rho, ux, uy


def leaf_mass(arrays, dx, masks):
    """Finite public leaves only; covered coarse, solid and halos are excluded."""
    if any(not np.isfinite(g[:, mask]).all() for g, mask in zip(arrays, masks)):
        raise ValueError("Nonfinite physical leaf population")
    return sum(
        float(np.sum(np.sum(g, axis=0, dtype=np.float64)[mask], dtype=np.float64))
        * (dx / 2**level) ** 2
        for level, (g, mask) in enumerate(zip(arrays, masks))
    )


def initial_state(shape, origin, spacing, height=164.0, speed=SPEED, root=False):
    """Same physical parabolic initial condition as uniform InitTubeFlow_v2."""
    y = origin[1] + spacing * (np.arange(shape[0]) + 0.5)
    u = 6 * speed * ((y - 1) / height) * (1 - (y - 1) / height)
    velocity = np.broadcast_to(u[:, None], shape).copy()
    if root:
        velocity[[0, -1], :] = 0
        velocity[:, [0, -1]] = 0
    cx = np.array([0, 1, -1, 0, 0, 1, -1, 1, -1])[:, None, None]
    w = np.array([4 / 9] + [1 / 9] * 4 + [1 / 36] * 4)[:, None, None]
    cu = cx * velocity
    return w * (1 + 3 * cu + 4.5 * cu * cu - 1.5 * velocity * velocity)


def host_geometry(source, h, collision="SRT"):
    """Production donor/transfer and complete cut-link checks, without CUDA."""
    from CelerisLab.config import load_lbm_config
    from CelerisLab.lbm.refinement import FixedPatch2D, prepare_fine_boundary
    from CelerisLab.lbm.blocks import BlockLayout2D
    from CelerisLab.lbm.channel import (
        validate_channel_config,
        validate_channel_transfers,
    )
    from CelerisLab.lbm.refinement_tables import blocked_base_tables
    from CelerisLab.body.geometry.circle import CircleGeometry
    from CelerisLab.body.objects import SimObject

    cfgdict, _, _ = config(h, 40, collision)
    with TemporaryDirectory() as tmp:
        p = Path(tmp) / "cfg.json"
        p.write_text(json.dumps(cfgdict))
        cfg = load_lbm_config(str(p))
    pair = FixedPatch2D.from_config(cfg, BOUNDS)
    layout = BlockLayout2D.build(cfg.ny, cfg.nx, *BASE_BLOCK)
    validate_channel_config(cfg, [BASE_BLOCK, FINE_BLOCK])
    targets = blocked_base_tables(pair, layout, np.argsort(pair.coalesce_coarse))[1]
    validate_channel_transfers(pair, layout, targets)
    BlockLayout2D.build(*pair.fine_shape, *FINE_BLOCK, periodic=(False, False))
    body = SimObject(
        0, geometry=CircleGeometry(80.5, 81.0, 20.0), center=(80.5, 81.0), radius=20.0
    )
    host = prepare_fine_boundary(cfg, pair, (body,), _shifted_blocked=True)
    return cfg, pair, body, host


def verify_source(source, collision="SRT"):
    if collision not in ("SRT", "MRT"):
        raise ValueError("Matched collision must be SRT or MRT")
    name = "source_manifest.json" if collision == "SRT" else "source_manifest_mrt.json"
    expected = json.loads((HERE / "data/refined_dfg" / name).read_text())
    for name, digest in expected["sha256"].items():
        if hashlib.sha256((source / name).read_bytes()).hexdigest() != digest:
            raise RuntimeError("Source mismatch: " + name)
    return expected


def run_case(a, h, periodic, case):
    from CelerisLab.simulation import Simulation
    from CelerisLab.cuda.context import CudaContext
    from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
    from CelerisLab.lbm.descriptors import FLUID
    from contextlib import ExitStack

    refined = case == "refined"
    cfgdict, bodydict, p = config(h, 40 if refined else 80, a.collision)
    cap = 5.0 if a.mode == "pilot" else 300.0
    max_steps = math.ceil(cap * p["diameter"] / SPEED)
    rows = []
    healthy = True
    step_seconds = 0.0
    started = time.monotonic()
    check_deadline(a.deadline)
    dest = a.output / case
    dest.mkdir()
    with ExitStack() as stack:
        if refined:
            cfg, pair, body, host = host_geometry(a.source, h, a.collision)
            context = stack.enter_context(CudaContext(0))
            backend = stack.enter_context(
                FixedPatchStepper(
                    cfg,
                    BOUNDS,
                    context=context,
                    coarse_boundary="channel",
                    base_block=BASE_BLOCK,
                    fine_block=FINE_BLOCK,
                    fine_objects=(body,),
                    observe_raw_exchange=True,
                )
            )
            from CelerisLab.lbm.channel import channel_masks
            from CelerisLab.lbm.refinement import ACTIVE
            from CelerisLab.lbm.descriptors import SOLID

            masks = [
                channel_masks(pair.coarse_shape)[0] & ~pair.covered,
                (pair.fine_kind == ACTIVE)
                & ((host.flags.reshape(pair.fine_shape) & SOLID) == 0),
            ]
            backend.initialize(
                initial_state(pair.coarse_shape, (0.0, 0.0), 1.0, root=True),
                initial_state(pair.fine_shape, pair.fine_origin, 0.5),
            )
            probe_p = dict(
                p,
                center=[
                    (80.5 - pair.fine_origin[0]) / 0.5 - 0.5,
                    (81.0 - pair.fine_origin[1]) / 0.5 - 0.5,
                ],
                radius=40.0,
            )
        else:
            tmp = Path(stack.enter_context(TemporaryDirectory()))
            (tmp / "lbm.json").write_text(json.dumps(cfgdict))
            (tmp / "body.json").write_text(json.dumps(bodydict))
            sim = stack.enter_context(
                Simulation(str(tmp / "lbm.json"), str(tmp / "body.json"), device_id=0)
            )
            sim.initialize()
            probe_p = p
            fluid = (sim.field.flag.reshape(p["ny"], p["nx"]) & FLUID) != 0
        startup = time.monotonic() - started
        check_deadline(a.deadline)
        begin = 0
        with (dest / "history.csv").open("w") as out:
            writer = None
            while begin < max_steps:
                check_deadline(a.deadline)
                width = (
                    0.02
                    if a.mode == "pilot" or begin * SPEED / p["diameter"] >= 240
                    else 0.5
                )
                n = min(max(1, round(width * p["diameter"] / SPEED)), max_steps - begin)
                impulse = np.zeros(2)
                tick = time.monotonic()
                if refined:
                    for _ in range(n):
                        check_deadline(a.deadline)
                        backend.step(1)
                        exchange = backend.read_raw_exchange(
                            0, level=1, window="last_base_step"
                        )
                        if (
                            exchange["units"] != "base_lattice_per_unit_depth"
                            or exchange["sample_count"] != 2
                            or exchange["time_end"] - exchange["time_start"] != 1
                        ):
                            raise RuntimeError("Unexpected force observation contract")
                        impulse += exchange["total"]["impulse"]
                        check_deadline(a.deadline)
                else:
                    done = 0
                    while done < n:
                        check_deadline(a.deadline)
                        batch = min(100, n - done)
                        sim.run(batch)
                        impulse += sim.read_force(0).astype(np.float64) * batch
                        done += batch
                        check_deadline(a.deadline)
                step_seconds += time.monotonic() - tick
                if refined:
                    arrays = backend.download()
                    fields = [macros(g) for g in arrays]
                    root_rho, root_ux, _ = fields[0]
                    rho, ux, uy = fields[1]
                    fluid = masks[1]
                    if any(
                        not np.isfinite(g[:, mask]).all()
                        for g, mask in zip(arrays, masks)
                    ):
                        healthy = False
                        break
                    mass = leaf_mass(arrays, p["dx"], masks)
                    vals = [
                        tuple(v[mask] for v in field)
                        for mask, field in zip(masks, fields)
                    ]
                else:
                    m = sim.get_macroscopic()
                    rho, ux, uy = [
                        np.asarray(m[k], np.float64).reshape(p["ny"], p["nx"])
                        for k in ("rho", "ux", "uy")
                    ]
                    root_rho, root_ux = rho, ux
                    mass = float(rho[fluid].sum()) * p["dx"] ** 2
                    vals = [(rho[fluid], ux[fluid], uy[fluid])]
                healthy = (
                    all(np.isfinite(v).all() for vs in vals for v in vs)
                    and np.isfinite(impulse).all()
                )
                if not healthy:
                    break
                den = SPEED**2 * p["diameter"] / 2
                row = dict(
                    step=begin + n,
                    tstar=(begin + n) * SPEED / p["diameter"],
                    tstar_force=(begin + n / 2) * SPEED / p["diameter"],
                    cd=float(impulse[0] / n / den),
                    cl=float(impulse[1] / n / den),
                    mass=mass,
                    rho_min=float(min(vs[0].min() for vs in vals)),
                    rho_max=float(max(vs[0].max() for vs in vals)),
                    max_mach=float(
                        max(np.sqrt(3 * (vs[1] ** 2 + vs[2] ** 2)).max() for vs in vals)
                    ),
                    flux_in=float(np.sum(root_rho[1:-1, 2] * root_ux[1:-1, 2]))
                    * p["dx"]
                    / SPEED,
                    flux_out=float(np.sum(root_rho[1:-1, -3] * root_ux[1:-1, -3]))
                    * p["dx"]
                    / SPEED,
                    **h.pressure_probes(rho, fluid, probe_p),
                )
                check_deadline(a.deadline)
                ensure_finite(row)
                rows.append(row)
                if writer is None:
                    writer = csv.DictWriter(out, fieldnames=row)
                    writer.writeheader()
                writer.writerow(row)
                out.flush()
                healthy = (
                    row["max_mach"] <= 0.2
                    and row["rho_min"] >= 0.5
                    and row["rho_max"] <= 1.5
                )
                begin += n
                if not healthy:
                    break
                if len(rows) == 1:
                    ready = dict(
                        case=case,
                        collision=a.collision,
                        **a.provenance,
                        mode=a.mode,
                        step=begin,
                        tstar=row["tstar"],
                        runner_sha256=runner_hash(),
                        all_runtime_and_observation_paths_exercised=True,
                    )
                    (dest / "ready.json").write_text(
                        json.dumps(ready, indent=2, allow_nan=False)
                    )
                    print("READY " + json.dumps(ready, allow_nan=False), flush=True)
    result = dict(
        case=case,
        collision=a.collision,
        **a.provenance,
        mode=a.mode,
        history=rows,
        healthy=bool(healthy),
        steps=begin,
        startup_seconds=startup,
        step_seconds=step_seconds,
        wall_seconds=time.monotonic() - started,
        force="raw stationary-body link momentum reaction, per-unit-depth; average over entire observation interval; normalized in each base lattice",
        mass_units="rho_reference * m^2 per unit depth",
        flux_units="rho_reference * m^2/s per unit depth",
        pressure="postcollision end-sample; physical pressure difference / (rho_reference * 1 m/s squared)",
        suite_validated=False,
        **periodic.assess(rows),
    )
    if not healthy:
        result["periodic"] = False
    reference = json.loads(
        (a.source / "tests/validation/data/dfg_re100/reference_cycle.json").read_text()
    )
    if result["cycles"]:
        result["reference_comparison"] = periodic.compare(
            result["cycles"][-1], reference
        )
    result["status"] = (
        "invalid_state"
        if not healthy
        else (
            "pilot_complete"
            if a.mode == "pilot"
            else ("periodic" if result["periodic"] else "incomplete_periodicity")
        )
    )
    (dest / "summary.json").write_text(json.dumps(result, indent=2, allow_nan=False))
    return result


def runner_hash():
    return hashlib.sha256(Path(__file__).read_bytes()).hexdigest()


def ensure_finite(value):
    if isinstance(value, dict):
        for item in value.values():
            ensure_finite(item)
    elif isinstance(value, list):
        for item in value:
            ensure_finite(item)
    elif isinstance(value, (float, int)) and not isinstance(value, bool):
        if not math.isfinite(value):
            raise ValueError("Nonfinite JSON value")


def read_finite_json(path):
    def reject(value):
        raise ValueError("Nonfinite JSON constant: " + value)

    value = json.loads(path.read_text(), parse_constant=reject)
    ensure_finite(value)
    return value


def check_deadline(deadline):
    if time.monotonic() >= deadline:
        raise TimeoutError(
            "Campaign wall-time ceiling reached; no continuation or acceptance"
        )


def verify_pilot(path, manifest, collision="SRT"):
    if path is None:
        raise ValueError(
            "Full comparison requires --pilot-results from a completed matched pilot"
        )
    previous = read_finite_json(path / "plan.json")
    if previous.get("collision") != collision:
        raise ValueError("Pilot collision differs")
    if previous["source_manifest"] != manifest or previous["bounds_root_edge"] != list(
        BOUNDS
    ):
        raise ValueError("Pilot source or frozen geometry differs")
    if previous.get("runner_sha256") != runner_hash():
        raise ValueError("Pilot runner differs")
    calibrated_seconds = 0.0
    for case, diameter in (("uniform", 80), ("refined", 40)):
        summary = read_finite_json(path / case / "summary.json")
        if (
            summary.get("collision") != collision
            or summary.get("case") != case
            or summary.get("mode") != "pilot"
            or summary.get("status") != "pilot_complete"
            or summary.get("healthy") is not True
        ):
            raise ValueError("Mislabeled or unhealthy pilot receipt")
        rows = summary.get("history")
        interval = round(0.02 * diameter / SPEED)
        total = round(5 * diameter / SPEED)
        if (
            not isinstance(rows, list)
            or len(rows) != 250
            or summary.get("steps") != total
        ):
            raise ValueError("Incomplete pilot history or step count")
        for i, row in enumerate(rows):
            step = (i + 1) * interval
            if type(row.get("step")) is not int or row["step"] != step:
                raise ValueError("Nonmonotone or incomplete pilot step history")
            for key, expected in (
                ("tstar", step * SPEED / diameter),
                ("tstar_force", (step - interval / 2) * SPEED / diameter),
            ):
                value = row.get(key)
                if type(value) not in (float, int) or not math.isclose(
                    value, expected, rel_tol=0, abs_tol=1e-12
                ):
                    raise ValueError("Incorrect pilot observation time")
            for key in (
                "cd",
                "cl",
                "dp",
                "mass",
                "rho_min",
                "rho_max",
                "max_mach",
                "flux_in",
                "flux_out",
                "rho_front",
                "rho_back",
            ):
                if type(row.get(key)) not in (int, float):
                    raise ValueError("Missing numeric pilot observation")
            if not (
                0 < row["mass"]
                and 0.5 <= row["rho_min"] <= row["rho_max"] <= 1.5
                and 0 <= row["max_mach"] <= 0.2
            ):
                raise ValueError("Unhealthy pilot observations")
        ready = read_finite_json(path / case / "ready.json")
        if (
            ready.get("collision") != collision
            or ready.get("case") != case
            or ready.get("mode") != "pilot"
            or ready.get("step") != interval
            or ready.get("runner_sha256") != runner_hash()
            or ready.get("all_runtime_and_observation_paths_exercised") is not True
        ):
            raise ValueError("Invalid pilot READY receipt")
        startup = summary.get("startup_seconds")
        wall = summary.get("wall_seconds")
        stepping = summary.get("step_seconds")
        if (
            any(type(v) not in (float, int) for v in (startup, wall, stepping))
            or not 0 <= startup < wall
            or not 0 < stepping <= wall
        ):
            raise ValueError("Invalid pilot calibration timings")
        # Full physical duration is 60 times pilot; dense pilot sampling makes
        # this conservative for the sparser early full observations. Not a guarantee.
        calibrated_seconds += startup + 60 * (wall - startup)
    return 3 * calibrated_seconds


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--collision", choices=["SRT", "MRT"], default="SRT")
    parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--mode", choices=["plan", "pilot", "full"], default="plan")
    parser.add_argument("--gpu-gates-approved", action="store_true")
    parser.add_argument("--pilot-results", type=Path)
    parser.add_argument(
        "--max-wall-seconds",
        type=float,
        help="Required explicit campaign wall-time ceiling for GPU execution",
    )
    a = parser.parse_args()
    a.source = a.source.resolve()
    manifest = verify_source(a.source, a.collision)
    h, periodic = helpers(a.source)
    from source_provenance import checkout_provenance

    a.provenance = checkout_provenance(a.source)
    calibration = (
        verify_pilot(a.pilot_results, manifest, a.collision)
        if a.mode == "full" else None
    )
    if a.mode != "plan" and (
        a.max_wall_seconds is None
        or not math.isfinite(a.max_wall_seconds)
        or a.max_wall_seconds <= 0
    ):
        raise ValueError("GPU execution requires finite positive --max-wall-seconds")
    budget = (
        min(a.max_wall_seconds, calibration)
        if calibration is not None
        else a.max_wall_seconds
    )
    if a.output.exists() and any(a.output.iterdir()):
        raise ValueError("Output must be empty/new")
    a.output.mkdir(parents=True, exist_ok=True)
    specification = plan(h, a.collision)
    specification.update(
        source_manifest=manifest,
        **a.provenance,
        gates=periodic.GATES,
        reference_gates=periodic.REFERENCE_GATES,
        runner_sha256=runner_hash(),
        wall_time_budget_seconds=budget,
        full_pilot_calibration_seconds=calibration,
    )
    host_geometry(a.source, h, a.collision)
    (a.output / "plan.json").write_text(json.dumps(specification, indent=2))
    if a.mode == "plan":
        return 0
    if not a.gpu_gates_approved or not os.environ.get("CUDA_VISIBLE_DEVICES"):
        raise RuntimeError(
            "GPU execution blocked: root GPU gates approval and explicit device selection required"
        )
    results = []
    a.deadline = time.monotonic() + budget
    for case in ("uniform", "refined"):
        try:
            result = run_case(a, h, periodic, case)
        except TimeoutError as exc:
            failure = dict(
                status="wall_time_limit",
                collision=a.collision,
                **a.provenance,
                case=case,
                matched_case_pass=False,
                error=str(exc),
            )
            (a.output / "comparison.json").write_text(
                json.dumps(failure, indent=2, allow_nan=False)
            )
            print(json.dumps(failure, allow_nan=False), flush=True)
            return 3
        results.append(result)
        if not result["healthy"]:
            break
    comparison = dict(
        collision=a.collision,
        **a.provenance,
        status="pilot_only" if a.mode == "pilot" else "incomplete",
        matched_case_pass=False,
    )
    if len(results) == 2 and all(r["healthy"] and r["periodic"] for r in results):
        comparison.update(
            status="periodic_comparison",
            refined_vs_uniform=periodic.compare(
                results[1]["cycles"][-1], results[0]["cycles"][-1]
            ),
        )
        comparison["matched_case_pass"] = comparison["refined_vs_uniform"][
            "reference_agreement"
        ]
    (a.output / "comparison.json").write_text(
        json.dumps(comparison, indent=2, allow_nan=False)
    )
    return (
        0
        if len(results) == 2
        and all(r["healthy"] for r in results)
        and (a.mode == "pilot" or comparison["matched_case_pass"])
        else 2
    )


if __name__ == "__main__":
    raise SystemExit(main())