"""One paired zero/PPO Karman pinball VIV experiment; see docs/karman_viv.md.

Research-only adapter: static precursor -> two fresh, identical moving plants.
No moving checkpoint/rollback, policy retraining, or force filtering.
"""
from __future__ import annotations

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

import numpy as np

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))
from research.steady_pinball.run import Structure, config, scaled_points, CENTRES, PROBES, COLUMNS

BIAS = np.array([0., -4., 4.])
MODEL_SHA = '148240c1dcb8b11d8e5a0a9e991f6874c4e799c9744707bf0fb32a8efeb0468d'
NORM_SHA = '534fda818c030a56960e95d0f3ac1e350d887d03db7f80f1d6453195a502b7f0'
RELEASE_SOURCE = '516975ef6482f983ace347cb10de7450d2fbd680'
STEP_COLUMNS = (*COLUMNS, 'omega_F', 'omega_U', 'omega_L')


def sha(path):
    h = hashlib.sha256()
    with open(path, 'rb') as f:
        for block in iter(lambda: f.read(1024*1024), b''):
            h.update(block)
    return h.hexdigest()


def array_sha(a):
    return hashlib.sha256(memoryview(np.ascontiguousarray(a)).cast('B')).hexdigest()


def observation(sensor_mean, force_mean, scale, norm):
    """Fine-grid area means/CME -> trained D20 sensor sums/forces -> frozen norm."""
    sensors = np.asarray(sensor_mean, np.float64)*scale*78.
    forces = np.asarray(force_mean, np.float64)*scale
    raw = np.asarray([*sensors, *forces], np.float32)
    preclip = np.hstack((raw[6:]/norm['force_norm_fact'],
                        (raw[:6]-np.asarray(norm['sens_deviation'], np.float32)) /
                        np.asarray(norm['sens_norm_fact'], np.float32)))
    if not np.isfinite(preclip).all():
        raise RuntimeError('nonfinite policy observation')
    return raw, np.clip(preclip, -1, 1).astype(np.float32), int((np.abs(preclip)>1).sum())


def wall_target(action, scale):
    return (8.*np.asarray(action)+BIAS)*(.01/scale)


def ema(eff, target, scale):
    alpha = .9**(1./scale**2)
    return alpha*eff + (1.-alpha)*target


def load_policy(model):
    # Same inference spaces as legacy_test.core.model_loader, without importing it.
    import gymnasium as gym
    from stable_baselines3 import PPO
    class InferenceEnv(gym.Env):
        observation_space = gym.spaces.Box(-1., 1., (12,), np.float32)
        action_space = gym.spaces.Box(-1., 1., (3,), np.float32)
        def reset(self, *, seed=None, options=None):
            super().reset(seed=seed)
            return np.zeros(12, np.float32), {}
        def step(self, action):
            return np.zeros(12, np.float32), 0., False, False, {}
    return PPO.load(str(model), env=InferenceEnv(), device='cpu')


def add_scene(sim, scale, moving=False, pinball=True):
    from CelerisLab.body.motion import Translation
    sim.add_body('circle', scaled_points(((200., 255.5),), scale)[0], 20.*scale)
    probes = [sim.add_body('sensor', c, 5.*scale) for c in scaled_points(PROBES, scale)]
    bodies = []
    if pinball:
        envelope = Translation((0., 0., -20.*scale, 20.*scale), .05/scale, shared=True) if moving else None
        bodies = [sim.add_body('circle', c, 10.*scale, translation=envelope)
                  for c in scaled_points(CENTRES, scale)]
    return probes, bodies


def capture_seed(sim):
    """Capture both FP32 buffers at fixed geometry; never a moving-state checkpoint."""
    import pycuda.driver as cuda
    with sim.ctx.activate():
        sim.stream.synchronize()
        buffers = [np.empty_like(sim.field.ddf) for _ in range(2)]
        for dst, src in zip(buffers, (sim.field.ddf_gpu, sim.field.temp_gpu)):
            cuda.memcpy_dtoh(dst, src)
    return (*buffers, sim.get_flags(), sim.stepper.step_count)


def install_seed(sim, seed):
    """Only a new, zero-offset instance may receive the static release state."""
    import pycuda.driver as cuda
    current, spare, flags, clock = seed
    if sim.stepper.step_count != 0 or any(a.shape != sim.field.ddf.shape or a.dtype != sim.field.ddf.dtype
                                         for a in (current, spare)):
        raise ValueError('release transfer requires a fresh matching field')
    if not np.array_equal(sim.get_flags(), flags):
        raise ValueError('static/moving release flags differ')
    for obj in sim.bodies.objects:
        if obj.motion is not None and (obj.motion.applied != (0., 0.) or obj.motion.pending is not None):
            raise ValueError('release transfer requires zero geometry and no target')
    with sim.ctx.activate():
        sim.stream.synchronize()
        cuda.memcpy_htod(sim.field.ddf_gpu, current)
        cuda.memcpy_htod(sim.field.temp_gpu, spare)
    sim.stepper._step_count = clock  # one-time research initialization, not restore()
    sim.field.invalidate_host_ddf_cache()
    copied = capture_seed(sim)
    if copied[3] != clock or any(not np.array_equal(a, b, equal_nan=True) for a, b in zip(copied[:3], seed[:3])):
        raise RuntimeError('release transfer did not preserve both buffers/flags')
    return dict(current=array_sha(copied[0]), spare=array_sha(copied[1]), flags=array_sha(copied[2]), clock=copied[3])


def check_time(start, cap):
    if time.monotonic()-start > cap:
        raise TimeoutError('paired experiment wall cap')


def warmup(sim, count, start, cap):
    while count:
        check_time(start, cap)
        n = min(count, 10000)
        sim.run(n)
        count -= n


def precursor(cfgpath, scale, warm, si, targets, out, start, cap):
    from CelerisLab import Simulation
    with Simulation(str(cfgpath)) as sim:
        probes, _ = add_scene(sim, scale, pinball=False)
        sim.initialize()
        warmup(sim, warm, start, cap)
        target = []
        for _ in range(targets):
            check_time(start, cap)
            sim.run(si)
            target.append([v*scale*78. for p in probes for v in sim.read_sensor(p)[:2]])
        np.save(out/'target.npy', np.asarray(target, np.float32))
        first = sim.bodies.count
        for c in scaled_points(CENTRES, scale):
            sim.add_body('circle', c, 10.*scale)
        sim.sync_bodies()
        bodies = list(range(first, first+3))
        warmup(sim, warm-1, start, cap)
        sim.run(1)
        force = sum(float(sim.read_force(p)[1]) for p in bodies)
        seed = capture_seed(sim)
        np.savez(out/'release.npz', current=seed[0], spare=seed[1], flags=seed[2], clock=seed[3], force=force)
        print(json.dumps(dict(event='precursor_complete', clock=seed[3], force=force,
                              elapsed_s=time.monotonic()-start)), flush=True)
        return seed, force


def run_arm(name, cfgpath, scale, si, intervals, seed, initial_force, model, norm, out, start, cap, *, mass_ratio=2.):
    from CelerisLab import Simulation
    from CelerisLab.body.motion import MovingTables
    cols = {k:i for i, k in enumerate(MovingTables.DIAG)}
    diameter = 20.*scale
    structure = Structure(mass_ratio, 1.e-4/scale**2, 0., diameter)
    arm = out/name
    arm.mkdir()
    rows = np.lib.format.open_memmap(arm/'steps.npy', mode='w+', dtype=np.float64,
                                    shape=(si*intervals, len(STEP_COLUMNS)))
    rows[:] = np.nan
    signals = {k:np.full((intervals, width), np.nan, np.float32) for k, width in
               [('sensors',6), ('forces',6), ('actions',3), ('observations',12)]}
    completed = completed_intervals = clipped = 0
    status, transfer = 'running', None
    force = initial_force
    obs = np.zeros(12, np.float32)
    arm_start = time.monotonic()
    try:
        with Simulation(str(cfgpath)) as sim:
            probes, bodies = add_scene(sim, scale, moving=True)
            sim.initialize()
            transfer = install_seed(sim, seed)
            for i in range(intervals):
                action = -BIAS/8 if name == 'zero' else np.asarray(model.predict(obs, deterministic=True)[0], np.float32)
                if action.shape != (3,) or not np.isfinite(action).all() or (np.abs(action)>1).any():
                    raise RuntimeError('invalid policy action')
                target, eff = wall_target(action, scale), np.zeros(3)
                sensor_sum, force_sum = np.zeros(6), np.zeros(6)
                for k in range(si):
                    if k % 1000 == 0:
                        check_time(start, cap)
                    y1, v1 = structure.candidate(force)
                    if not np.isfinite([y1,v1]).all() or abs(y1)>diameter:
                        raise RuntimeError('nonfinite structure or |y| > D; no retuning')
                    eff = ema(eff, target, scale)
                    omega = eff/(diameter/2)
                    for p, w in zip(bodies, omega):
                        sim.set_body(p, omega=float(w), offset=(0., y1), vy=y1-structure.y)
                    sim.run(1, zero_obs=True, sync_obs=True)
                    d = sim.moving_diagnostics()
                    cme, wen = d[:,cols['cme_y']], d[:,cols['wen_y']]
                    node = d[:,cols['cover_jy']]-d[:,cols['uncover_jy']]
                    row = (completed+1,y1,v1,*cme,*wen,*node,
                           d[:,cols['eq_only']].sum(),d[:,cols['no_donor']].sum(),d[:,cols['crossings']].sum(),*omega)
                    if not np.isfinite(row).all() or row[13] != 0:
                        raise RuntimeError('nonfinite moving diagnostic or no_donor')
                    f = np.array([v for p in bodies for v in sim.read_force(p)[:2]])
                    s = np.array([v for p in probes for v in sim.read_sensor(p)[:2]])
                    if not np.isfinite([f,s]).all():
                        raise RuntimeError('nonfinite force/sensor observation')
                    # Policy sees packed CME as in fixed-body training; structure sees CME+Wen.
                    force_sum += f
                    sensor_sum += s
                    force = float(cme.sum(dtype=np.float64)+wen.sum(dtype=np.float64))
                    rows[completed] = row
                    completed += 1
                    structure.y, structure.v = y1,v1
                raw, obs, nclip = observation(sensor_sum/si, force_sum/si, scale, norm)
                clipped += nclip
                for key, val in [('sensors',raw[:6]),('forces',raw[6:]),('actions',action),('observations',obs)]:
                    signals[key][i] = val
                completed_intervals = i+1
                if (i+1)%10 == 0 or i+1 == intervals:
                    rows.flush()
                    np.savez(arm/'signals.npz', **signals)
                    print(json.dumps(dict(event='progress', arm=name, interval=i+1, step=completed,
                                          tU_D=completed*.01/scale/diameter, y_D=y1/diameter,
                                          elapsed_s=time.monotonic()-arm_start)), flush=True)
            status = 'duration_completed_not_qualified'
    except BaseException as exc:
        status = f'error: {type(exc).__name__}: {exc}'
        raise
    finally:
        rows.flush()
        np.savez(arm/'signals.npz', **signals)
        report = dict(status=status, completed_steps=completed, completed_intervals=completed_intervals,
                      elapsed_s=time.monotonic()-arm_start, release_transfer=transfer, clipped_observation_channels=clipped,
                      columns=STEP_COLUMNS, M=structure.mass,K=structure.stiffness,fn=structure.frequency,damping=0.,
                      max_abs_y_D=float(np.max(np.abs(rows[:completed,1]))/diameter) if completed else None)
        (arm/'report.json').write_text(json.dumps(report,indent=2)+'\n')


def saved_release(directory):
    """Read the retained m*=2 static release; not a moving-state restart."""
    case = json.loads((directory/'case.json').read_text())
    expected = dict(source=RELEASE_SOURCE, model_sha256=MODEL_SHA, norm_sha256=NORM_SHA,
                    smoke=False, D=40, U=.005, nu=.004, warmup_each=2048000,
                    target_intervals=150, SI=3200)
    if any(case.get(k) != v for k, v in expected.items()):
        raise ValueError('saved release case does not match the frozen precursor')
    if json.loads((directory/'config.json').read_text()) != config(40):
        raise ValueError('saved release config differs')
    with np.load(directory/'release.npz') as data:
        seed = (data['current'], data['spare'], data['flags'], int(data['clock']))
        force = float(data['force'])
    transfer = json.loads((directory/'zero/report.json').read_text())['release_transfer']
    if seed[3] != transfer['clock'] or seed[3] != 4576000 or not np.isfinite(force):
        raise ValueError('saved release clock/force differs')
    if [array_sha(a) for a in seed[:3]] != [transfer[k] for k in ('current','spare','flags')]:
        raise ValueError('saved release buffer/flags hash differs')
    target = np.load(directory/'target.npy')
    if target.shape != (150,6) or not np.isfinite(target).all():
        raise ValueError('saved target is incomplete or nonfinite')
    return seed, force


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--model', type=Path, required=True)
    ap.add_argument('--norm', type=Path, required=True)
    ap.add_argument('--out', type=Path, required=True)
    ap.add_argument('--smoke', action='store_true', help='short plumbing only; no physical conclusion')
    ap.add_argument('--mass-ratio', type=float, choices=(2.,4.), default=2.,
                    help='4 is the single declared exploratory candidate; fn and damping unchanged')
    ap.add_argument('--release-dir', type=Path, help='retained m*=2 precursor; skip warmup')
    ap.add_argument('--wall-seconds', type=float, default=14400.)
    args = ap.parse_args()
    if not np.isfinite(args.wall_seconds) or args.wall_seconds<=0:
        ap.error('positive finite wall cap required')
    if args.release_dir is not None and args.smoke:
        ap.error('saved physical release cannot be used for the plumbing smoke')
    if sha(args.model)!=MODEL_SHA or sha(args.norm)!=NORM_SHA:
        ap.error('model/normalizer must match frozen provenance')
    out = args.out.resolve()
    if out == ROOT or ROOT in out.parents:
        ap.error('output must be outside checkout')
    out.mkdir(parents=True, exist_ok=False)
    norm = json.loads(args.norm.read_text())
    model = load_policy(args.model)
    scale = 2.
    warm, si, targets, intervals = (32,8,2,4) if args.smoke else (2048000,3200,150,400)
    cfgpath = out/'config.json'
    cfgpath.write_text(json.dumps(config(40),indent=2)+'\n')
    receipt = dict(source=subprocess.check_output(['git','-C',str(ROOT),'rev-parse','HEAD'],text=True).strip(),
                   model_sha256=sha(args.model),norm_sha256=sha(args.norm),smoke=args.smoke,D=40,U=.005,nu=.004,
                   warmup_each=warm,target_intervals=targets,SI=si,response_intervals=intervals,
                   response_steps=si*intervals,response_tU_D=si*intervals*.005/40,
                   rhs='previous native step sum CME_y+Wen_y; policy uses interval packed CME',
                   probe_frame='lab-fixed; fine area mean * 2 * trained 78 cells',
                   ema_native_alpha=.9**.25, structural_damping=0., mass_ratio=args.mass_ratio,
                   exploratory=args.mass_ratio==4., columns=STEP_COLUMNS)
    loaded = None
    if args.release_dir is not None:
        loaded = saved_release(args.release_dir)
        cfgpath.write_bytes((args.release_dir/'config.json').read_bytes())
        (out/'target.npy').write_bytes((args.release_dir/'target.npy').read_bytes())
        receipt.update(release_dir=str(args.release_dir.resolve()), release_source=RELEASE_SOURCE,
                       release_sha256=sha(args.release_dir/'release.npz'),
                       warmup_executed=0, target_sha256=sha(out/'target.npy'))
    (out/'case.json').write_text(json.dumps(receipt,indent=2)+'\n')
    start = time.monotonic()
    status = 'running'
    try:
        seed, force = loaded if loaded is not None else precursor(cfgpath,scale,warm,si,targets,out,start,args.wall_seconds)
        for name in ('zero','controlled'):
            run_arm(name,cfgpath,scale,si,intervals,seed,force,model,norm,out,start,args.wall_seconds,
                    mass_ratio=args.mass_ratio)
        status = 'smoke_completed' if args.smoke else 'pair_completed_not_qualified'
    except BaseException as exc:
        status = f'error: {type(exc).__name__}: {exc}'
        raise
    finally:
        (out/'terminal.json').write_text(json.dumps(dict(status=status,elapsed_s=time.monotonic()-start),indent=2)+'\n')


if __name__ == '__main__':
    main()