import streamlit as st
import torch
import torchvision.transforms as transforms
from torchvision.models.mobilenetv2 import MobileNetV2
from PIL import Image
import numpy as np
import cv2
from datetime import datetime
from gtts import gTTS
import os

# ====== Setup ======
os.makedirs("history", exist_ok=True)
torch.serialization.add_safe_globals({'MobileNetV2': MobileNetV2})

# Load model
model = torch.load("pkr_currency_classifier.pt", map_location='cpu', weights_only=False)
model.eval()
class_names = ['Fake', 'Not Currency', 'Real']

# Transforms
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])

# Text-to-speech
def speak_streamlit(text):
    tts = gTTS(text=text, lang='en')
    tts.save("temp.mp3")
    st.audio("temp.mp3", format="audio/mp3")

# Prediction
def predict(image):
    img = transform(image).unsqueeze(0)
    with torch.no_grad():
        outputs = model(img)
        probs = torch.nn.functional.softmax(outputs, dim=1)
        confidence, predicted = torch.max(probs, 1)
        print("Probabilities:", probs.numpy())
        print("Confidence:", confidence.item())
    return class_names[predicted.item()]


# Save history
def save_history(image, result):
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    image.save(f"history/{timestamp}_{result}.png")

# ====== PAGE CONFIG & CUSTOM STYLING ======
st.set_page_config(page_title="Currency Authenticity Detector", page_icon="💵", layout="centered")

# Custom CSS
st.markdown("""
    <style>
        .main {
            background-color: #f0f2f6;
        }
        .stButton button {
            background-color: #2e86de;
            color: white;
            border-radius: 8px;
            padding: 0.5em 1.5em;
            margin: 0.3em 0;
            transition: all 0.3s ease-in-out;
        }
        .stButton button:hover {
            background-color: #1b4f72;
            color: white;
            transform: scale(1.05);
        }
        .result-box {
            font-size: 20px;
            padding: 10px;
            border-radius: 10px;
            background-color: #eaf2f8;
            text-align: center;
        }
        .prediction-fake {
            color: red;
            font-weight: bold;
        }
        .prediction-real {
            color: green;
            font-weight: bold;
        }
        .prediction-unknown {
            color: orange;
            font-weight: bold;
        }
    </style>
""", unsafe_allow_html=True)

# ====== UI Layout ======
st.markdown("# 💵 Currency Authenticity Detector")
st.markdown("### ✅ Instantly check if your currency is **Real or Fake**!")

currency_type = st.selectbox("🔄 Select currency type:", ["PKR (Pakistani Rupees)", "USD (US Dollars)", "INR (Indian Rupees)"])
if "PKR" not in currency_type:
    st.warning("⚠️ Currently only Pakistani Rupees (PKR) is supported. Other currencies coming soon!")
    st.stop()

st.markdown("### 🖼️ Choose how to scan your currency note:")

# Choose input method
option = st.radio("🔍 Input Method:", ["Upload Image", "Scan via Camera"])

# ===== Upload Image =====
if option == "Upload Image":
    uploaded_file = st.file_uploader("📤 Upload a currency image", type=["jpg", "jpeg", "png"])
    if uploaded_file:
        image = Image.open(uploaded_file).convert("RGB")
        st.image(image, caption="🖼️ Uploaded Image", use_column_width=True)
        prediction = predict(image)

        if prediction == "Not Currency":
            st.markdown('<div class="result-box prediction-unknown">⚠️ This does not appear to be a currency note.</div>', unsafe_allow_html=True)
            speak_streamlit("This is not a currency note.")
        elif prediction == "Fake":
            st.markdown(f'<div class="result-box prediction-fake">❌ Prediction: {prediction} Currency</div>', unsafe_allow_html=True)
            speak_streamlit("This is a fake currency note.")
            save_history(image, prediction)
        else:
            st.markdown(f'<div class="result-box prediction-real">✔️ Prediction: {prediction} Currency</div>', unsafe_allow_html=True)
            speak_streamlit("This is a real currency note.")
            save_history(image, prediction)

# ===== Webcam Scan =====
elif option == "Scan via Camera":
    st.markdown("🎥 Press the button below to start your webcam.")
    start_camera = st.button("📷 Start Camera")
    if start_camera:
        cap = cv2.VideoCapture(0)
        stframe = st.empty()
        result_box = st.empty()
        stop_button = st.button("🛑 Stop Camera")

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            pil_image = Image.fromarray(frame_rgb)
            stframe.image(pil_image, channels="RGB", use_column_width=True)

            prediction = predict(pil_image)

            if prediction == "Not Currency":
                result_box.markdown('<div class="result-box prediction-unknown">⚠️ This does not appear to be a currency note.</div>', unsafe_allow_html=True)
                speak_streamlit("This is not a currency note.")
            elif prediction == "Fake":
                result_box.markdown(f'<div class="result-box prediction-fake">❌ Prediction: {prediction} Currency</div>', unsafe_allow_html=True)
                speak_streamlit("This is a fake currency note.")
                save_history(pil_image, prediction)
            else:
                result_box.markdown(f'<div class="result-box prediction-real">✔️ Prediction: {prediction} Currency</div>', unsafe_allow_html=True)
                speak_streamlit("This is a real currency note.")
                save_history(pil_image, prediction)

            if stop_button:
                break

        cap.release()
        stframe.empty()
        result_box.empty()

# ===== Scan History =====
if st.checkbox("📁 Show Scan History"):
    st.markdown("### 🕘 Recent Scans")
    history_files = sorted(os.listdir("history"))[::-1][:5]
    if not history_files:
        st.info("No scan history available.")
    else:
        for img_file in history_files:
            st.image(f"history/{img_file}", caption=img_file, width=250)

# ===== Footer =====
st.markdown("---")
st.markdown("<center>👨‍💻 Designed by <b>MERAJ GRAPHICS</b></center>", unsafe_allow_html=True)