import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

import argparse
import json
import yaml
import torch
import numpy as np
from data.smal_quadrupeds import SMALQuadrupedGenerator, PARENT_INDICES
from pipeline.hybrid_splicer import KinematicHybridSplicer
from models.point_dit import PointDiT
from models.flow_matching import OptimalTransportFlowMatching

def main():
    parser = argparse.ArgumentParser(description="Sample creatures and export flow trajectories for the viewer.")
    parser.add_argument("--config", type=str, default="configs/default_config.yaml")
    parser.add_argument("--checkpoint", type=str, default=None)
    args = parser.parse_args()

    with open(args.config, "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f)
        
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = PointDiT(
        point_dim=cfg["model"]["point_dim"],
        hidden_dim=cfg["model"]["hidden_dim"],
        num_layers=cfg["model"]["num_layers"],
        num_heads=cfg["model"]["num_heads"],
        num_joints=cfg["model"]["num_joints"],
        spatial_distance_scale=cfg["model"]["spatial_distance_scale"],
        fourier_octaves=cfg["model"].get("fourier_octaves", 0)
    )
    
    ckpt_path = args.checkpoint or os.path.join(cfg["training"]["checkpoint_dir"],
                                                f"model_epoch_{cfg['training']['num_epochs']:03d}.pt")
    if os.path.exists(ckpt_path):
        ckpt = torch.load(ckpt_path, map_location=device, weights_only=True)
        model.load_state_dict(ckpt.get("ema_state", ckpt["model_state"]))
        print(f"Loaded checkpoint from {ckpt_path}")
    else:
        print("Note: No trained checkpoint found yet. Using freshly initialized model for demonstration.")
        
    model.to(device)
    fm = OptimalTransportFlowMatching(model)
    generator = SMALQuadrupedGenerator(seed=42)
    splicer = KinematicHybridSplicer()
    
    specimens = [
        {"name": "Lion (Felidae)", "type": "pure", "family": "felidae"},
        {"name": "Horse (Equidae)", "type": "pure", "family": "equidae"},
        {"name": "Wolf (Canidae)", "type": "pure", "family": "canidae"},
        {"name": "Bison (Bovidae)", "type": "pure", "family": "bovidae"},
        {"name": "Chimera (Lion Body + Horse Neck/Head)", "type": "hybrid", "host": "felidae", "donor": "equidae", "junction": "neck_base"}
    ]
    
    exported = []
    num_points = cfg["data"]["num_points"]
    steps = cfg["sampling"]["num_steps"]
    solver = cfg["sampling"]["solver"]
    
    for spec in specimens:
        print(f"Generating {spec['name']}...")
        if spec["type"] == "pure":
            joints = generator.generate_instance(spec["family"])["joints"]
        else:
            skel_h = generator.generate_instance(spec["host"])["joints"]
            skel_d = generator.generate_instance(spec["donor"])["joints"]
            joints = splicer.splice(skel_h, skel_d, junction_joint=spec["junction"])
            
        j_tensor = torch.from_numpy(joints).unsqueeze(0).to(device)
        traj, _ = fm.sample(j_tensor, num_points=num_points, num_steps=steps, solver=solver, record_trajectory=True,
                            guidance_scale=cfg["sampling"].get("guidance_scale", 1.0))
        
        exported.append({
            "name": spec["name"],
            "type": spec["type"],
            "joints": joints.tolist(),
            "joint_parents": PARENT_INDICES,
            "trajectory": [step.squeeze(0).numpy().tolist() for step in traj]
        })
        
    out_dir = cfg["sampling"]["export_dir"]
    os.makedirs(out_dir, exist_ok=True)
    out_file = os.path.join(out_dir, "creatures.json")
    with open(out_file, "w", encoding="utf-8") as f:
        json.dump({"creatures": exported}, f)
    print(f"Exported {len(exported)} creature trajectories to {out_file}")

if __name__ == "__main__":
    main()