import trackio
import gradio as gr
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import StandardScaler
import hashlib
import time

# ===============================
# CORPORATE LOGIN SYSTEM
# ===============================

USERNAME = "admin"
PASSWORD = "manufacturing123"

def authenticate(user, pwd):
    if user == USERNAME and pwd == PASSWORD:
        return gr.update(visible=True), gr.update(visible=False)
    return gr.update(visible=False), gr.update(visible=True)

# ===============================
# LIVE IOT STREAM (PLC Simulation)
# ===============================

def generate_iot_data():
    return pd.DataFrame({
        "temp": np.random.normal(90, 15, 5),
        "vibration": np.random.normal(0.6, 0.3, 5),
        "output": np.random.normal(100, 20, 5)
    })

# ===============================
# NEURAL NETWORK MODEL
# ===============================

base_data = generate_iot_data()
base_data["failure"] = np.where(base_data["temp"] > 110, 1, 0)

scaler = StandardScaler()
X = scaler.fit_transform(base_data[["temp","vibration","output"]])
y = base_data["failure"]

model = MLPClassifier(hidden_layer_sizes=(32,16), max_iter=500)
model.fit(X, y)

# Save model hash for integrity
def get_model_hash():
    return hashlib.md5(str(model.coefs_).encode()).hexdigest()

original_hash = get_model_hash()

# ===============================
# FGSM-STYLE ATTACK
# ===============================

def fgsm_attack(data, epsilon=5):
    attacked = data.copy()
    attacked["temp"] += epsilon
    return attacked

# ===============================
# ATTACK ENGINE
# ===============================

def apply_attack(data, attack):

    if attack == "Data Poisoning":
        data["temp"] -= 40

    elif attack == "Model Evasion":
        data["temp"] = np.where(data["temp"] > 100, 99.5, data["temp"])

    elif attack == "FGSM Adversarial":
        data = fgsm_attack(data)

    elif attack == "Model Replacement":
        global model
        model = MLPClassifier()  # replace with empty weak model

    return data

# ===============================
# DEFENSE LAYER
# ===============================

def defense_layer(data):

    if data["temp"].mean() < 50:
        return "⚠ Data Poisoning Detected"

    if get_model_hash() != original_hash:
        return "⚠ Model Integrity Compromised"

    if data["temp"].std() < 3:
        return "⚠ Adversarial Pattern Detected"

    return "✅ System Secure"

# ===============================
# MAIN SIMULATION
# ===============================

def run_system(attack, defense_toggle):

    data = generate_iot_data()

    before_scaled = scaler.transform(data[["temp","vibration","output"]])
    before_pred = model.predict(before_scaled)

    if attack != "None":
        data = apply_attack(data, attack)

    after_scaled = scaler.transform(data[["temp","vibration","output"]])
    after_pred = model.predict(after_scaled)

    risk_score = int(sum(after_pred) * 15)

    if attack != "None" and defense_toggle == "OFF":
        risk_score += 40

    if defense_toggle == "ON":
        defense_status = defense_layer(data)
    else:
        defense_status = "❌ Defense Disabled"

    # Risk Gauge
    gauge = go.Figure(go.Indicator(
        mode="gauge+number",
        value=risk_score,
        title={'text': "Enterprise Risk Index"},
        gauge={
            'axis': {'range': [0,100]},
            'steps': [
                {'range':[0,30],'color':"green"},
                {'range':[30,70],'color':"yellow"},
                {'range':[70,100],'color':"red"}
            ]
        }
    ))

    # SCADA Live Chart
    scada_chart = go.Figure()
    scada_chart.add_trace(go.Bar(y=data["temp"], name="PLC Temperature"))
    scada_chart.update_layout(title="Live SCADA - PLC Temperature")

    return gauge, scada_chart, defense_status

# ===============================
# CORPORATE EXECUTIVE UI
# ===============================

with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.Markdown("# 🏭 Smart Factory AI Security Command Center")
    gr.Markdown("Industry 4.0 | AI Attack Simulation | Executive Monitoring")

    login_box = gr.Column()
    dashboard_box = gr.Column(visible=False)

    with login_box:
        user = gr.Textbox(label="Username")
        pwd = gr.Textbox(label="Password", type="password")
        login_btn = gr.Button("Login")

    with dashboard_box:
        attack = gr.Dropdown(["None","Data Poisoning","Model Evasion","FGSM Adversarial","Model Replacement"], label="Select AI Attack")
        defense = gr.Radio(["ON","OFF"], label="Defense System")

        run_btn = gr.Button("Run Simulation")

        risk_output = gr.Plot(label="Risk Gauge")
        scada_output = gr.Plot(label="SCADA Live View")
        defense_text = gr.Textbox(label="Defense Status")

    login_btn.click(authenticate, [user, pwd], [dashboard_box, login_box])
    run_btn.click(run_system, [attack, defense], [risk_output, scada_output, defense_text])

demo.launch()
trackio.show()