import gradio as gr
from datasets import load_dataset
import numpy as np
from PIL import Image
from sentence_transformers import SentenceTransformer

# reuse the same grayscale conversion
def flux_to_gray(flux_array):
    a = np.array(flux_array, dtype=np.float32)
    a = np.squeeze(a)
    if a.ndim == 3:
        axis = int(np.argmin(a.shape))
        a = np.nanmean(a, axis=axis)
    a = np.nan_to_num(a, nan=0.0, posinf=0.0, neginf=0.0)
    lo = np.nanpercentile(a, 1)
    hi = np.nanpercentile(a, 99)
    if not np.isfinite(lo) or not np.isfinite(hi) or hi <= lo:
        lo, hi = float(np.nanmin(a)), float(np.nanmax(a))
    norm = np.clip((a - lo) / (hi - lo + 1e-9), 0, 1)
    arr = (norm * 255).astype(np.uint8)
    return Image.fromarray(arr, mode="L")

# load a well-known CLIP model
model = SentenceTransformer("clip-ViT-B-32")

def test_single_embedding():
    ds = load_dataset("MultimodalUniverse/jwst", split="train", streaming=True)
    rec = next(iter(ds))
    pil = flux_to_gray(rec["image"]["flux"]).convert("RGB")  # CLIP expects RGB
    emb = model.encode(pil, convert_to_numpy=True)
    info = f"OK. Image embedding shape: {emb.shape}"
    caption = f"object_id: {rec.get('object_id')}"
    return pil, caption, info

demo = gr.Interface(
    fn=test_single_embedding,
    inputs=None,
    outputs=[
        gr.Image(type="pil", label="Preview"),
        gr.Textbox(label="Info", lines=1),
        gr.Textbox(label="Embedding", lines=1),
    ],
    title="JWST → CLIP embedding check",
    description="Embeds one JWST image with CLIP to confirm the pipeline."
)

demo.launch()