import gradio as gr
import torch
import numpy as np
from PIL import Image
import os

print("Starting Lung Ultrasound AI App...")
print(f"Current directory: {os.getcwd()}")
print(f"Files in directory: {os.listdir('.')}")

CLASS_NAMES = ['covid', 'other', 'healthy']
device = torch.device('cpu')

model_path = "pytorch_model.bin"
model = None

# Load model
try:
    from model import LungUltrasoundModel
    model = LungUltrasoundModel(num_classes=3, num_seg_classes=1)
    if os.path.exists(model_path):
        print(f"Loading model from {model_path}")
        state_dict = torch.load(model_path, map_location=device)
        model.load_state_dict(state_dict, strict=False)
        model.eval()
        print("Model loaded successfully!")
    else:
        print(f"Model file not found: {model_path}")
except Exception as e:
    print(f"Error loading model: {e}")

def predict(image):
    if image is None:
        return "No image uploaded"
    
    if model is None:
        return "Model not loaded"
    
    # Preprocess
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    image = image.resize((224, 224))
    img_array = np.array(image).astype(np.float32) / 255.0
    img_tensor = torch.from_numpy(img_array).permute(2, 0, 1).unsqueeze(0)
    
    with torch.no_grad():
        logits, _ = model(img_tensor)
        probs = torch.softmax(logits, dim=1)
        pred_class = torch.argmax(probs, dim=1).item()
        confidence = probs[0, pred_class].item()
    
    result = f"""
    ========================================
    PREDICTION: {CLASS_NAMES[pred_class]}
    CONFIDENCE: {confidence:.1%}
    ========================================
    COVID-19:  {probs[0,0]:.1%}
    Other:     {probs[0,1]:.1%}
    Healthy:   {probs[0,2]:.1%}
    """
    
    return result

# Simple interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Ultrasound Image"),
    outputs=gr.Textbox(label="Results", lines=10),
    title="Lung Ultrasound AI",
    description="Upload a lung ultrasound image to get disease prediction."
)

iface.launch(server_name="0.0.0.0", server_port=7860, share=True)