import streamlit as st
import tensorflow as tf
from transformers import BertTokenizer, TFBertForSequenceClassification
import numpy as np
import os

# --- Page Configuration ---
st.set_page_config(
    page_title="Disaster Tweet Classifier",
    page_icon="🚨",
    layout="centered"
)

# --- Title and Description ---
st.title("🚨 Disaster Tweet Classifier")
st.markdown("""
This AI application determines whether a tweet is about a **real disaster** (e.g., earthquake, fire) 
or just uses **metaphorical language** (e.g., 'This movie was a disaster').
\n**Powered by:** BERT (Bidirectional Encoder Representations from Transformers)
""")

# --- Model Loading Logic ---
@st.cache_resource
def load_model():
    # model files are now in the same folder as this script (src/)
    model_path = "src/"

    try:
        tokenizer = BertTokenizer.from_pretrained(model_path)
        model = TFBertForSequenceClassification.from_pretrained(model_path, from_pt=False)
        return tokenizer, model
    except Exception as e:
        st.error(f"❌ Error loading model: {e}")
        st.info("Make sure config.json, tf_model.h5, vocab.txt and tokenizer files are in the same folder.")
        return None, None

# Load model
with st.spinner("Loading AI Model... Please wait..."):
    tokenizer, model = load_model()

# --- Prediction Logic ---
def predict_tweet(text, tokenizer, model, max_len=80):
    encoded = tokenizer.encode_plus(
        text,
        add_special_tokens=True,
        max_length=max_len,
        padding='max_length',
        return_attention_mask=True,
        truncation=True
    )

    input_ids = tf.convert_to_tensor([encoded['input_ids']])
    attention_mask = tf.convert_to_tensor([encoded['attention_mask']])

    # Forward pass
    outputs = model([input_ids, attention_mask])
    logits = outputs.logits

    # Softmax
    probs = tf.nn.softmax(logits, axis=1).numpy()[0]

    pred_class = int(np.argmax(probs))
    confidence = float(probs[pred_class])

    return pred_class, confidence

# --- User Interface ---
st.divider()
user_input = st.text_area("Enter a tweet below:", height=100, placeholder="Example: The forest fire is spreading rapidly!")

col1, col2 = st.columns([1, 2])
with col1:
    analyze_button = st.button("Analyze Tweet", type="primary", use_container_width=True)

if analyze_button:
    if not user_input.strip():
        st.warning("⚠️ Please enter some text first.")
    elif model is None:
        st.error("Model failed to load.")
    else:
        pred_class, confidence = predict_tweet(user_input, tokenizer, model)

        st.divider()
        if pred_class == 1:
            st.error(f"### 🚨 Prediction: REAL DISASTER")
            st.write(f"Confidence Score: **%{confidence * 100:.2f}**")
            st.progress(int(confidence * 100))
        else:
            st.success(f"### 😌 Prediction: NOT A DISASTER")
            st.write(f"Confidence Score: **%{confidence * 100:.2f}**")
            st.progress(int(confidence * 100))