"""Bounded analytic Poiseuille validation; run on an explicitly selected GPU.

Reference: ref/lbmpy/src/lbmpy/analytical_solutions.py::poiseuille_flow and
ref/lbmpy/tests/poiseuille.py (x periodic, half-way walls, postcollision F/2
velocity correction). D2Q9 FP32/shifted, no LES, nu=.1, Re=Umax*H/nu=4.8.
Half-way walls are y=.5 and H+.5; u=F*s*(H-s)/(2*rho0*nu).
No fully periodic Taylor-Green claim: current step kernels always correct the
rows y=1,NY-2 as walls, even if their neighboring flags are made fluid.
"""
import argparse
import json
import os
from pathlib import Path
from tempfile import TemporaryDirectory
import time

import numpy as np
import pycuda.driver as cuda
import CelerisLab
from CelerisLab.simulation import Simulation
from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY, FLUID, SOLID, BC_WALL

ROOT=Path(__file__).resolve().parents[2]
CX=np.asarray(D2Q9_EX)
CY=np.asarray(D2Q9_EY)
W=np.asarray([4/9]+[1/9]*4+[1/36]*4,dtype=np.float32)


def spatial_populations(sim):
    f=np.empty((9,sim.field.ny,sim.field.nx),np.float32)
    with sim.ctx.activate():
        cuda.memcpy_dtoh(f,sim.field.ddf_gpu)
    if sim.lbm_cfg.streaming=="esopull":
        p=sim.stepper.step_count-1
        out=np.empty_like(f);out[0]=f[0]
        for i in range(1,9,2):
            out[i]=np.roll(f[i+1 if p%2 else i],(-CY[i],-CX[i]),(0,1))
            out[i+1]=f[i if p%2 else i+1]
        f=out
    return f.astype(np.float64)+W[:,None,None]


def run(height,collision,streaming,reynolds=4.8):
    nu=.1;umax=reynolds*nu/height;force=8*nu*umax/height**2
    cfg=json.loads((ROOT/'src/CelerisLab/configs/config_lbm.json').read_text())
    cfg['grid'].update(nx=16,ny=height+2,nz=1,lattice_model='D2Q9')
    cfg['physics'].update(viscosity=nu,velocity=0.,rho=1.)
    cfg['method'].update(collision=collision,streaming=streaming,store_precision='FP32',
                         ddf_shifting=True,y_wall_bc='bounce_back')
    cfg['method']['les']['enabled']=False
    started=time.monotonic()
    with TemporaryDirectory(prefix='celeris-poiseuille-') as directory:
        lbm=Path(directory)/'lbm.json';body=Path(directory)/'body.json'
        lbm.write_text(json.dumps(cfg));body.write_text('{"objects":[]}')
        sim=Simulation(str(lbm),str(body),device_id=0)
    try:
        sim.initialize()
        # Use implemented modulo neighbors in x, retaining real y walls.
        flag=sim.field.flag.reshape(height+2,16)
        flag[:]=FLUID;flag[[0,-1],:]=SOLID|BC_WALL
        sim.field.upload_flags()
        sim.field.update_params(fx=force)
        decay_time=height**2/(np.pi**2*nu)
        chunk=int(np.ceil(.5*decay_time));previous=np.zeros(height)
        history=[]
        for _ in range(24):  # at most 12 slowest diffusion decay times
            sim.run(chunk)
            f=spatial_populations(sim)[:,1:-1,:]
            rho=f.sum(axis=0)
            ux=(np.einsum('q,qyx->yx',CX,f)-force/2)/rho
            uy=np.einsum('q,qyx->yx',CY,f)/rho
            profile=ux.mean(axis=1)
            residual=float(np.linalg.norm(profile-previous)/np.sqrt(height)/umax)
            history.append(dict(step=sim.stepper.step_count,residual=residual))
            if residual<2e-5: break
            previous=profile
        s=np.arange(height)+.5
        exact=force*s*(height-s)/(2*nu)
        error=float(np.linalg.norm(profile-exact)/np.linalg.norm(exact))
        mass_drift=float(abs(rho.mean()-1))
        transverse=float(np.max(np.abs(uy)))
        limit={16:.01,32:.003,64:.0015}.get(height,.01)
        result=dict(height=height,collision=collision,streaming=streaming,
                    nu=nu,umax=umax,force_density=force,re=umax*height/nu,
                    ma=umax*np.sqrt(3),steps=sim.stepper.step_count,
                    diffusion_times=sim.stepper.step_count/decay_time,
                    relative_l2=error,relative_l2_limit=limit,
                    mass_drift=mass_drift,max_transverse=transverse,
                    steady_residual=residual,profile=profile.tolist(),
                    analytic_profile=exact.tolist(),history=history,
                    seconds=time.monotonic()-started)
        result['passed']=bool(np.isfinite(f).all() and error<=limit and mass_drift<=5e-4
                              and transverse<=1e-5 and residual<=2e-5)
        return result
    finally:
        sim.close()


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--heights',type=int,nargs='+',default=[16,32,64])
    parser.add_argument('--collisions',nargs='+',default=['SRT','TRT','MRT'])
    parser.add_argument('--streaming',nargs='+',default=['double_buffer','esopull'])
    parser.add_argument('--re',type=float,default=4.8,help='Reynolds number; default is the refinement baseline')
    parser.add_argument('--output',type=Path,default=Path('/tmp/celeris-poiseuille.json'))
    args=parser.parse_args()
    if not os.environ.get('CUDA_VISIBLE_DEVICES'):
        parser.error('Select a physical GPU explicitly with CUDA_VISIBLE_DEVICES')
    assert Path(CelerisLab.__file__).resolve()==ROOT/'src/CelerisLab/__init__.py'
    report=dict(package=CelerisLab.__file__,cuda_visible_devices=os.environ['CUDA_VISIBLE_DEVICES'],device_id=0,
                reference='ref/lbmpy/tests/poiseuille.py',results=[])
    for h in args.heights:
        for collision in args.collisions:
            for streaming in args.streaming:
                item=run(h,collision,streaming,args.re)
                report['results'].append(item)
                args.output.write_text(json.dumps(report,indent=2))
                print(json.dumps({k:v for k,v in item.items() if k not in
                    ('profile','analytic_profile','history')}),flush=True)
    return 0 if all(r['passed'] for r in report['results']) else 1


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