import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
import sys
import tempfile
from pathlib import Path
import gradio as gr
import numpy as np
import trimesh
import torch
from sklearn.cluster import KMeans
from huggingface_hub import hf_hub_download

# Setup paths for partfield
ROOT = Path(__file__).parent
sys.path.append(str(ROOT / "partfield"))
CKPT_DIR = ROOT / "partfield" / "model"
CKPT_PATH = CKPT_DIR / "model_objaverse.ckpt"

from partfield.config import default_argument_parser, setup
from partfield.model_trainer_pvcnn_only_demo import Model

def ensure_checkpoint():
    CKPT_DIR.mkdir(parents=True, exist_ok=True)
    if CKPT_PATH.is_file():
        return str(CKPT_PATH)
    print("Downloading PartField checkpoint (~1.24 GB)...")
    downloaded = hf_hub_download(
        repo_id="mikaelaangel/partfield-ckpt",
        filename="model_objaverse.ckpt",
        token=os.environ.get("HF_TOKEN"),
        local_dir=str(CKPT_DIR),
    )
    return downloaded

ensure_checkpoint()

# Initialize Model Configuration
CONFIG_FILE = ROOT / "partfield" / "configs" / "final" / "demo.yaml"
parser = default_argument_parser()
args = parser.parse_args([
    "-c", str(CONFIG_FILE),
    "--opts",
    "continue_ckpt", str(CKPT_PATH),
    "result_name", "gradio_output",
])
cfg = setup(args, freeze=False)

# Load Model
print("Loading PartField model...")
model = Model(cfg)
state_dict = torch.load(CKPT_PATH, map_location="cpu")["state_dict"]
model.load_state_dict(state_dict, strict=False)
if torch.cuda.is_available():
    model = model.cuda()
model.eval()
print("Model loaded successfully.")

def process_mesh(mesh_file, num_clusters=5):
    try:
        import gc
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            
        num_clusters = int(num_clusters)
        import uuid
        uid = f"gradio_upload_{uuid.uuid4().hex[:8]}"
        
        # Load mesh
        mesh = trimesh.load(mesh_file.name, force='mesh')
        
        # Fix quads if necessary
        if mesh.faces.shape[1] == 4:
            new_faces = []
            for face in mesh.faces:
                new_faces.append([face[0], face[1], face[2]])
                new_faces.append([face[0], face[2], face[3]])
            mesh.faces = np.array(new_faces)

        # Scale and Center (Crucial for PartField)
        vertices = mesh.vertices
        bbmin = vertices.min(0)
        bbmax = vertices.max(0)
        center = (bbmin + bbmax) * 0.5
        scale = 2.0 * 0.9 / (bbmax - bbmin).max()
        vertices = (vertices - center) * scale
        mesh.vertices = vertices

        # Sample point cloud
        pc, _ = trimesh.sample.sample_surface(mesh, 20000)

        # Prepare Batch
        device = "cuda" if torch.cuda.is_available() else "cpu"
        batch = {
            'uid': [uid],
            'pc': torch.tensor(pc, dtype=torch.float32).unsqueeze(0).to(device),
            'vertices': [torch.tensor(mesh.vertices, dtype=torch.float32).to(device)],
            'faces': [torch.tensor(mesh.faces, dtype=torch.int64).to(device)]
        }

        # Run inference (will save features to disk inside exp_results)
        save_dir = f"exp_results/{cfg.result_name}"
        feat_path_1 = f'{save_dir}/part_feat_{uid}_0.npy'
        feat_path_2 = f'{save_dir}/part_feat_{uid}_0_batch.npy'
        pca_path = f'{save_dir}/feat_pca_{uid}_0.ply'
        if os.path.exists(feat_path_1): os.remove(feat_path_1)
        if os.path.exists(feat_path_2): os.remove(feat_path_2)

        with torch.no_grad():
            model.predict_step(batch, 0)

        # Load Features
        feat_path = feat_path_2 if os.path.exists(feat_path_2) else feat_path_1
        if not os.path.exists(feat_path):
            return None
        
        point_feat = np.load(feat_path)
        point_feat = point_feat / np.linalg.norm(point_feat, axis=-1, keepdims=True)

        # Cluster Features
        clustering = KMeans(n_clusters=num_clusters, random_state=0).fit(point_feat)
        labels = clustering.labels_

        # Split Mesh by Cluster Labels
        submeshes = []
        for i in range(num_clusters):
            face_mask = (labels == i)
            sub_faces = mesh.faces[face_mask]
            if len(sub_faces) > 0:
                submesh = trimesh.Trimesh(vertices=mesh.vertices, faces=sub_faces, process=True)
                # Ensure the submesh centers back to original space? 
                # Leaving in normalized space is fine, the Blender add-on can handle scale
                submeshes.append(submesh)

        scene = trimesh.Scene(submeshes)
        out_path = tempfile.mktemp(suffix=".glb")
        scene.export(out_path)
        
        return out_path
    except Exception as e:
        print("Error processing mesh:", e)
        return None
    finally:
        # Clean up temporary disk files
        for p in [feat_path_1, feat_path_2, pca_path]:
            try:
                if 'p' in locals() and os.path.exists(p):
                    os.remove(p)
            except Exception:
                pass

        import gc
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

with gr.Blocks() as demo:
    gr.Markdown("# PartField Space — API Endpoint")
    with gr.Row():
        inp_file = gr.File(label="Input 3D Mesh (.glb/.obj)")
        inp_clusters = gr.Number(value=5, label="Number of Parts")
    out_file = gr.File(label="Segmented Mesh (.glb)")
    gr.Button("Segment Mesh").click(process_mesh, inputs=[inp_file, inp_clusters], outputs=[out_file], api_name="/segment")

demo.queue(default_concurrency_limit=1).launch(ssr_mode=False)