import streamlit as st
import numpy as np
import jax
import jax.numpy as jnp
from PIL import Image

from utils import load_model


def app(model_name):
    model, processor = load_model(f"koclip/{model_name}")

    st.title("Zero-shot Image Classification")
    st.markdown(
        """
        Some text goes in here.
    """
    )

    query = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
    captions = st.text_input("사용하실 캡션을 쉼표 단위로 구분해서 적어주세요", value="고양이,강아지,느티나무...")

    if st.button("질문 (Query)"):
        if query is None:
            st.error("Please upload an image query.")
        else:
            image = Image.open(query)
            st.image(image)
            # pixel_values = processor(
            #     text=[""], images=image, return_tensors="jax", padding=True
            # ).pixel_values
            # pixel_values = jnp.transpose(pixel_values, axes=[0, 2, 3, 1])
            # vec = np.asarray(model.get_image_features(pixel_values))
            captions = captions.split(",")
            inputs = processor(text=captions, images=image, return_tensors="jax", padding=True)
            inputs["pixel_values"] = jnp.transpose(
                inputs["pixel_values"], axes=[0, 2, 3, 1]
            )
            outputs = model(**inputs)
            probs = jax.nn.softmax(outputs.logits_per_image, axis=1)

            for idx, prob in sorted(enumerate(*probs), key=lambda x: x[1], reverse=True):
                st.text(f"Score: `{prob}`, {captions[idx]}")