import os
import time
import argparse
import numpy as np
import torch
from src.utils import load_config
from src.mind3d import MinD3D
from src.utils import torch_init_model, set_random_seed


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default="./configs/mind3d.yaml")
    parser.add_argument('--check_point_path', type=str, default="ckpt/mind3d_sub01.pt")
    parser.add_argument('--fmri_dir', type=str, default="/home/hubin/data/fMRI-Shape/sub-0001/npy_data")
    parser.add_argument('--test_list', type=str, default="/home/hubin/data/fMRI-Shape/annotations/core_test_list.txt")
    parser.add_argument('--out_dir', type=str, default="/home/hubin/trellis_work/outputs/mind3d_sub01")
    parser.add_argument('--topk', type=int, default=250)
    parser.add_argument('--seed', type=int, default=100)
    parser.add_argument('--shard', type=int, default=0)
    parser.add_argument('--num_shards', type=int, default=1)
    args = parser.parse_args()

    cfg = load_config(args.config, 'configs/default.yaml')
    device = torch.device("cuda")
    os.makedirs(os.path.join(args.out_dir, "meshes"), exist_ok=True)

    model_full = MinD3D(cfg, device=device)
    state_dict = torch.load(args.check_point_path, map_location='cpu', weights_only=False)
    torch_init_model(model_full, state_dict)
    del state_dict
    model_full.eval()

    ids = [l.strip() for l in open(args.test_list) if l.strip()][args.shard::args.num_shards]
    for i, obj in enumerate(ids):
        out_path = os.path.join(args.out_dir, "meshes", obj.replace("/", "_") + ".obj")
        if os.path.exists(out_path):
            continue
        set_random_seed(args.seed)
        t0 = time.time()
        fmri_data = np.load(os.path.join(args.fmri_dir, obj + ".npy"))
        fmri_data = torch.from_numpy(fmri_data[2:8]).unsqueeze(0).cuda()
        with torch.no_grad():
            fmri_embedding = model_full.get_fmri_embedding_batch(fmri_data, diffusion_bs=128, decoder_bs=1)
            mesh_list = model_full.generate_mesh_from_embedding(fmri_embedding, bs=1, top_k=args.topk)
        if mesh_list and mesh_list[0].vertices.shape[0]:
            mesh_list[0].export(out_path)
        print(f"[{i + 1}/{len(ids)}] {obj} {time.time() - t0:.1f}s", flush=True)