import json
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from qwen_vl_utils import process_vision_info

CACHE = Path("/home/s56fasla/MOTIONX/cache")
OUT = Path("/home/s56fasla/MOTIONX/example")
STEM = "Playing_with_a_water_gun_while_sitting_clip1"
SETTINGS = [("current_fps1_max8", 1.0, 8), ("proposed_fps4_max32", 4.0, 32)]

OUT.mkdir(parents=True, exist_ok=True)
man = json.load(open(CACHE / "manifest_train_ok.json"))
item = next((m for m in man if m["stem"] == STEM), man[0])
print("video:", item["video"])
print("caption:", item["text"])
json.dump(item, open(OUT / "item.json", "w"), indent=1)

for name, fps, maxf in SETTINGS:
    msg = [{"role": "user", "content": [{"type": "video", "video": item["video"],
                                         "fps": fps, "max_frames": maxf,
                                         "max_pixels": 200704}]}]
    _, vids = process_vision_info(msg)
    v = vids[0]
    if isinstance(v, torch.Tensor):
        arr = v.permute(0, 2, 3, 1).cpu().numpy()
    else:
        arr = np.asarray(v)
    arr = arr.astype(np.uint8)
    T, H, W, _ = arr.shape
    cols = min(T, 8)
    rows = (T + cols - 1) // cols
    sheet = Image.new("RGB", (cols * W, rows * H), "black")
    for i in range(T):
        sheet.paste(Image.fromarray(arr[i]), ((i % cols) * W, (i // cols) * H))
    scale = min(1.0, 1600 / sheet.width)
    if scale < 1.0:
        sheet = sheet.resize((int(sheet.width * scale), int(sheet.height * scale)))
    sheet.save(OUT / f"{name}.png")
    print(name, "frames", T, "size", (H, W), "->", OUT / f"{name}.png", flush=True)