import streamlit as st
from transformers import pipeline
import torch
import plotly.express as px
import pandas as pd
import time

# --- UI CONFIGURATION ---
st.set_page_config(page_title="NEURAL-X | AI Classifier", layout="wide", initial_sidebar_state="expanded")

# --- CUSTOM CSS ---
st.markdown("""
    <style>
    .stApp { background: radial-gradient(circle at top right, #1e293b, #0f172a); color: #f8fafc; }
    h1 { font-family: 'Inter', sans-serif; font-weight: 800; background: -webkit-linear-gradient(#38bdf8, #818cf8); -webkit-background-clip: text; -webkit-text-fill-color: transparent; letter-spacing: -1px; }
    
    div.stTextArea { width: 100% !important; box-sizing: border-box !important; }
    div[data-baseweb="base-input"], div[data-baseweb="textarea"] { background-color: #1e293b !important; border: 1px solid rgba(255, 255, 255, 0.2) !important; border-radius: 15px !important; }
    textarea { color: #ffffff !important; -webkit-text-fill-color: #ffffff !important; caret-color: #ffffff !important; }

    [data-testid="stMetricValue"] { color: #38bdf8; font-size: 2rem !important; white-space: normal !important; word-wrap: break-word !important; line-height: 1.2; }
    
    .stButton > button { width: 100%; background: linear-gradient(90deg, #6366f1 0%, #a855f7 100%); color: white; border: none; padding: 0.75rem; border-radius: 12px; font-weight: 600; transition: all 0.3s ease; text-transform: uppercase; letter-spacing: 1px; }
    .stButton > button:hover { transform: translateY(-2px); box-shadow: 0 10px 20px rgba(99, 102, 241, 0.4); }
    </style>
    """, unsafe_allow_html=True)

# --- AI BACKEND ---
@st.cache_resource(show_spinner=False)
def load_ai():
    return pipeline("zero-shot-classification", model="cross-encoder/nli-distilroberta-base", device=-1)

classifier = load_ai()

HIDDEN_LABELS = [
    "Business", "Macroeconomics", "Startups", "Stock Market", "Politics", "Geopolitics", 
    "War & Conflict", "Human Rights", "Technology", "Artificial Intelligence", 
    "Cybersecurity", "Science", "Healthcare", "Climate Change", "Entertainment", 
    "Music", "Video Games", "Sports", "Crime", "Religion", "Psychology", 
    "Social Issues", "Mental Health", "Youth & Development", "Education"
]

# --- SIDEBAR (History) ---
with st.sidebar:
    st.title("🗂️ Analysis History")
    if 'history' not in st.session_state:
        st.session_state.history = []
    
    for item in reversed(st.session_state.history):
        st.info(f"**{item['label']}**\n\n{item['text'][:50]}...")

# --- MAIN UI ---
st.title("NEURAL-X CLASSIFIER")
st.write("#### Enterprise-grade semantic analysis powered by Deep Learning.")

col1, col2 = st.columns([1.2, 0.8], gap="large")

with col1:
    text_input = st.text_area("INPUT STREAM", height=250, placeholder="Paste raw text here...")
    analyze_btn = st.button("Execute Neural Scan")

with col2:
    if analyze_btn:
        if not text_input.strip():
            st.warning("⚠️ Neural input buffer empty.")
        else:
            with st.spinner("Processing Tensors..."):
                start_time = time.time()
                result = classifier(text_input, HIDDEN_LABELS, multi_label=True)
                end_time = time.time()
                
                # Dynamic Logic with 60% Threshold
                valid_pairs = [(l, s) for l, s in zip(result['labels'], result['scores']) if s >= 0.60]
                
                # THE REVERT: Blindly guessing the highest score if nothing hits 60%
                if not valid_pairs:
                    valid_pairs = [(result['labels'][0], result['scores'][0])]
                
                st.subheader("Analysis Results")
                m1, m2 = st.columns(2)
                
                m1.metric("Top Category", valid_pairs[0][0])
                m2.metric("Inference Time", f"{round(end_time - start_time, 2)}s")
                
                df = pd.DataFrame({
                    "Concept": [p[0] for p in valid_pairs],
                    "Confidence": [p[1] for p in valid_pairs]
                }).sort_values("Confidence")
                
                fig = px.bar(df, x="Confidence", y="Concept", orientation='h',
                             color="Confidence", color_continuous_scale="Tealgrn",
                             template="plotly_dark")
                             
                fig.update_yaxes(automargin=True)
                fig.update_layout(paper_bgcolor='rgba(0,0,0,0)', plot_bgcolor='rgba(0,0,0,0)')
                st.plotly_chart(fig, use_container_width=True)
                
                st.session_state.history.append({"label": valid_pairs[0][0], "text": text_input})
    else:
        st.info("💡 Enter text and initiate scan to view semantic mapping.")