import streamlit as st
import torch
import torchvision.transforms as T
import cv2
import numpy as np
from PIL import Image
import easyocr

st.markdown("""
    <style>
    .main {
        background: linear-gradient(135deg, #1e3c72, #2a5298);
        padding: 20px;
        border-radius: 15px;
        box-shadow: 0 8px 16px rgba(0, 0, 0, 0.2);
    }
    .title {
        font-size: 3rem;
        font-weight: bold;
        color: #ffffff;
        text-align: center;
        text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.3);
    }
    .subtitle {
        font-size: 1.2rem;
        color: #d1e8ff;
        text-align: center;
        margin-bottom: 20px;
    }
    .upload-box {
        background-color: rgba(255, 255, 255, 0.1);
        padding: 20px;
        border-radius: 10px;
        border: 2px dashed #ffffff;
    }
    .stButton>button {
        background-color: #00cc99;
        color: white;
        font-weight: bold;
        border-radius: 25px;
        padding: 10px 20px;
        transition: all 0.3s ease;
    }
    .stButton>button:hover {
        background-color: #00b386;
        transform: scale(1.05);
    }
    .image-caption {
        color: #ffffff;
        font-style: italic;
        text-align: center;
    }
    </style>
""", unsafe_allow_html=True)

@st.cache_resource
def load_model():
    try:
        model = torch.load("numberplate_model.pth", map_location=torch.device('cpu'), weights_only=False)
        model.eval()
        return model
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None

model = load_model()

@st.cache_resource
def load_ocr_reader():
    return easyocr.Reader(['en'], gpu=False)

reader = load_ocr_reader()

def detect_and_crop_numberplate(image):
    if model is None:
        return image, None
    
    transform = T.ToTensor()
    img_tensor = transform(image).unsqueeze(0)

    with torch.no_grad():
        predictions = model(img_tensor)

    boxes = predictions[0]['boxes']
    scores = predictions[0]['scores']

    image_np = np.array(image)
    cropped_plate = None
    for box, score in zip(boxes, scores):
        if score > 0.5:
            x1, y1, x2, y2 = map(int, box.tolist())
            cv2.rectangle(image_np, (x1, y1), (x2, y2), (0, 255, 153), 3)
            cv2.putText(image_np, f"{score:.2f}", (x1, y1-10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
            cropped_plate = image.crop((x1, y1, x2, y2))
            break

    return Image.fromarray(image_np), cropped_plate

def extract_text(cropped_plate):
    if cropped_plate is None:
        return "No number plate detected"
    cropped_np = np.array(cropped_plate)
    result = reader.readtext(cropped_np)
    text = " ".join([res[1] for res in result])
    return text if text else "Unable to read text"

st.markdown('<div class="main">', unsafe_allow_html=True)
st.markdown('<p class="title">🚗 Number Plate Wizard</p>', unsafe_allow_html=True)
st.markdown('<p class="subtitle">Unveil the magic of AI-powered detection!</p>', unsafe_allow_html=True)

with st.container():
    st.markdown('<div class="upload-box">', unsafe_allow_html=True)
    uploaded_file = st.file_uploader("📸 Drop your image here", type=["jpg", "png", "jpeg"], 
                                   help="Upload a clear image of a vehicle to detect its number plate!")
    st.markdown('</div>', unsafe_allow_html=True)

if uploaded_file is not None:
    image = Image.open(uploaded_file).convert("RGB")
    
    if st.button("✨ Detect Number Plate"):
        with st.spinner("🔮 Casting detection spell..."):
            result_image, cropped_plate = detect_and_crop_numberplate(image)
            detected_text = extract_text(cropped_plate)
        
        st.markdown("<h3 style='color: #ffffff; text-align: center;'>🎉 Detection Complete!</h3>", unsafe_allow_html=True)
        col1, col2 = st.columns(2)
        with col1:
            st.image(image, caption="Original Snapshot", use_container_width=True)
            st.markdown('<p class="image-caption">The journey begins here</p>', unsafe_allow_html=True)
        with col2:
            st.image(result_image, caption="Magic Revealed", use_container_width=True)
            st.markdown('<p class="image-caption">Number plate spotted!</p>', unsafe_allow_html=True)
        
        st.markdown(f"<h4 style='color: #ffffff; text-align: center;'>Detected Text: {detected_text}</h4>", unsafe_allow_html=True)
        
        st.success("Boom! The number plate has been summoned! 🚘")

st.markdown('<p style="color: #d1e8ff; text-align: center; margin-top: 20px;">Powered by Abhijeet Singh</p>', 
            unsafe_allow_html=True)
st.markdown('</div>', unsafe_allow_html=True)