import gradio as gr
import torch
import numpy as np
from PIL import Image
from transformers import AutoModel, AutoConfig

# The 11 classes your model predicts
FAULT_CLASSES = ['AG','BG','CG','AB','AC','BC','ABG','ACG','BCG','ABCG','NF']

# 1. Load the model directly from your new Hugging Face repository
repo_id = "Sanath2709/chronogrid-fusionnet"
print(f"Loading model from {repo_id}...")

config = AutoConfig.from_pretrained(repo_id, trust_remote_code=True)
model = AutoModel.from_pretrained(repo_id, config=config, trust_remote_code=True)
model.eval()

def predict(image, noise_std):
    if image is None:
        return None
        
    # Preprocess
    img = image.convert('L').resize((227, 227))
    arr = np.array(img, dtype=np.float32) / 255.0  
    
    # Add Gaussian Noise if requested (Robustness Testing)
    if noise_std > 0:
        noise = np.random.normal(0, noise_std, arr.shape).astype(np.float32)
        arr = np.clip(arr + noise, 0.0, 1.0)
        
    # Create a visual representation of the noisy image to show the user
    noisy_image_visual = Image.fromarray((arr * 255).astype(np.uint8))
    
    # Add batch dimension -> Shape becomes [1, 227, 227]
    x = torch.tensor(arr).unsqueeze(0)
    
    # Run Inference
    with torch.no_grad():
        outputs = model(x)
        logits = outputs[0] if isinstance(outputs, tuple) else outputs
        probs = torch.softmax(logits, dim=1).squeeze().numpy()
        
    # Format output dictionary
    confidences = {FAULT_CLASSES[i]: float(probs[i]) for i in range(len(FAULT_CLASSES))}
    return confidences, noisy_image_visual

with gr.Blocks(title="ChronoGrid FusionNet") as interface:
    with gr.Row():
        gr.Image("logo.svg", width=60, show_label=False, container=False, scale=0)
        gr.Markdown("# ChronoGrid FusionNet Simulator\n### IEEE 9-Bus Fault Detection via Time-Frequency S-Transform Signatures")
    
    with gr.Row():
        # LEFT COLUMN (Inputs)
        with gr.Column():
            input_image = gr.Image(type="pil", label="Upload S-Transform Heatmap", height=350)
            noise_slider = gr.Slider(
                minimum=0.0, maximum=0.5, step=0.05, value=0.0, 
                label="Gaussian Noise Level (std)", 
                info="Add noise to test model robustness"
            )
            
            with gr.Row():
                clear_btn = gr.Button("Clear")
                submit_btn = gr.Button("Analyze Signal", variant="primary")
                
            gr.Markdown("""
            **Dataset Simulation:** 
            Generate custom test signals using the [Al-Roomi Repository](https://al-roomi.org/power-flow) or [MATLAB File Exchange](https://www.mathworks.com/matlabcentral/fileexchange/65385-wscc-9-bus-test-system-ieee-benchmark) Simulink models.
            """)
                
        # RIGHT COLUMN (Outputs & Info)
        with gr.Column():
            output_label = gr.Label(num_top_classes=3, label="Network Prediction")
            noisy_image_display = gr.Image(label="Signal fed to the model (with noise)", interactive=False, height=250)
            
            with gr.Accordion("Fault Types Reference (Dataset Examples)", open=True):
                # Dynamically load the 11 actual dataset examples we copied
                import os
                gallery_images = []
                gallery_dir = "examples/gallery"
                if os.path.exists(gallery_dir):
                    for img_name in sorted(os.listdir(gallery_dir)):
                        if img_name.endswith('.jpg'):
                            caption = img_name.replace('.jpg', '')
                            gallery_images.append((os.path.join(gallery_dir, img_name), caption))
                
                gr.Gallery(value=gallery_images, columns=3, height="auto", show_label=False, object_fit="contain")
            
            gr.Markdown("""
            *Powered by ChronoGrid FusionNet (1D CNN + BiLSTM + Attention)*  
            *View model on [Hugging Face](https://huggingface.co/Sanath2709/chronogrid-fusionnet)*
            """)
            
            # Allow users to download the paper
            gr.File(value="ChronoGrid_FusionNet_Full_Paper.docx", label="Download Full Research Manuscript")

    # EXAMPLES SECTION
    gr.Markdown("### Try it yourself (Sample Signals)")
    gr.Examples(
        examples=[
            ["examples/AG.jpg", 0.0],
            ["examples/ABG.jpg", 0.1],
            ["examples/NF.jpg", 0.0]
        ],
        inputs=[input_image, noise_slider],
        outputs=[output_label, noisy_image_display],
        fn=predict,
        cache_examples=False,
    )

    # Event Listeners
    submit_btn.click(fn=predict, inputs=[input_image, noise_slider], outputs=[output_label, noisy_image_display])
    clear_btn.click(lambda: (None, 0.0, None, None), inputs=None, outputs=[input_image, noise_slider, output_label, noisy_image_display])

if __name__ == "__main__":
    interface.launch(share=False)