import os
import sys
import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image
from scipy.ndimage import rotate, gaussian_filter
import gradio as gr
from huggingface_hub import hf_hub_download

hf_hub_download(repo_id="MultivexAI/RobustMNIST-v1.0", filename="model.py", local_dir=".")
hf_hub_download(repo_id="MultivexAI/RobustMNIST-v1.0", filename="model.pt", local_dir=".")

sys.path.append(os.path.abspath("."))
from model import HierarchicalNetwork

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = HierarchicalNetwork(out_dims=11).to(DEVICE)
model.load_state_dict(torch.load("model.pt", map_location=DEVICE))
model.eval()

def preprocess_and_predict(sketch_data, rotation_val, noise_val, blur_val):
    if sketch_data is None:
        return None, {}

    if isinstance(sketch_data, dict):
        img_array = sketch_data.get("composite", None)
        if img_array is None:
            layers = sketch_data.get("layers", [])
            img_array = layers[0] if layers else None
    else:
        img_array = sketch_data

    if img_array is None:
        return None, {}

    pil_img = Image.fromarray(img_array.astype('uint8'))
    if pil_img.mode == 'RGBA':
        canvas_bg = Image.new("RGB", pil_img.size, (255, 255, 255))
        canvas_bg.paste(pil_img, mask=pil_img.split()[3])
        gray_img = canvas_bg.convert('L')
    else:
        gray_img = pil_img.convert('L')

    resized_img = gray_img.resize((28, 28), Image.Resampling.LANCZOS)
    np_img = np.array(resized_img).astype(np.float32)

    border_average = (np_img[0, :].mean() + np_img[-1, :].mean() + np_img[:, 0].mean() + np_img[:, -1].mean()) / 4.0
    if border_average > 127.5:
        np_img = 255.0 - np_img

    if rotation_val > 0:
        np_img = rotate(np_img, rotation_val, reshape=False, order=1, mode='constant', cval=0.0)

    if blur_val > 0:
        np_img = gaussian_filter(np_img, sigma=blur_val)

    if noise_val > 0:
        variance_scale = noise_val * 255.0
        additive_noise = np.random.normal(0, variance_scale, np_img.shape)
        np_img = np.clip(np_img + additive_noise, 0.0, 255.0)

    normalized_array = np_img / 255.0
    tensor_input = torch.tensor(normalized_array, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(DEVICE)

    with torch.inference_mode():
        logits = model(tensor_input)
        probabilities = F.softmax(logits, dim=1).cpu().numpy()[0]

    class_labels = [str(i) for i in range(10)] + ["Unknown"]
    distribution = {class_labels[i]: float(probabilities[i]) for i in range(11)}
    
    preview_pil = Image.fromarray(np.clip(np_img, 0, 255).astype(np.uint8))
    preview_output = preview_pil.resize((280, 280), Image.Resampling.NEAREST)

    return preview_output, distribution

with gr.Blocks(title="Robust MNIST Classifier") as interface:
    gr.Markdown("## Robust Hierarchical Classifier")
    gr.Markdown("Draw a single digit, adjust the sliders to apply synthetic environmental distortions, and observe the robustness profile.")
    
    with gr.Row():
        with gr.Column():
            canvas = gr.Sketchpad(
                label="Draw Digit", 
                type="numpy"
            )
            rotation = gr.Slider(minimum=0, maximum=180, value=0, step=1, label="Rotation Angle (Degrees)")
            noise = gr.Slider(minimum=0.0, maximum=1.0, value=0.0, step=0.05, label="Gaussian Noise Level")
            blur = gr.Slider(minimum=0.0, maximum=5.0, value=0.0, step=0.1, label="Gaussian Blur (Sigma)")
            run_btn = gr.Button("Evaluate Signature", variant="primary")
            
        with gr.Column():
            preview = gr.Image(label="Model-View Reconstruction (280x280)", image_mode="L")
            probabilities_output = gr.Label(num_top_classes=5, label="Probability Map Output")

    run_btn.click(
        fn=preprocess_and_predict,
        inputs=[canvas, rotation, noise, blur],
        outputs=[preview, probabilities_output]
    )

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