app.py

import streamlit as st
import pandas as pd
import numpy as np
import joblib
import matplotlib.pyplot as plt
import seaborn as sns
from sentence_transformers import SentenceTransformer
from google import genai  

# Import custom extraction logic from your helper file
from feature_engine import extract_stylometric_features, extract_linguistic_features

if "GEMINI_API_KEY" in st.secrets:
    ai_client = genai.Client(api_key=st.secrets["GEMINI_API_KEY"])
else:
    st.error("GEMINI_API_KEY missing from .streamlit/secrets.toml")


# 1. PAGE CONFIGURATION
st.set_page_config(
    page_title="ModelTrace | AI Forensics",
    page_icon="🕵️‍♂️",
    layout="wide"
)


# 2. ASSET CACHING
@st.cache_resource
def load_engine():
    ensemble = joblib.load("modeltrace_exports/production_ensemble.pkl")
    scaler = joblib.load("modeltrace_exports/feature_scaler.pkl")
    pca = joblib.load("modeltrace_exports/semantic_pca.pkl")
    feature_names = joblib.load("modeltrace_exports/feature_names.pkl")
    emb_model = SentenceTransformer('all-MiniLM-L6-v2')
    return ensemble, scaler, pca, feature_names, emb_model

ensemble, scaler, pca, feature_names, emb_model = load_engine()


# 3. SIDEBAR & UI HEADER
with st.sidebar:
    st.title("ModelTrace Dashboard")
    st.markdown("### Production Engine v1.1")
    st.info("Accuracy: 83.2%")
    st.divider()
    st.markdown("**Detecting Origins For:**")
    st.caption("GPT-4o, Llama-3, Qwen-2.5, DeepSeek, Gemma")

st.title("ModelTrace: Explainable AI Forensics")
st.markdown("Identify the generative origin of any text and receive an automated forensic breakdown.")


# 4. USER INPUT
raw_text = st.text_area("Paste Text for Analysis", height=250, placeholder="Minimum 50 words recommended...")
analyze_btn = st.button("Run Forensic Analysis", type="primary", use_container_width=True)


# 5. INFERENCE & XAI PIPELINE
if analyze_btn and raw_text:
    with st.spinner("Extracting 88-dimensional fingerprint..."):
        
        # --- A. FEATURE EXTRACTION ---
        stylo = extract_stylometric_features(raw_text)
        ling = extract_linguistic_features(raw_text)
        
        raw_emb = emb_model.encode([raw_text])
        sem_feats = pca.transform(raw_emb)[0].tolist()
        
        all_features = stylo + ling + sem_feats
        feature_df = pd.DataFrame([all_features], columns=feature_names)
        scaled_features = pd.DataFrame(scaler.transform(feature_df), columns=feature_names)
        
        # --- B. CLASSIFICATION ---
        prediction = ensemble.predict(scaled_features)[0]
        probs = ensemble.predict_proba(scaled_features)[0]
        
        # --- C. XAI EXTRACTION ---
        lgb_submodel = ensemble.named_estimators_['lgb']
        raw_contribs = lgb_submodel.booster_.predict(scaled_features, pred_contrib=True)
        
        num_classes = len(ensemble.classes_)
        num_feats = len(feature_names)
        reshaped = raw_contribs[0].reshape(num_classes, num_feats + 1)
        
        class_idx = list(ensemble.classes_).index(prediction)
        winning_contribs = reshaped[class_idx][:-1]
        
        sorted_reasons = sorted(dict(zip(feature_names, winning_contribs)).items(), 
                                key=lambda x: abs(x[1]), reverse=True)[:5]

    
    # 6. RESULTS DISPLAY
    st.divider()
    res_col1, res_col2 = st.columns(2)
    
    with res_col1:
        st.subheader("Attribution Result")
        st.success(f"**Predicted Origin:** {prediction.upper()}")
        
        prob_df = pd.DataFrame({"Model": ensemble.classes_, "Match %": probs})
        st.bar_chart(prob_df.set_index("Model"))

    with res_col2:
        st.subheader("Mathematical Drivers")
        fig, ax = plt.subplots()
        feats = [r[0] for r in sorted_reasons]
        impacts = [r[1] for r in sorted_reasons]
        sns.barplot(x=impacts, y=feats, palette="viridis", ax=ax)
        st.pyplot(fig)

    
    # 7. AUTOMATED ANALYST REPORT 
    st.divider()
    st.subheader("Automated AI Analyst Report")
    st.info("Generating a plain English forensic breakdown with gemini-2.5-flash...")
    
    feature_summary = "\n".join([f"- {r[0]}: Impact {r[1]:.4f}" for r in sorted_reasons])
    
    prompt = f"""
    You are an AI Forensics Expert. Our model predicts this text belongs to {prediction}.
    Explain why based on these top driving features:
    {feature_summary}
    Write 2 concise paragraphs. Translate technical terms into plain English for a security team.
    """
    
    try:
        response = ai_client.models.generate_content(
            model='gemini-2.5-flash',
            contents=prompt,
        )
        st.write(response.text)
    except Exception as e:
        try:
            response = ai_client.models.generate_content(
                model='gemini-1.5-flash',
                contents=prompt,
            )
            st.write(response.text)
        except Exception as fallback_error:
            st.error(f"Could not generate AI report: {fallback_error}")