#!/usr/bin/env python3
"""Run one independent T-Rex Track-Force 16-step chunk from an NPZ prefix.

Required NPZ arrays:
  head_left, left_wrist, right_wrist: uint8 [9,H,W,3]
  state_eef62: float [62]
  track_past_xy: float [16,250,2] in [0,1]
  track_past_visibility: float/bool [16,250]
  tactile_force_history: float [16,10,6] (raw sensor units)

The last force-history sample is used as current force.  Output contains the
normalized model action, physical delta-base EEF62 action, absolute EEF62
targets, and predicted future tracks.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import torch
from transformers import AutoTokenizer

from groot.vla.model.dreamzero.transform.dreamzero_cotrain import (
    basic_clean,
    whitespace_clean,
)
from groot.vla.model.trex_track_force.runtime import (
    TrexRuntimeStatistics,
    delta_base_to_absolute,
)
from groot.vla.model.trex_track_force.vla import TrexTrackForceVLA


def _grid_three_views(archive: np.lib.npyio.NpzFile) -> np.ndarray:
    views = [
        np.asarray(archive[name], dtype=np.uint8)
        for name in ("head_left", "left_wrist", "right_wrist")
    ]
    if any(view.ndim != 4 or view.shape[0] != 9 or view.shape[-1] != 3 for view in views):
        raise ValueError("each RGB view must be uint8 [9,H,W,3]")
    if len({view.shape for view in views}) != 1:
        raise ValueError("all three RGB histories must have the same shape")
    _, height, width, channels = views[0].shape
    grid = np.zeros((9, 2 * height, 2 * width, channels), dtype=np.uint8)
    grid[:, :height, :width] = views[0]
    grid[:, height:, :width] = views[1]
    grid[:, :height, width:] = views[2]
    return grid


def _pad64(values: np.ndarray) -> np.ndarray:
    return np.pad(values, ((0, 0), (0, 2)), mode="constant")


def run(args: argparse.Namespace) -> None:
    device = torch.device(args.device)
    dtype = torch.bfloat16 if args.bf16 else torch.float32
    stats = TrexRuntimeStatistics.from_dataset(args.dataset_root)
    with np.load(args.input_npz, allow_pickle=False) as archive:
        history_images = _grid_three_views(archive)
        reference_state = np.asarray(archive["state_eef62"], dtype=np.float32)
        track_xy = np.asarray(archive["track_past_xy"], dtype=np.float32)
        track_visibility = np.asarray(
            archive["track_past_visibility"], dtype=np.float32
        )
        force_history_raw = np.asarray(
            archive["tactile_force_history"], dtype=np.float32
        )
    if reference_state.shape != (62,):
        raise ValueError("state_eef62 must be [62]")
    if track_xy.shape != (16, 250, 2) or track_visibility.shape != (16, 250):
        raise ValueError("past tracks must be [16,250,2] and [16,250]")
    if force_history_raw.shape != (16, 10, 6):
        raise ValueError("tactile_force_history must be [16,10,6]")

    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
    instruction = whitespace_clean(basic_clean(args.instruction))
    text = tokenizer(
        instruction,
        max_length=512,
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    )
    model = TrexTrackForceVLA.load_lora(str(args.checkpoint))
    model.eval().requires_grad_(False)
    model.to(device=device, dtype=dtype)

    normalized_state = _pad64(stats.normalize_state(reference_state)[None])
    normalized_force_history = stats.normalize_force(force_history_raw)
    inputs = {
        "history_images": torch.from_numpy(history_images[None]).to(device),
        "state": torch.from_numpy(normalized_state[:, None]).to(device, dtype),
        "track_past_xy": torch.from_numpy(track_xy[None, None]).to(device, dtype),
        "track_past_visibility": torch.from_numpy(
            track_visibility[None, None]
        ).to(device, dtype),
        "current_force": torch.from_numpy(
            normalized_force_history[-1:][None]
        ).to(device, dtype),
        "tactile_force_history": torch.from_numpy(
            normalized_force_history[None, None]
        ).to(device, dtype),
        "text": text.input_ids.to(device),
        "text_attention_mask": text.attention_mask.to(device),
    }
    with torch.inference_mode():
        prediction = model.get_action(inputs)
    normalized_action = prediction["action_pred"].float().cpu().numpy()[0]
    delta_base_action = stats.denormalize_action(normalized_action)
    absolute_action = delta_base_to_absolute(reference_state, delta_base_action)
    output = {
        "normalized_action64": normalized_action,
        "delta_base_action62": delta_base_action,
        "absolute_action62": absolute_action,
        "track_pred": prediction["track_pred"].float().cpu().numpy()[0],
    }
    if "video_latents_pred" in prediction:
        output["video_latents_pred"] = (
            prediction["video_latents_pred"].float().cpu().numpy()[0]
        )
    args.output_npz.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(args.output_npz, **output)
    print(f"wrote {args.output_npz} with 16 actions at 20 Hz")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--input-npz", type=Path, required=True)
    parser.add_argument("--output-npz", type=Path, required=True)
    parser.add_argument("--dataset-root", type=Path, required=True)
    parser.add_argument("--tokenizer-path", type=Path, required=True)
    parser.add_argument(
        "--instruction",
        default="Perform the requested bimanual manipulation.",
    )
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--bf16", action=argparse.BooleanOptionalAction, default=True)
    return parser.parse_args()


if __name__ == "__main__":
    run(parse_args())