app.py

import streamlit as st
import pandas as pd
import plotly.express as px
import joblib
import numpy as np
import os
import json

# --- PAGE CONFIG ---
st.set_page_config(
    page_title="LogAI Diagnostics (XGBoost)",
    layout="wide",
    page_icon="🛡️",
    initial_sidebar_state="expanded"
)

# --- CUSTOM CSS ---
st.markdown("""
<style>
.stApp { background-color: #0e1117; color: #ffffff; }
.main-header { font-size: 2.5rem; color: #60a5fa; font-weight: 700; }
section[data-testid="stSidebar"] { background-color: #262730; }
</style>
""", unsafe_allow_html=True)

# --- LOAD MODELS ---
@st.cache_resource
def load_models():
    vectorizer = joblib.load("tfidf-xgboost/model_v3/tfidf_vectorizer.joblib")
    source_model = joblib.load("tfidf-xgboost/model_v3/xgb_source_model.joblib")  
    content_model = joblib.load("tfidf-xgboost/model_v3/xgb_content_model.joblib") 

    with open("tfidf-xgboost/model_v3/xgb_label_maps.json") as f:
        label_maps = joblib.load(f) if f.name.endswith(".joblib") else json.load(f)

    return vectorizer, source_model, content_model, label_maps

# --- SIDEBAR ---
with st.sidebar:
    st.title("🛡️ LogAI Controls")
    st.info("Model: TF-IDF + XGBoost")

    required_files = [
        "tfidf-xgboost/model_v3/tfidf_vectorizer.joblib",
        "tfidf-xgboost/model_v3/xgb_source_model.joblib",
        "tfidf-xgboost/model_v3/xgb_content_model.joblib",
        "tfidf-xgboost/model_v3/xgb_label_maps.json"
    ]

    missing = [f for f in required_files if not os.path.exists(f)]
    if missing:
        st.error(f"❌ Missing files: {missing}")
        st.stop()
    else:
        st.success("✅ Models Loaded")

    st.markdown("---")
    conf_threshold = st.slider("Confidence Threshold", 0.0, 1.0, 0.5)


# --- MAIN ---
st.markdown('<div class="main-header">🛡️ LogAI Diagnostics</div>', unsafe_allow_html=True)
st.markdown("### Intelligent Log Classification (XGBoost)")

vectorizer, src_model, cnt_model, label_maps = load_models()

# --- FILE UPLOAD ---
uploaded_file = st.file_uploader(
    "📂 Upload Log File (CSV / TXT / LOG / JSON)",
    type=["csv", "txt", "log", "json"]
)

if uploaded_file:
    # --- READ FILE ---
    try:
        if uploaded_file.name.endswith(".csv"):
            df = pd.read_csv(uploaded_file)
        elif uploaded_file.name.endswith(".json"):
            df = pd.read_json(uploaded_file)
        else:
            lines = uploaded_file.read().decode("utf-8").splitlines()
            df = pd.DataFrame(lines, columns=["content"])

        df.columns = [c.lower() for c in df.columns]

        if "content" not in df.columns:
            df.rename(columns={df.columns[-1]: "content"}, inplace=True)

    except Exception as e:
        st.error(f"File read error: {e}")
        st.stop()

    st.write(f"**Loaded {len(df):,} logs**")
    st.dataframe(df.head(5), width='content')

    # --- RUN INFERENCE ---
    if st.button("🚀 Run AI Analysis", type="primary"):
        with st.spinner("Running SVM inference..."):
            X = vectorizer.transform(df["content"].astype(str))

            # --- SOURCE PREDICTION ---
            if hasattr(src_model, "predict_proba"):
                src_probs = src_model.predict_proba(X)
            else:
                df_scores = src_model.decision_function(X)
                src_probs = np.exp(df_scores) / np.exp(df_scores).sum(axis=1, keepdims=True)

            # --- CONTENT PREDICTION ---
            if hasattr(cnt_model, "predict_proba"):
                cnt_probs = cnt_model.predict_proba(X)
            else:
                df_scores = cnt_model.decision_function(X)
                cnt_probs = np.exp(df_scores) / np.exp(df_scores).sum(axis=1, keepdims=True)

            src_idx = np.argmax(src_probs, axis=1)
            cnt_idx = np.argmax(cnt_probs, axis=1)

            src_conf = np.max(src_probs, axis=1)
            cnt_conf = np.max(cnt_probs, axis=1)

            df["Source_Type"] = [
                label_maps["source"][str(i)] if c >= conf_threshold else "LowConfidence"
                for i, c in zip(src_idx, src_conf)
            ]

            df["Content_Type"] = [
                label_maps["content"][str(i)] if c >= conf_threshold else "LowConfidence"
                for i, c in zip(cnt_idx, cnt_conf)
            ]

            df["Source_Confidence"] = src_conf.round(3)
            df["Content_Confidence"] = cnt_conf.round(3)

        st.success("✅ Analysis Complete")

        # --- METRICS ---
        m1, m2, m3 = st.columns(3)
        m1.metric("Total Logs", len(df))
        m2.metric("Source Types", df["Source_Type"].nunique())
        m3.metric("Content Types", df["Content_Type"].nunique())

        # --- VISUALS ---
        c1, c2 = st.columns(2)

        with c1:
            st.subheader("📍 Source Distribution")
            src_counts = df["Source_Type"].value_counts().reset_index()
            src_counts.columns = ["Source", "Count"]
            fig1 = px.pie(src_counts, values="Count", names="Source", hole=0.4)
            fig1.update_layout(template="plotly_dark")
            st.plotly_chart(fig1, width='content')

        with c2:
            st.subheader("🔍 Content Distribution")
            cnt_counts = df["Content_Type"].value_counts().reset_index()
            cnt_counts.columns = ["Content", "Count"]
            fig2 = px.bar(cnt_counts, x="Count", y="Content", orientation="h")
            fig2.update_layout(template="plotly_dark")
            st.plotly_chart(fig2, width='content')

        # --- TABLE ---
        st.subheader("📋 Classified Logs")
        st.dataframe(df, width='content')

        # --- DOWNLOAD ---
        st.download_button(
            "📥 Download Results",
            df.to_csv(index=False).encode("utf-8"),
            "logai_results.csv",
            "text/csv",
            type="primary"
        )