"""Smoke test for 3-PoinTr: load a tiny subset of our prepared Stack-Can data,
build the model, run 1 forward pass, verify loss is finite."""
import os, sys, json, glob
sys.path.insert(0, "/home/hongyi/PWM/reward_model/3PoinTr_repo")

import torch
from util.training_utils.train_util import (
    build_common_arg_parser, prepare_datasets_and_loaders, normalize_and_augment,
    extract_diffusion_kwargs,
)
from util.training_utils.arch import TrajectoryModel


def main():
    DATA_ROOT = "/home/hongyi/PWM/external/3pointr_data"
    SPLIT_FILE = os.path.join(DATA_ROOT, "smoke_test_split.json")

    eps = sorted(glob.glob(os.path.join(DATA_ROOT, "no_actions/stack_can/[0-9][0-9][0-9][0-9]")))
    if len(eps) < 4:
        raise RuntimeError(f"Need >=4 episodes built; only {len(eps)} available.")
    train_eps = eps[:max(2, len(eps) - 1)]
    val_eps = eps[-1:]
    payload = {
        "actions_train": [], "actions_val": [], "actions_test": [],
        "flows_train": train_eps, "flows_val": val_eps, "flows_test": [],
        "meta": {"smoke": True}
    }
    with open(SPLIT_FILE, "w") as f:
        json.dump(payload, f, indent=2)
    print(f"[smoke] split: train={len(train_eps)}, val={len(val_eps)}")

    parser = build_common_arg_parser("smoke")
    args = parser.parse_args([
        "--dataset", DATA_ROOT + "/",
        "--split_file", SPLIT_FILE,
        "--subsample", "512",
        "--batch_size", "2",
        "--epochs", "1",
        "--num_queries", "10",
        "--arch_dim", "256", "--arch_depth", "2", "--arch_heads", "4",
        "--arch_query_depth", "1", "--arch_query_heads", "1",
        "--flow_data_only",
        "--tasks", "stack_can",
        "--lr", "1e-4",
    ])

    arts = prepare_datasets_and_loaders(args)
    train_loader = arts.train_loader
    val_loader = arts.val_loader
    num_tasks = len(arts.tasks)
    traj_len = arts.train_set.traj_len
    print(f"[smoke] traj_len={traj_len}, num_tasks={num_tasks}, train_size={len(arts.train_set)}")
    dataset_artifacts = type("X", (), {"traj_len": traj_len, "num_tasks": num_tasks})()

    diff_kwargs = extract_diffusion_kwargs(args)
    diff_kwargs["n_action_steps"] = dataset_artifacts.traj_len - 1
    model = TrajectoryModel(
        dim=args.arch_dim, depth=args.arch_depth, num_heads=args.arch_heads,
        query_depth=args.arch_query_depth, query_heads=args.arch_query_heads,
        horizon=dataset_artifacts.traj_len - 1,
        traj_len=dataset_artifacts.traj_len,
        use_task_conditioning=False, num_tasks=dataset_artifacts.num_tasks,
        decoder_only_modules=("decoder_only.point_flow_head",),
        **diff_kwargs,
    ).cuda().train()
    print(f"[smoke] model params: {sum(p.numel() for p in model.parameters())/1e6:.2f}M")

    batch = next(iter(train_loader))
    (x0, cum, ee_pos, ee_ori6d, grip, step_mask, task_ids, action_mask, point_mask) = batch
    print(f"[smoke] batch: x0={tuple(x0.shape)} cum={tuple(cum.shape)} step_mask={tuple(step_mask.shape)} pt_mask={tuple(point_mask.shape)}")

    norm = normalize_and_augment(args, x0, cum, None, None, None,
                                 step_mask, action_mask, point_mask, "cuda")
    feed = {
        "x0": norm["x0n"],
        "cum_target": norm["cum_gtn"],
        "step_mask": norm["step_mask"],
        "point_mask": norm["point_mask"],
        "num_steps": norm["step_mask"].shape[1],
        "flow_loss_w": 10.0,
    }
    out = model(feed, mode="train")
    print(f"[smoke] outputs:", {k: (v.item() if v.dim()==0 else tuple(v.shape)) for k, v in out.items() if torch.is_tensor(v)})
    print("[smoke] OK")


if __name__ == "__main__":
    main()