"""Serial Celeris diagnostic/scan runner; outputs live outside this package."""
from __future__ import annotations
import argparse, json, os, tempfile
from pathlib import Path
import numpy as np
from .contract import CanonicalContract
from .metrics import profile_errors


def _cfg(c: CanonicalContract, D: float):
    scale=D/c.diameter; nx=int(round(c.nx*scale)); ny=int(round(c.ny*scale))
    return {"grid":{"lattice_model":"D2Q9","nx":nx,"ny":ny,"nz":1},
      "physics":{"data_type":"FP32","viscosity":c.u_inf*D/c.re_d,"velocity":c.u_inf,"rho":1.0},
      "method":{"collision":"MRT","streaming":"double_buffer","store_precision":"FP32","ddf_shifting":False,
        "les":{"enabled":False,"cs":0.16,"closed_form":True},"trt":{"magic_param":0.1875},
        "inlet":{"profile":"uniform","scheme":"regularized","trt_neq_damp":0.5,"regularized_neq_damp":0.5},
        "outlet":{"mode":"neq_extrap","backflow_clamp":True,"blend_alpha":0.7,"srt_neq_damp":0.5},
        "y_wall_bc":"free_slip","omega_guard":{"min":0.01,"max":1.99}},
      "cuda":{"threads_per_block":256,"compute_capability":"auto"}}


def run_case(out: Path, *, branch: str, s: float, D: float, t_end: float, sample_dt: float, half_height_d: float = 15.0, upstream_d: float = 50.0, downstream_d: float = 50.0):
    from CelerisLab import Simulation
    c=CanonicalContract(diameter=D, half_height_d=half_height_d, upstream_d=upstream_d, downstream_d=downstream_d); cfg=_cfg(c,D); centers=c.centers
    action=np.zeros(3) if branch=="qin" else c.action(s,branch)
    body={"objects":[]} if branch=="qin" else {"objects":[{"type":"cylinder","center":list(z),"radius":c.radius,"omega":float(w)} for z,w in zip(centers,action)]}
    out.mkdir(parents=True,exist_ok=False)
    (out/'contract.json').write_text(json.dumps({"branch":branch,"s":s,"D":D,"t_end":t_end,"sample_dt":sample_dt,"action":action.tolist(),"centers":centers,"oracle":c.oracle(s)},indent=2))
    with tempfile.TemporaryDirectory(prefix='steady_cfd_') as td:
        lp=Path(td)/'config_lbm.json'; bp=Path(td)/'config_body.json'
        lp.write_text(json.dumps(cfg)); bp.write_text(json.dumps(body))
        sim=Simulation(lbm_config_path=str(lp),body_config_path=str(bp)); sim.initialize()
        steps=int(round(t_end*D/c.u_inf)); stride=max(1,int(round(sample_dt*D/c.u_inf)))
        hist=[]
        for stop in range(stride,steps+1,stride):
            sim.run(min(stride,steps-(stop-stride)),zero_obs=True)
            m=sim.get_macroscopic(); ux=np.asarray(m['ux']); uy=np.asarray(m['uy'])
            ix=int(round(centers[0][0]+c.primary_station_d*D))
            prof=profile_errors(ux[2:-2,ix],uy[2:-2,ix],u_inf=c.u_inf)
            hist.append([stop,prof['E_inf_vector'],prof['E_L2_vector']])
        m=sim.get_macroscopic(); sim.close()
    ux=np.asarray(m['ux'],np.float32); uy=np.asarray(m['uy'],np.float32); rho=np.asarray(m['rho'],np.float32)
    np.savez_compressed(out/'endpoint.npz',ux=ux,uy=uy,rho=rho,history=np.asarray(hist),action=action)
    from CelerisLab.common.render import compute_vorticity,render_vorticity_field
    render_vorticity_field(compute_vorticity(ux,uy),nx=cfg['grid']['nx'],ny=cfg['grid']['ny'],out_path=str(out/'vorticity.png'),cylinders=[] if branch=='qin' else [(z,c.radius) for z in centers])
    result={"branch":branch,"action":action.tolist(),"steps":steps,"final_uniform_diagnostic":dict(zip(["step","E_inf_vector","E_L2_vector"],hist[-1]))}
    (out/'result.json').write_text(json.dumps(result,indent=2)); return result


def main(argv=None):
    p=argparse.ArgumentParser(); p.add_argument('command',choices=['oracle','diagnostic']); p.add_argument('--out',type=Path)
    p.add_argument('--D',type=float,default=10.0); p.add_argument('--s',type=float,default=5.0); p.add_argument('--t-end',type=float,default=40.0); p.add_argument('--sample-dt',type=float,default=2.0)
    a=p.parse_args(argv); c=CanonicalContract(diameter=a.D)
    if a.command=='oracle': print(json.dumps(c.oracle(a.s),indent=2)); return 0
    if a.out is None: p.error('--out required')
    a.out.mkdir(parents=True,exist_ok=False)
    rows=[]
    for b in ('qin','plus-minus','minus-plus'):
        rows.append(run_case(a.out/b,branch=b,s=a.s,D=a.D,t_end=a.t_end,sample_dt=a.sample_dt))
    
    q=np.load(a.out/'qin'/'endpoint.npz'); ix=int(round(c.centers[0][0]+c.primary_station_d*a.D)); refux=q['ux'][2:-2,ix]; refuy=q['uy'][2:-2,ix]
    for row in rows:
        e=np.load(a.out/row['branch']/'endpoint.npz'); du=(e['ux'][2:-2,ix]-refux)/c.u_inf; dv=(e['uy'][2:-2,ix]-refuy)/c.u_inf; vec=np.hypot(du,dv)
        row['direct_qin']={"E_inf_x":float(np.max(np.abs(du))),"E_inf_y":float(np.max(np.abs(dv))),"E_inf_vector":float(np.max(vec)),"E_L2_vector":float(np.sqrt(np.mean(vec*vec)))}
    (a.out/'summary.json').write_text(json.dumps(rows,indent=2)); return 0

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