import streamlit as st
import pandas as pd
import numpy as np
import joblib
from tensorflow.keras.models import load_model

from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import MinMaxScaler

st.set_page_config(page_title="Anomaly Detection", layout="wide")

def load_models():
    iso = joblib.load("iso_model.pkl")
    ae = load_model("autoencoder_model.h5")
    scaler = joblib.load("scaler.pkl")
    return iso, ae, scaler

def main():
    st.title("🚨 Network Traffic Anomaly Detection")
    st.markdown("Upload a CSV file and choose a model to detect anomalies in network traffic.")

    uploaded_file = st.file_uploader("Upload your CSV file", type=["csv"])

    if uploaded_file:
        iso_model, ae_model, scaler = load_models()

        df = pd.read_csv(uploaded_file)
        X = scaler.transform(df)

        iso_pred = iso_model.predict(X)
        iso_pred = np.where(iso_pred == 1, 0, 1)

        ae_recon = ae_model.predict(X)
        ae_mse = np.mean(np.power(X - ae_recon, 2), axis=1)
        ae_pred = np.where(ae_mse > 0.0096, 1, 0)

        model_choice = st.selectbox("Choose a model", ["Isolation Forest", "Autoencoder", "Hybrid"])

        if model_choice == "Isolation Forest":
            final_pred = iso_pred
        elif model_choice == "Autoencoder":
            final_pred = ae_pred
        else:
            final_pred = np.logical_or(iso_pred, ae_pred).astype(int)

        df['anomaly'] = final_pred
        st.write(df.head())

        st.subheader("📊 Anomaly Distribution")
        st.bar_chart(df['anomaly'].value_counts())

        csv = df.to_csv(index=False).encode()
        st.download_button("📥 Download Predictions", csv, "predictions.csv", "text/csv")

        st.success("✅ Prediction complete!")

if __name__ == "__main__":
    main()