import streamlit as st
from PIL import Image
import torch
from model import ModelColorization
from utils import process_gs_image, inverse_transform_cs

# Custom CSS for styling
st.markdown(
    """
    <style>
    .main {
        background-color: #f9f9f9;
    }
    .title {
        color: #ffffff;
        font-size: 2.5em;
        text-align: center;
        margin-bottom: 0.5em;
    }
    .subheader {
        color: #5a5a5a;
        font-size: 1.1em;
        text-align: center;
        margin-bottom: 2em;
    }
    .upload-box {
        background-color: #ffffff;
        border-radius: 10px;
        padding: 2em;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        margin-bottom: 2em;
    }
    .result-box {
        background-color: #ffffff;
        border-radius: 10px;
        padding: 2em;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        margin-top: 2em;
    }
    .stButton>button {
        background-color: #4CAF50;
        color: white;
        border-radius: 5px;
        padding: 0.5em 1em;
        font-size: 1em;
        width: 100%;
    }
    .stButton>button:hover {
        background-color: #45a049;
    }
    .group-banner {
        text-align: center;
        font-size: 1.5em;
        color: #ffffff;
        font-weight: bold;
        margin-top: 2em;
    }
    </style>
    """,
    unsafe_allow_html=True
)

# Load model
model = ModelColorization().from_pretrained("sebastiansarasti/AutoEncoderImageColorization")

# App header
st.markdown('<p class="title">🎨 Neural Image Colorizer</p>', unsafe_allow_html=True)
st.markdown('<p class="subheader">Bring black & white photos to life with AI</p>', unsafe_allow_html=True)

# Upload section
with st.container():
    st.markdown("### 📤 Upload Your Image")
    uploaded_file = st.file_uploader(
        "Choose a black & white photo...", 
        type=["jpg", "jpeg", "png"],
        label_visibility="collapsed"
    )

# Processing section
if uploaded_file is not None:
    with st.container():
        col1, col2 = st.columns(2)
        with col1:
            st.markdown("### ⬆ Original")
            original_img = Image.open(uploaded_file)
            st.image(original_img, use_container_width=True)
        
        with col2:
            st.markdown("### 🎨 Colorized")
            if st.button("✨ Colorize Image", type="primary"):
                with st.spinner("Colorizing your image..."):
                    # Process image
                    image, original_size = process_gs_image(original_img)
                    
                    # Run model
                    model.eval()
                    with torch.no_grad():
                        result = model(image)
                    
                    # Get colorized image
                    colorized_image = inverse_transform_cs(result.squeeze(0), original_size)
                    
                    # Display result
                    st.image(colorized_image, use_container_width=True)
                    st.success("Colorization complete!")

# Footer
st.markdown(
    '<p class="group-banner">Developed with  by Group 9 |  Computer Vision Project</p>',
    unsafe_allow_html=True
)