import os
import sys
import json
import torch
import numpy as np
import cv2
from PIL import Image
from torchvision import transforms
from safetensors.torch import load_model
from huggingface_hub import hf_hub_download
import gradio as gr

# Add src to path
sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))
from model import SEMViTAutoencoder

# --- Hub Configuration ---
REPO_ID = "morinousagi/sem-vit-anomaly" 

def load_resources():
    device = torch.device("cpu")
    
    # 1. Download files from the Model Repo
    print("Downloading model weights and config from Hub...")
    model_path = hf_hub_download(repo_id=REPO_ID, filename="model.safetensors")
    config_path = hf_hub_download(repo_id=REPO_ID, filename="config.json")
    
    # 2. Initialize Model
    model = SEMViTAutoencoder()
    load_model(model, model_path, strict=False)
    model.to(device)
    model.eval()
    
    # 3. Load Threshold
    with open(config_path, "r") as f:
        config = json.load(f)
        threshold = config.get("threshold", 1.0)   # changed to 1.0
        
    return model, threshold, device

# Global initialization
MODEL, THRESHOLD, DEVICE = load_resources()



def predict(img):
    if img is None:
        return None, None
    
    # 1. Preprocess (matching dataset.py)
    transform = transforms.Compose([
        transforms.Resize((512, 512)),
        transforms.Grayscale(num_output_channels=3),
        transforms.ToTensor()
    ])
    
    pil_img = Image.fromarray(img.astype('uint8'))
    input_tensor = transform(pil_img).unsqueeze(0).to(DEVICE)
    
    # 2. Inference
    with torch.no_grad():
        output_tensor = MODEL(input_tensor)
    
    # 3. Calculate Anomaly Map (matching Peak Score Logic)
    # Move to CPU and numpy
    orig = input_tensor.squeeze().cpu().numpy() # [3, 512, 512]
    recon = output_tensor.squeeze().cpu().numpy() # [3, 512, 512]
    
    # Pixel-wise MSE across channels
    diff = np.mean(np.square(orig - recon), axis=0) # [512, 512]
    
    # 4. Scoring Logic (matching train_eval.py)
    # Apply Gaussian Blur to aggregate the defect signal (mimics avg_pool2d)
    smoothed_diff = cv2.GaussianBlur(diff, (15, 15), 0)
    
    # Peak Score (Max value in the smoothed map)
    score = np.max(smoothed_diff)
    
    is_defective = score > THRESHOLD
       
    status = "DEFECTIVE" if is_defective else "NORMAL"
    bg_color = "#fee2e2" if is_defective else "#dcfce7"
    text_color = "#b91c1c" if is_defective else "#15803d"
    
    html_result = f"""
    <div style="background-color: {bg_color}; padding: 20px; border-radius: 10px; text-align: center;">
        <h2 style="color: {text_color}; margin: 0;">{status}</h2>
        <p style="color: {text_color}; opacity: 0.8;">Peak Score: {score:.5f} | Threshold: {THRESHOLD:.5f}</p>
    </div>
    """
    
    # 5. Generate Heatmap
    # Use the smoothed_diff for the heatmap so the "glow" matches the score location
    heatmap_gray = cv2.normalize(smoothed_diff, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
    heatmap_color = cv2.applyColorMap(heatmap_gray, cv2.COLORMAP_JET)
    heatmap_rgb = cv2.cvtColor(heatmap_color, cv2.COLOR_BGR2RGB)
    
    return html_result, heatmap_rgb




# --- Minimal UI Layout ---
with gr.Blocks() as demo:
    gr.Markdown("# 🔬 Wafer SEM Anomaly Inspection")
    gr.Markdown("For more information, refer to [README.md](https://huggingface.co/spaces/morinousagi/pytorch-vit-defect-detect/blob/main/README.md)")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_ui = gr.Image(label="Input SEM Image (512x512)", type="numpy")
            with gr.Row():
                reset_btn = gr.Button("Reset")
                run_btn = gr.Button("Run Analysis", variant="primary")
                
        with gr.Column(scale=1):
            result_ui = gr.HTML()
            heatmap_ui = gr.Image(label="Defect Location Heatmap")

    # Interactivity
    run_btn.click(predict, inputs=input_ui, outputs=[result_ui, heatmap_ui])
    
    # Reset functionality
    reset_btn.click(
        lambda: (None, "", None),
        inputs=None,
        outputs=[input_ui, result_ui, heatmap_ui]
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Default(), server_name="0.0.0.0", server_port=7860)