#!/usr/bin/env python3
# Busca organelas REAIS de uma célula HeLa inteira, reconstruída em 3D a
# partir de FIB-SEM (microscopia eletrônica volumétrica) e segmentada por
# machine learning pelo projeto CellMap/COSEM do Janelia Research Campus
# (HHMI) — publicado como OpenOrganelle, CC BY 4.0.
# https://openorganelle.janelia.org/  |  dataset "jrc_hela-2"
# DOI da coleção: https://doi.org/10.25378/janelia.c.7456966
#
# POR QUE ISSO E NÃO GEOMETRIA PROCEDURAL: é o mesmo princípio já usado pro
# resto do atlas (malhas reais do Z-Anatomy) — aqui não existe um "Z-Anatomy
# da célula" pronto pra converter de FBX, mas o Janelia publica a segmentação
# 3D real como malhas no formato Neuroglancer "precomputed" (multi-resolução,
# comprimido com Draco). Este script usa a biblioteca `cloud-volume` (Seung
# Lab) pra baixar e decodificar essas malhas direto do bucket público, sem
# fabricar nenhuma forma.
#
# Escolhas de curadoria (documentadas pra poder ajustar depois):
#   - Cada organela é uma "camada" de segmentação própria no dataset
#     (nucleus_seg, mito_seg, er_seg, golgi_seg, pm_seg), com um ID por
#     instância. Pegamos a maior instância de núcleo/membrana (a célula
#     mais completa no campo de visão) e TODAS as instâncias de
#     mitocôndria/RE/Golgi (são realmente múltiplas organelas de verdade,
#     não fragmentos de ruído — juntar todas reproduz a distribuição real).
#   - LOD (nível de detalhe) escolhido por organela: mais fino pra
#     estruturas pequenas/poucas (núcleo, Golgi), mais grosso pras que têm
#     MUITA geometria agregada (mitocôndrias, RE) — o pipeline Node.js
#     (build-cell-model.mjs) simplifica/comprime tudo de qualquer forma.
#   - Todas as peças são recentralizadas usando o MESMO deslocamento (o
#     centro da membrana plasmática) — cada organela NÃO pode ser
#     recentralizada individualmente, isso destruiria a posição relativa
#     real dentro da célula.
#   - Escala: nanômetros (unidade nativa do dataset) → micrômetros (÷1000),
#     só por conveniência numérica — a câmera do app se ajusta ao tamanho
#     real do conteúdo de qualquer forma (ver Scene.tsx/CameraController).
#
# Vesículas (~40 mil instâncias) e endossomos (~5.600) ficaram de fora
# desta primeira versão por volume de dados — dá pra adicionar depois como
# mais uma organela seguindo exatamente o mesmo padrão abaixo.
import os
import time

import numpy as np
import cloudvolume as cv
import trimesh

DATASET = "jrc_hela-2"
BASE = f"precomputed://https://janelia-cosem-datasets.s3.amazonaws.com/{DATASET}/neuroglancer/mesh"
OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "converted", "citologia")

# (nome do nó no glb, camada de segmentação, IDs de instância, LOD)
GROUPS = [
    ("celula-membrana-plasmatica", "pm_seg", [2], 4),
    ("celula-nucleo", "nucleus_seg", [1], 3),
    ("celula-reticulo-endoplasmatico", "er_seg", list(range(1, 25)), 4),
    ("celula-complexo-de-golgi", "golgi_seg", list(range(1, 9)), 3),
    ("celula-mitocondrias", "mito_seg", list(range(1, 422)), 4),
]

SCALE_NM_TO_UM = 1.0 / 1000.0


def make_volume(seg_layer):
    # CloudVolume espera o "info" completo de uma camada de imagem/
    # segmentação (com "scales"), mas só precisamos da parte de malha —
    # os valores de "scales" abaixo são um placeholder mínimo válido só
    # pra passar da checagem interna da biblioteca; a resolução REAL usada
    # pra posicionar os vértices vem do "transform" dentro do info.json de
    # cada camada de malha (lido automaticamente pelo cloud-volume).
    info = {
        "type": "segmentation",
        "data_type": "uint64",
        "num_channels": 1,
        "scales": [{
            "chunk_sizes": [[64, 64, 64]],
            "encoding": "raw",
            "key": "1_1_1",
            "resolution": [4, 4, 5.24],
            "size": [1, 1, 1],
            "voxel_offset": [0, 0, 0],
        }],
        "mesh": seg_layer,
    }
    return cv.CloudVolume(BASE, info=info, use_https=True)


def fetch_group(name, seg_layer, ids, lod):
    t0 = time.time()
    vol = make_volume(seg_layer)
    meshes = vol.mesh.get(ids, lod=lod)
    parts = []
    for m in meshes.values():
        if m is None or len(m.vertices) == 0 or len(m.faces) == 0:
            continue
        # A decodificação Draco em algumas LODs mais grosseiras às vezes
        # produz vértices NaN/Inf pra fragmentos degenerados — zera-los
        # em vez de excluir a organela inteira mantém a peça utilizável.
        vertices = np.nan_to_num(m.vertices, nan=0.0, posinf=0.0, neginf=0.0)
        parts.append(trimesh.Trimesh(vertices=vertices, faces=m.faces, process=False))
    if not parts:
        raise RuntimeError(f"{name}: nenhuma geometria válida encontrada (verifique os IDs/LOD)")
    combined = trimesh.util.concatenate(parts)
    print(f"  {name}: {len(combined.vertices)} vértices, {len(combined.faces)} faces ({time.time() - t0:.1f}s)", flush=True)
    return combined


def main():
    print(f"Baixando organelas reais de {DATASET} (Janelia OpenOrganelle/CellMap, CC BY 4.0)...")
    results = {}
    for name, seg_layer, ids, lod in GROUPS:
        results[name] = fetch_group(name, seg_layer, ids, lod)

    center = results["celula-membrana-plasmatica"].bounds.mean(axis=0)

    scene = trimesh.Scene()
    for name, mesh in results.items():
        mesh.vertices = (mesh.vertices - center) * SCALE_NM_TO_UM
        scene.add_geometry(mesh, node_name=name, geom_name=name)

    os.makedirs(OUT_DIR, exist_ok=True)
    out_path = os.path.join(OUT_DIR, "citologia-raw.glb")
    scene.export(out_path)
    print(f"✓ {out_path} salvo — pronto para o build-cell-model.mjs simplificar/comprimir.")


if __name__ == "__main__":
    main()