# Written by yams_any4d agent, 2026-07-11.
"""
Live inference server for the video-model + V4Head deploy path.

Runs in the a4d-bw env (the robot deploy env cannot host torch 2.11 / Any4D).
File-based handshake with deploy_v4head_video.py (yam_local_deploy):

  bridge_dir/request_<id>.npz   from client: bgr_jpeg (uint8), joints7 (7,),
                                T_ee_in_rbase (4,4)
  bridge_dir/response_<id>.npz  from server: xyz (41,3) world, quat (41,4) wxyz,
                                grip (41,), gen_jpegs (T,) object of uint8 arrays,
                                gen_hw (2,)

Per request: inject the live scene frame into a template dataset batch (all 41
rgb slots), generate with cond_frames_raw=1 (condition ONLY on the live frame),
read the final-denoise-step DiT features, run the cached scene-only head, and
decode 41 (xyz, quat, grip) actions + the generated frames.

Usage (yukon):
  cd ~/any4d_work/Any4D && source ~/miniconda3/bin/activate a4d-bw
  PYTHONPATH=$PWD:$PWD/externals/AnyData python custom/eval/v4head_deploy_server.py \
    --exp_cfg=custom/experiment/basile/rwm4_pp70_scene_lora.py \
    --ckpt_path=$HOME/any4d_work/ckpts/pp70_scene_merged.pt \
    --head_ckpt=$HOME/any4d_work/v4head_cached/best.pt \
    --dset_cfg=custom/config/anydata/web/debug/yamyukon_pp70_scene.yaml \
    --bridge_dir=$HOME/any4d_work/deploy_bridge --output_dir=/tmp/v4h_server
"""

import argparse
import copy
import glob
import os
import time

import cv2
import numpy as np
import torch

import custom.eval.infer_anydata as base
from custom.eval.infer_v4head_rerun import merge_lora_state_dict, to_u8
from custom.eval.visuals import cached_decode_latent_video_auto
from custom.any4d.v4_head import V4Head
from custom.any4d.v4_head_viz import pca_strip, bin_grid_panel, marginal_heatmap_grid


def build_action_stub(T_ee, t_out=41):
    """(41, 20) action tensor with only frame-0 EE pose filled (eef_feat anchor).
    rot6d = first two ROWS of R_base_ee."""
    a = np.zeros((t_out, 20), dtype=np.float32)
    a[0, 0:3] = T_ee[:3, 3]
    a[0, 3:6] = T_ee[:3, :3][0]   # row 0
    a[0, 6:9] = T_ee[:3, :3][1]   # row 1
    a[0, 18] = 0.5
    return torch.from_numpy(a)


def inject_live_frame(batch, bgr):
    """Replace every scene rgb entry in the anydata batch with the live frame,
    resized to the template's resolution. Returns modified deep copy."""
    b = copy.deepcopy(batch)
    rgb = b['anydata']['rgb']
    k0 = sorted(rgb.keys())[0]
    tpl = rgb[k0]
    tpl_t = tpl[0] if tpl.ndim == 4 else tpl          # (3, H, W)
    (H, W) = tpl_t.shape[-2:]
    img = cv2.cvtColor(cv2.resize(bgr, (W, H), interpolation=cv2.INTER_AREA),
                       cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    lo = float(tpl_t.min())
    if lo < -0.01:                                     # template in [-1, 1]
        img = img * 2.0 - 1.0
    t = torch.from_numpy(img.transpose(2, 0, 1))
    for k in list(rgb.keys()):
        src = rgb[k]
        rgb[k] = t.unsqueeze(0).to(src.dtype) if src.ndim == 4 else t.to(src.dtype)
    return b


def main(args):
    (model, config, _, device, exp_cfg) = base.setup_runtime(args)
    sd = torch.load(args.ckpt_path, map_location='cpu', weights_only=False)
    sd, merged = merge_lora_state_dict(sd)
    model.pipe.dit.load_state_dict(sd, strict=False)
    print(f'[cyan]video model loaded (merged={merged})')

    head = V4Head(prep_path=args.prep, use_wrist=False,
                  temporal_mode=args.temporal_mode,
                  pred_size=args.pred_size).to(device).float()
    head.load_state_dict(torch.load(args.head_ckpt, map_location='cpu', weights_only=False))
    head.eval()
    if getattr(model.pipe.dit, 'v4head', None) is None:
        model.pipe.dit.v4head = torch.nn.Identity()
        model.pipe.dit._v4head_feats = None
    print(f'[cyan]head loaded from {args.head_ckpt}')

    dataloader = base.load_dataset(args, exp_cfg)
    template = next(iter(dataloader))
    print('[cyan]template batch ready')

    (g_lo, g_hi) = (float(head.grip_range[0]), float(head.grip_range[1]))
    directives = dict(config.val_directives) if isinstance(getattr(config, 'val_directives', None), dict) else {}
    directives['cond_frames_raw'] = 1                  # condition ONLY on the live frame

    os.makedirs(args.bridge_dir, exist_ok=True)
    os.makedirs(args.output_dir, exist_ok=True)
    print(f'[green]SERVER READY — watching {args.bridge_dir}')

    while True:
        reqs = sorted(glob.glob(f'{args.bridge_dir}/request_*.npz'))
        if not reqs:
            time.sleep(0.3)
            continue
        req_path = reqs[0]
        rid = os.path.basename(req_path)[len('request_'):-len('.npz')]
        try:
            time.sleep(0.2)                            # let the writer finish
            req = np.load(req_path, allow_pickle=True)
            bgr = cv2.imdecode(req['bgr_jpeg'], cv2.IMREAD_COLOR)
            T_ee = req['T_ee_in_rbase'].astype(np.float64)
            print(f'[cyan]request {rid}: frame {bgr.shape}, generating...')
            t0 = time.time()

            batch = inject_live_frame(template, bgr)
            with torch.no_grad():
                val_dict, _ = model.validation_step(
                    data_batch=batch, iteration=-1, dataloader_key='deploy',
                    local_path=args.output_dir, directives=dict(directives),
                    val_iter=0, seed=int(time.time()) % 100000)
            feats = model.pipe.dit._v4head_feats
            (Hp, Wp) = feats[0].shape[2], feats[0].shape[3]
            pixel_hw = (Hp * 16, Wp * 16)
            act = build_action_stub(T_ee)
            with torch.no_grad():
                out = head({0: feats[0].float()}, pixel_hw, act[None].to(device))

            vpos = out['voxel_positions']
            if vpos.ndim == 4:
                vpos = vpos.unsqueeze(0)
            xyz = head.decode_pred_xyz(out['volume_logits'].float(), vpos.float())[0].cpu().numpy()
            rot_bin = out['rot_logits'][0].float().argmax(-1).cpu().numpy()
            quat = head.rot_centroids.float().cpu().numpy()[rot_bin]
            g_bin = out['grip_logits'][0].float().argmax(-1).cpu().numpy()
            grip = g_lo + (g_bin + 0.5) * (g_hi - g_lo) / head.n_gripper_bins

            gen_jpegs = []
            gen_bgr = []
            gen_hw = (0, 0)
            if 'rgb0' in val_dict.get('y0_pred_entries', {}):
                vid = cached_decode_latent_video_auto(val_dict['y0_pred_entries'], 'rgb0', model.vae)[0]
                gen_hw = (int(vid.shape[2]), int(vid.shape[3]))
                for t in range(vid.shape[1]):
                    fr = to_u8(vid[:, t])
                    gen_bgr.append(fr)
                    ok, buf = cv2.imencode('.jpg', fr, [cv2.IMWRITE_JPEG_QUALITY, 85])
                    gen_jpegs.append(np.frombuffer(buf.tobytes(), dtype=np.uint8))

            # diagnostic panels (token PCA / upsampled PCA / heatmap grid on GEN
            # frames / grip + rot bin grids) — rendered here, logged by the client
            def _jp(img_bgr):
                ok, b = cv2.imencode('.jpg', img_bgr, [cv2.IMWRITE_JPEG_QUALITY, 88])
                return np.frombuffer(b.tobytes(), dtype=np.uint8)
            panels = {}
            try:
                tok = out['scene_feats_per_frame'][0].detach().float().cpu().numpy()
                panels['panel_pca_tok'] = _jp(cv2.cvtColor(pca_strip(tok), cv2.COLOR_RGB2BGR))
                fma = out['feat_maps_all'][0].detach().float().cpu().numpy()
                panels['panel_pca_up'] = _jp(cv2.cvtColor(
                    pca_strip(fma, cell_w=160, max_frames=fma.shape[0]), cv2.COLOR_RGB2BGR))
                vol_np = out['volume_logits'][0].detach().float().cpu().numpy()
                bgf = gen_bgr if gen_bgr else [np.full((64, 64, 3), 80, np.uint8)]
                panels['panel_heatmap'] = _jp(marginal_heatmap_grid(vol_np, bgf))
                panels['panel_grip'] = _jp(bin_grid_panel(
                    out['grip_logits'][0].float().cpu().numpy()))
                panels['panel_rot'] = _jp(bin_grid_panel(
                    out['rot_logits'][0].float().cpu().numpy()))
            except Exception as pe:
                print(f'[yellow]panel render failed: {pe}')

            tmp = f'{args.bridge_dir}/response_{rid}.tmp.npz'
            np.savez(tmp, xyz=xyz.astype(np.float32), quat=quat.astype(np.float32),
                     grip=grip.astype(np.float32),
                     gen_jpegs=np.array(gen_jpegs, dtype=object),
                     gen_hw=np.array(gen_hw), status=np.array('ok'), **panels)
            os.rename(tmp, f'{args.bridge_dir}/response_{rid}.npz')
            print(f'[green]response {rid} in {time.time()-t0:.1f}s | '
                  f'xyz[0]={xyz[0].round(3)} grip[0]={grip[0]:.2f}')
        except Exception as e:
            import traceback
            traceback.print_exc()
            np.savez(f'{args.bridge_dir}/response_{rid}.npz',
                     status=np.array(f'error: {type(e).__name__}: {e}'))
        finally:
            try:
                os.remove(req_path)
            except OSError:
                pass


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--exp_cfg', required=True)
    parser.add_argument('--ckpt_path', required=True)
    parser.add_argument('--head_ckpt', required=True)
    parser.add_argument('--dset_cfg', required=True)
    parser.add_argument('--bridge_dir', required=True)
    parser.add_argument('--output_dir', required=True)
    parser.add_argument('--prep', default='/home/robot-lab/any4d_work/v4head_prep.npz')
    parser.add_argument('--temporal_mode', default='per_step', choices=['per_step', 'stack_mlp'])
    parser.add_argument('--pred_size', type=int, default=64)
    parser.add_argument('--num_steps', type=int, default=25)
    parser.add_argument('--batch_size', type=int, default=1)
    parser.add_argument('--device', default='cuda')
    parser.add_argument('--gpu_id', type=int, default=0)
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--data_overrides', default='{"subsample": null, "single_sample": "first"}')
    parser.add_argument('--data_defaults', default=None)
    parser.add_argument('--any4d_overrides', default=None)
    _a = parser.parse_args()
    _d = dict(exp_overrides=None, cond_aug_sigma=None, guidance=None, infer_metrics=False,
              metric_resolution=None, num_samples=1, num_segments=4, perturb_action=None,
              perturb_per_sample=False, perturb_traj=None, run_autoregressive=False,
              shard=None, visual_detail=0, viz_extra_modes=None, donor_rgb_path=None,
              donor_rgb_cam=0, extrapolation_strategy='backtrack', stop_after=-1,
              seed=1234, cond_frames_raw=None)
    for (k, v) in _d.items():
        if not hasattr(_a, k):
            setattr(_a, k, v)
    main(_a)