# ----------------------------
# FIX FOR HUGGING FACE TIMEOUT
# ----------------------------
import os
os.environ["STREAMLIT_BROWSER_GATHER_USAGE_STATS"] = "false"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

# ----------------------------
# IMPORTS
# ----------------------------
import streamlit as st
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
from PIL import Image

# ----------------------------
# PAGE CONFIG
# ----------------------------
st.set_page_config(
    page_title="Captcha OCR",
    page_icon="🔐",
    layout="centered"
)

st.title("🔐 Captcha OCR")
st.markdown("CRNN + CTC Model Deployment")

# ----------------------------
# LOAD CHARACTERS
# ----------------------------
@st.cache_resource
def load_characters():
    with open("characters.txt", "r") as f:
        characters = list(f.read().strip())
    return characters

characters = load_characters()

charToNum = layers.StringLookup(vocabulary=characters, mask_token=None)
numToChar = layers.StringLookup(
    vocabulary=charToNum.get_vocabulary(),
    mask_token=None,
    invert=True
)

# ----------------------------
# LOAD MODEL (LAZY + SAFE)
# ----------------------------
@st.cache_resource
def load_model():
    model = keras.models.load_model(
        "ocr_model.keras",
        compile=False  # IMPORTANT for memory reduction
    )
    return model

# Lazy loading (prevents reload crash)
if "model" not in st.session_state:
    st.session_state.model = load_model()

model = st.session_state.model

# ----------------------------
# PREPROCESS FUNCTION
# ----------------------------
def preprocess_image(image):
    image = image.convert("L")  # grayscale
    image = image.resize((200, 50))  # same as training
    image = np.array(image).astype("float32") / 255.0
    image = np.expand_dims(image, axis=-1)
    image = np.transpose(image, (1, 0, 2))  # IMPORTANT (match training)
    image = np.expand_dims(image, axis=0)
    return image

# ----------------------------
# DECODE FUNCTION
# ----------------------------
def decode_prediction(pred):
    input_len = np.ones(pred.shape[0]) * pred.shape[1]

    results = keras.backend.ctc_decode(
        pred,
        input_length=input_len,
        greedy=True
    )[0][0]

    output_text = []
    for res in results:
        res = tf.gather(res, tf.where(res != -1))
        res = tf.squeeze(res)
        text = tf.strings.reduce_join(numToChar(res)).numpy().decode("utf-8")
        output_text.append(text)

    return output_text[0]

# ----------------------------
# FILE UPLOADER
# ----------------------------
uploaded_file = st.file_uploader(
    "Upload Captcha Image",
    type=["png", "jpg", "jpeg"]
)

if uploaded_file is not None:
    image = Image.open(uploaded_file)

    st.image(image, caption="Uploaded Image", use_column_width=True)

    processed = preprocess_image(image)

    with st.spinner("Predicting..."):
        prediction = model.predict(processed)
        text = decode_prediction(prediction)

    st.success(f"🎯 Prediction: {text}")