import streamlit as st
import cv2
import numpy as np
from PIL import Image
from io import BytesIO
from openvino.runtime import Core
import random

st.set_page_config(page_title="Anime & Cartoon Stylizer", layout="wide")
st.title("🖼️ Anime & Cartoon Stylizer with OpenVINO")

persona = st.radio("Choose your vibe:", ["Classic", "ChatBoy 💘", "Poetic 🌸"], horizontal=True)
style = st.selectbox("Choose your style:", ["AnimeGAN", "CartoonGAN"])

@st.cache_resource
def load_model(style_name):
    ie = Core()
    model_path = "animegan.xml" if style_name == "AnimeGAN" else "cartoongan.xml"
    model = ie.read_model(model=model_path)
    compiled_model = ie.compile_model(model=model, device_name="CPU")
    return compiled_model, compiled_model.input(0), compiled_model.output(0)

def preprocess(image: Image.Image):
    img = np.array(image.resize((256, 256))).astype(np.float32)
    img = img / 127.5 - 1.0
    img = np.transpose(img, (2, 0, 1))
    return np.expand_dims(img, axis=0)

def postprocess(output):
    result = output.squeeze().transpose(1, 2, 0)
    result = (result + 1.0) * 127.5
    return np.clip(result, 0, 255).astype(np.uint8)

def get_compliment():
    return random.choice([
        "✨ You look like the protagonist of a dreamy anime romance.",
        "💘 That transformation? Utterly magical.",
        "🌸 Your photo just bloomed into a masterpiece.",
        "🎀 If Studio Ghibli saw this, they'd cast you instantly.",
        "🫶 This anime version of you? It's giving main character energy."
    ])

uploaded_file = st.file_uploader("Upload a photo", type=["jpg", "jpeg", "png"])

if uploaded_file:
    image = Image.open(uploaded_file).convert("RGB")
    st.subheader("📸 Original vs Stylized")
    col1, col2 = st.columns(2)
    col1.image(image, caption="Original", use_column_width=True)

    if st.button("Generate Style"):
        compiled_model, input_layer, output_layer = load_model(style)
        input_tensor = preprocess(image)

        # ✅ Use blob names for inference:
        input_name = input_layer.get_any_name()
        output_name = output_layer.get_any_name()

        # Run inference
        result_dict = compiled_model.infer({input_name: input_tensor})
        output = result_dict[output_name]

        result = postprocess(output)
        result_pil = Image.fromarray(result)
        col2.image(result_pil, caption=f"{style} Style", use_column_width=True)

        buf = BytesIO()
        result_pil.save(buf, format="PNG")
        st.download_button(
            "Download Stylized Image",
            data=buf.getvalue(),
            file_name=f"{style.lower()}_style.png",
            mime="image/png"
        )

        # Persona reaction
        if persona == "ChatBoy 💘":
            st.markdown(f"**{get_compliment()}**")
        elif persona == "Poetic 🌸":
            st.markdown("🌙 *Your image now dances in the moonlight of a painted dream.*")