"""One bounded GPU case. No production edits, plots, checkpoints or field dumps."""

import argparse
import hashlib
import json
import os
from pathlib import Path
import subprocess
import sys
import time

import numpy as np
from oracle import U, TIMES, LIMITS, initial, metrics, rectangles


def provenance(repo):
    git = lambda *a: subprocess.check_output(
        ["git", "-C", str(repo), *a], text=True
    ).strip()
    source_hash = hashlib.sha256()
    for path in sorted((repo / "src").rglob("*")):
        if path.is_file() and "__pycache__" not in path.parts and path.suffix != ".pyc":
            source_hash.update(
                str(path.relative_to(repo)).encode() + b"\0" + path.read_bytes() + b"\0"
            )
    return dict(
        revision=git("rev-parse", "HEAD"),
        dirty=bool(git("status", "--porcelain")),
        source_tree_sha256=source_hash.hexdigest(),
        harness_sha256={
            p.name: hashlib.sha256(p.read_bytes()).hexdigest()
            for p in Path(__file__).parent.glob("*.py")
        },
    )


def config(n, re):
    from CelerisLab.config import LBMConfig

    cfg = LBMConfig(
        nx=n,
        ny=n,
        viscosity=U * n / re,
        velocity=0.0,
        rho=1.0,
        lattice_model="D2Q9",
        collision="SRT",
        streaming="double_buffer",
        store_precision="FP32",
        ddf_shifting=True,
        les_enabled=False,
    )
    cfg.validate()
    return cfg


def workload(masks, base_steps):
    cells = [int(np.count_nonzero(m)) for m in masks]
    hops = [base_steps * 2**level for level in range(len(masks))]
    updates = [count * hop for count, hop in zip(cells, hops)]
    return dict(
        active_cells_by_level=cells,
        level_hops=hops,
        active_collision_updates_by_level=updates,
        active_collision_updates=sum(updates),
        root_intervals=base_steps,
        excludes="ghost streaming, halo copies, transfers, initialization and diagnostics",
    )


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--repo", type=Path, required=True)
    p.add_argument("--re", type=int, choices=[100, 1000], required=True)
    p.add_argument("--n", type=int, choices=[32, 64, 128, 256], required=True)
    p.add_argument("--depth", type=int, choices=[1, 2, 3], required=True)
    p.add_argument("--output", type=Path, required=True)
    args = p.parse_args()
    if not os.environ.get("CUDA_VISIBLE_DEVICES"):
        p.error("Explicit physical CUDA_VISIBLE_DEVICES selection required")
    repo = args.repo.resolve()
    sys.path.insert(0, str(repo / "src"))
    import CelerisLab

    if Path(CelerisLab.__file__).resolve() != repo / "src/CelerisLab/__init__.py":
        raise RuntimeError("Wrong imported source checkout")
    from CelerisLab.cuda.context import CudaContext
    from CelerisLab.lbm.blocks import BlockLayout2D
    from CelerisLab.lbm.block_stepper import PeriodicBlockStepper
    from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
    from CelerisLab.lbm.refinement import ACTIVE
    from CelerisLab.lbm.hierarchy import build_hierarchy, leaf_partition

    cfg = config(args.n, args.re)
    refinements = [] if args.depth == 1 else rectangles(args.n, args.depth)
    h = build_hierarchy(cfg, refinements)
    np.testing.assert_array_equal(leaf_partition(h), np.ones((args.n, args.n)))
    masks = [k == ACTIVE for k in h.kinds]
    report = dict(
        schema=1,
        source=provenance(repo),
        n=args.n,
        re_L=args.re,
        depth=args.depth,
        refinements=refinements,
        U=U,
        mach=U * np.sqrt(3),
        viscosity=cfg.viscosity,
        tau_by_level=[0.5 + 3 * c.viscosity for c in h.configs],
        initialization="physical FP64 equilibrium G at actual cell centers; rho=1+3p",
        cuda_visible_devices=os.environ["CUDA_VISIBLE_DEVICES"],
        device_id=0,
        limits=LIMITS,
        workload=workload(masks, round(TIMES[-1] * args.n / U)),
        rows=[],
        completed=False,
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    save = lambda: args.output.write_text(
        json.dumps(report, indent=2, allow_nan=False) + "\n"
    )
    save()
    started = time.monotonic()
    with CudaContext(0) as ctx:
        if args.depth == 1:
            layout = BlockLayout2D.build(args.n, args.n, 8, 8, periodic=(True, True))
            run = PeriodicBlockStepper(cfg, layout, context=ctx)
        else:
            run = FixedPatchStepper(
                cfg,
                context=ctx,
                coarse_boundary="periodic",
                refinements=refinements,
                level_blocks=[(8, 8)] * args.depth,
            )
        with run:
            run.initialize(
                *[
                    initial(m.shape, level, args.n, cfg.viscosity)
                    for level, m in enumerate(masks)
                ]
            )
            initial_mass = None
            for tstar in TIMES:
                step = round(tstar * args.n / U)
                run.step(step - run.step_count)
                arrays = (run.download(),) if args.depth == 1 else run.download()
                row = metrics(arrays, masks, step, args.n, cfg.viscosity, initial_mass)
                row["step"] = step
                if initial_mass is None:
                    initial_mass = row["mass"]
                report["rows"].append(row)
                save()
                print(
                    json.dumps(
                        dict(
                            event="PROGRESS",
                            n=args.n,
                            re_L=args.re,
                            depth=args.depth,
                            **row,
                        )
                    ),
                    flush=True,
                )
                if step == 0:
                    # Exercise real kernels/halos/transfers on every level before READY.
                    # This is the first measured trajectory step, not discarded warmup.
                    run.step(1)
                    probe = (run.download(),) if args.depth == 1 else run.download()
                    metrics(probe, masks, 1, args.n, cfg.viscosity, initial_mass)
                    report["ready_after_step"] = 1
                    save()
                    print(
                        json.dumps(
                            dict(
                                event="READY",
                                step=1,
                                n=args.n,
                                re_L=args.re,
                                depth=args.depth,
                                workload=report["workload"],
                            )
                        ),
                        flush=True,
                    )
    report.update(
        completed=True,
        seconds=time.monotonic() - started,
        screening_pass=all(
            row[key] <= limit for row in report["rows"] for key, limit in LIMITS.items()
        ),
    )
    save()
    print(
        json.dumps(
            {k: report[k] for k in ("n", "re_L", "depth", "screening_pass", "seconds")}
        ),
        flush=True,
    )
    return 0 if report["screening_pass"] else 1


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