import gradio as gr
import numpy as np
import pickle
import pygad.nn
import pygad.gann

# Load model and create Gradio interface
def load_model():
    with open('smart_grid_model (1).pkl', 'rb') as f:
        model_data = pickle.load(f)
    return model_data['GANN_instance'], model_data['best_solution_idx'], model_data['scaler']

def predict_stability(tau1, tau2, tau3, tau4, p1, p2, p3, p4, g1, g2, g3, g4):
    try:
        # Load model if not already loaded
        if 'model' not in predict_stability.__dict__:
            predict_stability.model, predict_stability.best_idx, predict_stability.scaler = load_model()

        # Prepare and scale input
        inputs = np.array([[tau1, tau2, tau3, tau4, p1, p2, p3, p4, g1, g2, g3, g4]], dtype=np.float32)
        inputs_scaled = predict_stability.scaler.transform(inputs)

        # Make prediction
        prediction = pygad.nn.predict(
            last_layer=predict_stability.model.population_networks[predict_stability.best_idx],
            data_inputs=inputs_scaled
        )

        # Process results
        stability = "Stable" if prediction[0] > 0.5 else "Unstable"
        confidence = float(abs(prediction[0] - 0.5) * 2)

        return stability, confidence

    except Exception as e:
        print(f"Prediction error: {str(e)}")
        return "Error", 0.0

# Create Gradio interface
iface = gr.Interface(
    fn=predict_stability,
    inputs=[
        gr.Number(label="tau1: Reaction time of supplier (power generation)"),
        gr.Number(label="tau2: Reaction time of consumer 1"),
        gr.Number(label="tau3: Reaction time of consumer 2"),
        gr.Number(label="tau4: Reaction time of consumer 3"),
        gr.Number(label="p1: Power produced by supplier"),
        gr.Number(label="p2: Power consumed by consumer 1"),
        gr.Number(label="p3: Power consumed by consumer 2"),
        gr.Number(label="p4: Power consumed by consumer 3"),
        gr.Number(label="g1: Willingness to adapt (supplier)"),
        gr.Number(label="g2: Willingness to adapt (consumer 1)"),
        gr.Number(label="g3: Willingness to adapt (consumer 2)"),
        gr.Number(label="g4: Willingness to adapt (consumer 3)")
    ],
    outputs=[
        gr.Label(label="Grid Stability"),
        gr.Number(label="Confidence Score")
    ],
    title="Smart Grid Stability Predictor",
    description="""
    Enter grid parameters to predict stability.""",
    examples=[
        [2.9590, 3.0798, 8.3810, 9.78075, 3.7630, -0.78260, -1.2573, -1.72308, 0.65045, 0.85957, 0.88744, 0.95803]
    ]
)

if __name__ == "__main__":
    iface.launch()