"""Bounded J1 Couette-throughflow validation; gates fixed before first run.

See research/j1-2026-10-12/REVIEW.md. No production audit kernel or extra launch
is added: a temporary test build records actual incoming populations and the
roundoff of the real collision, then is compared with uninstrumented stepping.
"""
import argparse
import hashlib
import json
import os
from pathlib import Path
import sys
from tempfile import TemporaryDirectory
from unittest.mock import patch

import numpy as np
import pycuda.driver as cuda

import CelerisLab
from CelerisLab.cuda import compiler_v2 as compiler
from CelerisLab.simulation import Simulation

ROOT = Path(__file__).resolve().parents[2]
CX = np.array([0,1,-1,0,0,1,-1,1,-1])
CY = np.array([0,0,0,1,-1,1,-1,-1,1])
GATES = dict(u_relative_l2=1e-3, v_max_error=2e-6, flux_relative_error=1e-3,
             mass_offset=1e-3, mass_window_span=6e-6, mass_chunk_change=2e-6,
             u_change=2e-5, v_change=2e-7, budget_relative_residual=2e-7)


def make_sim(v, mode='on_site_velocity', nx=16, top_u=.02, slot_width=0, top_mode='velocity'):
    config = json.loads((ROOT/'src/CelerisLab/configs/config_lbm.json').read_text())
    config['grid'].update(nx=nx, ny=35, nz=1, lattice_model='D2Q9')
    config['physics'].update(viscosity=.1, velocity=top_u, rho=1.)
    config['method'].update(collision='SRT', streaming='double_buffer',
                            store_precision='FP32', ddf_shifting=False,
                            y_wall_bc=mode, wall_normal_velocity=v, wall_slot_width=slot_width,
                            top_wall_bc=top_mode)
    config['method']['les']['enabled'] = False
    with TemporaryDirectory(prefix='celeris-j1-config-') as folder:
        lbm, body = Path(folder)/'lbm.json', Path(folder)/'body.json'
        lbm.write_text(json.dumps(config))
        body.write_text('{"objects":[]}')
        return Simulation(str(lbm), str(body), device_id=0)


def populations(sim):
    return sim.get_ddf().reshape(9, sim.field.ny, sim.field.nx).astype(np.float64)


def observables(f):
    interior = f[:, 1:-1, :]
    rho = interior.sum(axis=0)
    ux = np.einsum('q,qyx->yx', CX, interior)/rho
    uy = np.einsum('q,qyx->yx', CY, interior)/rho
    return rho, ux, uy


AUDIT = r'''
__device__ double j1_incoming[2*NX];
__device__ double j1_collision[NX*NY];
__device__ void j1_collide_audit(float* f, float rho, float ux, float uy) {
    const unsigned x=threadIdx.x+blockIdx.x*blockDim.x, y=blockIdx.y;
    double before=0.0;
    for(int i=0;i<NQ;++i) before+=(double)f[i];
    if(y==1) j1_incoming[x]=(double)f[3]+f[5]+f[8];
    if(y==NY-2) j1_incoming[NX+x]=(double)f[4]+f[6]+f[7];
    collide_dispatch(f,rho,ux,uy);
    double after=0.0;
    for(int i=0;i<NQ;++i) after+=(double)f[i];
    j1_collision[y*NX+x]=after-before;
}
'''


def population_budget(sim, bottom_slot=None):
    """16 real steps, separate incoming/outgoing/collision records, no targets."""
    saved = sim.get_ddf().copy()
    count = sim.stepper.step_count
    sim.run(16)
    plain = sim.get_ddf().copy()
    sim.field.ddf[:] = saved
    sim.field.upload_ddf()
    sim.stepper._step_count = count
    build = compiler.generate_config(sim.lbm_cfg)
    original_step = sim.stepper.step_fn
    try:
        source = Path(build)/'step/one_step_double.cu'
        code = source.read_text().replace('collide_dispatch(f, rho_n, ux, uy);',
                                         'j1_collide_audit(f, rho_n, ux, uy);')
        marker = '__global__ __launch_bounds__(NT, 2)'
        assert code.count(marker) == 1
        source.write_text(code.replace(marker, AUDIT+'\n'+marker))
        audit_hash = hashlib.sha256(source.read_bytes()).hexdigest()
        ptx = compiler.compile_kernel(build_dir=build)
        rows = []
        with sim.ctx.activate():
            module = compiler.load_module(ptx)
            params, size = module.get_global('d_params')
            cuda.memcpy_dtod(params, sim.field.module.get_global('d_params')[0], size)
            incoming_ptr, _ = module.get_global('j1_incoming')
            collision_ptr, size = module.get_global('j1_collision')
            cuda.memset_d8(collision_ptr, 0, size)
            incoming = np.empty(2*sim.field.nx, dtype=np.float64)
            collision = np.empty(sim.field.n, dtype=np.float64)
            sim.stepper.step_fn = module.get_function('OneStep')
            for _ in range(16):
                before = populations(sim)
                mass = before[:, 1:-1].sum()
                outgoing_bottom = before[[4,6,7], 1, :].sum()
                outgoing_top = before[[3,5,8], -2, :].sum()
                sim.run(1)
                after = populations(sim)
                cuda.memcpy_dtoh(incoming, incoming_ptr)
                cuda.memcpy_dtoh(collision, collision_ptr)
                delta = after[:, 1:-1].sum()-mass
                bottom = incoming[:sim.field.nx].sum()-outgoing_bottom
                top = incoming[sim.field.nx:].sum()-outgoing_top
                physical_residual = delta-bottom-top
                residual = physical_residual-collision.sum()
                rows.append(dict(step=sim.stepper.step_count, mass=mass, delta_mass=delta,
                                 bottom_net_in=bottom, top_net_in=top,
                                 collision_roundoff=float(collision.sum()),
                                 physical_residual=physical_residual,
                                 budget_residual=residual))
                if bottom_slot is not None:
                    outgoing = before[[4,6,7], 1, :].sum(axis=0)
                    rows[-1]['exchange'] = {
                        name: dict(incoming=float(incoming[:sim.field.nx][mask].sum()),
                                   outgoing=float(outgoing[mask].sum()))
                        for name, mask in (('slot', bottom_slot), ('flanks', ~bottom_slot))}
                    rows[-1]['exchange']['top'] = dict(
                        incoming=float(incoming[sim.field.nx:].sum()), outgoing=float(outgoing_top))
        identical = np.array_equal(sim.get_ddf(), plain)
        return dict(steps=rows, instrumented_matches_plain=identical,
                    audit_source_sha256=audit_hash,
                    max_relative_residual=max(abs(r['budget_residual'])/r['mass'] for r in rows),
                    max_physical_relative_residual=max(abs(r['physical_residual'])/r['mass'] for r in rows))
    finally:
        sim.stepper.step_fn = original_step
        compiler.cleanup_build_dir(build)


def run_case(v):
    height, nu, speed, nx = 32, .1, .02, 16
    chunk = int(np.ceil(.5*height**2/(np.pi**2*nu)))
    with make_sim(v) as sim:
        sim.initialize()
        # Reject unsupported entrypoints before observing the reference flow.
        try:
            sim.add_body('circle', center=(8,16), radius=2)
        except ValueError:
            pass
        else:
            raise AssertionError('J1 accepted an object')
        rho, ux, uy = observables(populations(sim))
        initial_mass = rho.sum()
        previous_u, previous_v, previous_mass = ux.mean(axis=1), uy.mean(axis=1), initial_mass
        history = [dict(step=0, mass=float(initial_mass))]
        steady = False
        for sample in range(1, 41):
            sim.run(chunk)
            f = populations(sim)
            if not np.isfinite(f).all():
                raise AssertionError('Nonfinite J1 populations')
            rho, ux, uy = observables(f)
            u, normal, mass = ux.mean(axis=1), uy.mean(axis=1), rho.sum()
            item = dict(step=sim.stepper.step_count, mass=float(mass),
                        u_change=float(np.sqrt(np.mean((u-previous_u)**2))/speed),
                        v_change=float(np.sqrt(np.mean((normal-previous_v)**2))),
                        mass_change=float(abs(mass-previous_mass)/initial_mass))
            history.append(item)
            previous_u, previous_v, previous_mass = u.copy(), normal.copy(), mass
            tail = history[-4:]
            steady = sample >= 8 and all(
                r['u_change'] <= GATES['u_change'] and r['v_change'] <= GATES['v_change']
                and r['mass_change'] <= GATES['mass_chunk_change'] for r in tail)
            steady = steady and np.ptp([r['mass'] for r in history[-5:]])/initial_mass <= GATES['mass_window_span']
            if steady:
                break
        y = np.arange(height+1)
        exact = speed*y/height if v == 0 else speed*np.expm1(v*y/nu)/np.expm1(v*height/nu)
        relative_l2 = float(np.linalg.norm(u[1:-1]-exact[1:-1])/np.linalg.norm(exact[1:-1]))
        v_error = float(np.max(np.abs(uy-v)))
        sections = [8,16,24]  # Distances above physical bottom; none is a BC row.
        flux = [float(np.sum(rho[s]*uy[s])) for s in sections]
        target_flux = float(rho.mean()*v*nx)
        flux_error = max(abs(q-target_flux)/abs(target_flux) for q in flux) if v else max(abs(q)/nx for q in flux)
        mass_span = float(np.ptp([r['mass'] for r in history[-5:]])/initial_mass)
        offset = float(abs(mass-initial_mass)/initial_mass)
        result = dict(v=v, height=height, nx=nx, nu=nu, top_u=speed,
                      pe=v*height/nu, re=speed*height/nu,
                      ma=float(np.hypot(speed,v)*np.sqrt(3)),
                      steps=sim.stepper.step_count, max_steps=40*chunk, steady=bool(steady),
                      relative_l2=relative_l2, v_max_error=v_error,
                      mass_offset=offset, mass_window_span=mass_span,
                      rho_range=[float(rho.min()),float(rho.max())],
                      section_y=sections, measured_flux=flux, reference_flux=target_flux,
                      flux_error=flux_error, profile_u=u.tolist(), profile_v=normal.tolist(),
                      analytic_u=exact.tolist(), history=history)
        budget = population_budget(sim)
        result['population_budget'] = budget
        result['passed'] = bool(steady and relative_l2 <= GATES['u_relative_l2']
                                and v_error <= GATES['v_max_error']
                                and flux_error <= (GATES['flux_relative_error'] if v else GATES['v_max_error'])
                                and offset <= GATES['mass_offset'] and mass_span <= GATES['mass_window_span']
                                and budget['instrumented_matches_plain']
                                and budget['max_relative_residual'] <= GATES['budget_relative_residual'])
        return result


def off_path_check(baseline):
    outcomes = []
    for kernel_root in (baseline, Path(compiler.kernel_path(''))):
        with patch.object(compiler, '_KERNEL_DIR', str(kernel_root)):
            with make_sim(0., mode='bounce_back') as sim:
                sim.initialize()
                sim.run(7)
                outcomes.append((sim.get_ddf().copy(), sim.get_flags().copy()))
    return bool(all(np.array_equal(a,b) for a,b in zip(*outcomes)))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--baseline-kernels', type=Path,
                        help='Optional pre-change CUDA source for the off-mode comparison')
    args = parser.parse_args()
    assert Path.cwd() == ROOT == Path('/home/frank14f/CelerisLab')
    assert Path(CelerisLab.__file__).resolve() == ROOT/'src/CelerisLab/__init__.py'
    assert Path(compiler.kernel_path('kernel_v2.cu')).resolve() == ROOT/'src/CelerisLab/lbm/kernels/kernel_v2.cu'
    assert os.environ.get('CUDA_VISIBLE_DEVICES') in ('0','1')
    report = dict(python=sys.executable, package=CelerisLab.__file__,
                  cuda_source=compiler.kernel_path('kernel_v2.cu'),
                  physical_gpu=os.environ['CUDA_VISIBLE_DEVICES'], gates=GATES, cases=[])
    report['off_path_bitwise_equal'] = (off_path_check(args.baseline_kernels)
                                        if args.baseline_kernels else None)
    for v in (0., .001, -.001, .01, -.01):
        case = run_case(v)
        report['cases'].append(case)
        args.output.write_text(json.dumps(report, indent=2)+'\n')
        print(json.dumps({k:case[k] for k in ('v','steps','steady','relative_l2','v_max_error',
                                             'mass_offset','mass_window_span','flux_error','passed')}), flush=True)
    return 0 if report['off_path_bitwise_equal'] is not False and all(c['passed'] for c in report['cases']) else 1


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