import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
import torch
import torchvision.transforms as transforms
import timm

# URL for the Hugging Face checkpoint
CHECKPOINT_URL = "https://huggingface.co/ReefNet/beit_global/resolve/main/checkpoint-60.pth"

# Class labels
all_classes = [
    'Acanthastrea', 'Acropora', 'Agaricia', 'Alveopora', 'Astrea', 'Astreopora',
    'Caulastraea', 'Coeloseris', 'Colpophyllia', 'Coscinaraea', 'Ctenactis',
    'Cycloseris', 'Cyphastrea', 'Dendrogyra', 'Dichocoenia', 'Diploastrea',
    'Diploria', 'Dipsastraea', 'Echinophyllia', 'Echinopora', 'Euphyllia',
    'Eusmilia', 'Favia', 'Favites', 'Fungia', 'Galaxea', 'Gardineroseris',
    'Goniastrea', 'Goniopora', 'Halomitra', 'Herpolitha', 'Hydnophora',
    'Isophyllia', 'Isopora', 'Leptastrea', 'Leptoria', 'Leptoseris',
    'Lithophyllon', 'Lobactis', 'Lobophyllia', 'Madracis', 'Meandrina', 'Merulina',
    'Montastraea', 'Montipora', 'Mussa', 'Mussismilia', 'Mycedium', 'Orbicella',
    'Oulastrea', 'Oulophyllia', 'Oxypora', 'Pachyseris', 'Pavona', 'Pectinia',
    'Physogyra', 'Platygyra', 'Plerogyra', 'Plesiastrea', 'Pocillopora',
    'Podabacia', 'Porites', 'Psammocora', 'Pseudodiploria', 'Sandalolitha',
    'Scolymia', 'Seriatopora', 'Siderastrea', 'Stephanocoenia', 'Stylocoeniella',
    'Stylophora', 'Tubastraea', 'Turbinaria'
]

# Example image paths
example_images = {
    "Acropora": "coral_images/Acropora_millepora.jpg",
    "Agaricia": "coral_images/Agaricia_agaricites.jpg",
    "Acropora aculeus": "coral_images/Acropora_aculeus.jpg",
    "Montipora": "coral_images/Montipora_patula.jpg",
    "Pocillopora": "coral_images/Pocillopora_acuta.jpg",
    "Porites": "coral_images/porities_lobata.jpg",
    "Favites": "coral_images/Favites_abdita.jpg",
    "Fungia": "coral_images/Fungia_concinna.jpg",
}

# Function to load the BeIT model
def load_model(model_name):
    print(f"Loading {model_name} model...")
    args = type('', (), {})()
    args.model = 'beitv2_large_patch16_224.in1k_ft_in22k_in1k'
    args.nb_classes = len(all_classes)
    args.drop_path = 0.1

    # Create model
    model = timm.create_model(
        args.model,
        pretrained=False,
        num_classes=args.nb_classes,
        drop_path_rate=args.drop_path,
        use_rel_pos_bias=True,
        use_abs_pos_emb=True,
    )

    # Load checkpoint from Hugging Face
    checkpoint = torch.hub.load_state_dict_from_url(CHECKPOINT_URL, map_location="cpu")
    state_dict = checkpoint.get('model', checkpoint)
    filtered_state_dict = {k: v for k, v in state_dict.items() if "relative_position_index" not in k}
    model.load_state_dict(filtered_state_dict, strict=False)

    # Move model to CUDA if available
    model.eval()
    if torch.cuda.is_available():
        model.cuda()
    return model

# Preprocessing transforms
preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Initialize model
model = load_model('beit')

def predict_label(image):
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    input_tensor = preprocess(image).unsqueeze(0)
    if torch.cuda.is_available():
        input_tensor = input_tensor.cuda()

    with torch.no_grad():
        outputs = model(input_tensor)
        predicted_class = torch.argmax(outputs, dim=1).item()

    return all_classes[predicted_class]

def draw_rectangle(image, x, y, size=224):
    """Draw a clear red rectangle with increased thickness."""
    image_pil = image.copy()  # Create a copy to avoid modifying the original image
    draw = ImageDraw.Draw(image_pil)
    x1, y1 = x, y
    x2, y2 = x + size, y + size
    draw.rectangle([x1, y1, x2, y2], outline="red", width=6)  # Increase the width for clarity
    return image_pil


def crop_image(image, x, y, size=224):
    image_np = np.array(image)
    h, w, _ = image_np.shape
    x = min(max(x, 0), w - size)
    y = min(max(y, 0), h - size)
    cropped = image_np[y:y+size, x:x+size]
    return Image.fromarray(cropped)

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## Coral Classification with BeIT Model")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Image", interactive=True)
            x_slider = gr.Slider(0, 1000, step=1, value=0, label="X Coordinate")
            y_slider = gr.Slider(0, 1000, step=1, value=0, label="Y Coordinate")
        with gr.Column():
            interactive_image = gr.Image(label="Interactive Image")
            cropped_image = gr.Image(label="Cropped Patch")
            label_output = gr.Textbox(label="Predicted Label")
    
    # Crop and Predict buttons
    crop_button = gr.Button("Crop")
    predict_button = gr.Button("Predict")

    # Example table
    def load_example(example_path):
        return Image.open(example_path).convert("RGB")

        # Generate table of examples
    with gr.Row():
        gr.Markdown("### Example Images for Quick Testing")

    with gr.Row():
        for genus, path in example_images.items():
            with gr.Column():
                thumbnail = gr.Image(value=path, interactive=False, label=genus)
                select_button = gr.Button(value=f"Select {genus}")
                select_button.click(fn=lambda p=path: load_example(p), inputs=None, outputs=image_input)

    # Button functionality
    crop_button.click(fn=lambda img, x, y: (draw_rectangle(img, x, y), crop_image(img, x, y)), 
                      inputs=[image_input, x_slider, y_slider], outputs=[interactive_image, cropped_image])
    predict_button.click(fn=predict_label, inputs=cropped_image, outputs=label_output)

    def update_sliders(image):
        if image:
            width, height = image.size
            return gr.update(maximum=width - 224), gr.update(maximum=height - 224)
        return gr.update(), gr.update()

    image_input.change(fn=update_sliders, inputs=image_input, outputs=[x_slider, y_slider])

# demo.launch()
demo.launch(server_name="0.0.0.0", server_port=7860)