import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader
import torchvision
from PIL import Image
import gradio as gr
from huggingface_hub import hf_hub_download
import torch.nn as nn
import timm 

REPO_ID = "Raaniel/model-smoke"
MODEL_FILE_NAME = "model_smoke.pt"
USE_CUDA = torch.cuda.is_available()
num_classes = 3

# Download the model
checkpoint_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILE_NAME)

# Load the checkpoint
state = torch.load(checkpoint_path, map_location=torch.device('cuda' if USE_CUDA else 'cpu'))

# Create the model and modify it
model = timm.create_model('mobilenetv3_small_050', pretrained=True)
num_features = model.classifier.in_features

# Additional linear and dropout layers
model.classifier = nn.Sequential(
    nn.Linear(num_features, 256),  # Additional linear layer
    nn.ReLU(inplace=True),
    nn.Dropout(0.5),
    nn.Linear(256, num_classes)  # Final classification layer
)

# Load the model weights
model.load_state_dict(state['weights'])

# Move model to the appropriate device
device = torch.device('cuda' if USE_CUDA else 'cpu')
model = model.to(device)

classes = ["chmury", 'inne', "dym"]

def predict(image, model=model, classes=classes, device=device, transform=transform):
    model.eval()

    print(type(image))
    # Check if the image is a PyTorch Tensor, if so, use it directly
    if isinstance(image, torch.Tensor):
        img_batch = image.unsqueeze(0).to(device)
    elif isinstance(image, np.ndarray):  # Check if the image is a numpy ndarray
        # Convert numpy ndarray to PIL Image
        img = Image.fromarray(image)
        # Transform the image
        img_transformed = transform(img)
        # Convert to a batch of 1 and send to device
        img_batch = img_transformed.unsqueeze(0).to(device)
    else:
        # Load the image and apply transformations
        img = Image.open(image)
        img_transformed = transform(img)
        img_batch = img_transformed.unsqueeze(0).to(device)

    # Make predictions
    with torch.no_grad():
        _, predicted_idx = model(img_batch).max(1)

    # Map the index to the class name
    predicted_class = classes[predicted_idx.item()]

    return predicted_class

examples = ["https://img.freepik.com/free-photo/fantasy-style-clouds_23-2151057636.jpg?size=338&ext=jpg&ga=GA1.1.87170709.1707609600&semt=sph",
            "https://energyeducation.ca/wiki/images/5/51/Smoke_column_-_High_Park_Wildfire_%281%29.jpg",
            "https://img-aws.ehowcdn.com/360x267p/s3-us-west-1.amazonaws.com/contentlab.studiod/getty/31a4debc7443411195df509e38a5f9a3.jpg",
            "https://thumb.bibliocad.com/images/content/00000000/9000/9813.jpg",
            "https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcRC7j2LoW8D13BOgbT_9J2SI_krX0sadT4oaSuyFjNb3jElJdU-J7DpPgCYvEfFzqoD6c0&usqp=CAU"]


css = """
h1 {
    text-align: center;
    display:block;
}
"""

with gr.Blocks(theme=gr.themes.Base(primary_hue="zinc",
                                    secondary_hue="neutral",
                                    neutral_hue="slate",
                                    font =  gr.themes.GoogleFont("Montserrat")),
               css = css,
               title="Smoke Detection") as demo:

    demo.load(None, None, js="""
  () => {
  const params = new URLSearchParams(window.location.search);
  if (!params.has('__theme')) {
    params.set('__theme', 'light');
    window.location.search = params.toString();
  }
  }""",
    )

    markdown_content = """
    <img src='file/dd_logo.png' width='200'>
    """
    gr.Markdown(markdown_content)
    gr.Markdown("# 🔥 Wczesne rozpoznawanie pożarów 🔥")

    gr.Markdown(""" ## Zauważ Ogień, Ocal Naturę! W mgnieniu oka rozróżnij dym od chmur dzięki naszemu inteligentnemu narzędziu do wykrywania pożarów.
    ### Napędzane przez obszerną bazę ponad 14 000 obrazów i zaawansowane modele uczenia maszynowego, nasze narzędzie umożliwia wczesne rozpoznawanie ognia. Szybko, sprytnie i czujnie – chronimy nasze środowisko przed pierwszymi oznakami zagrożenia.""")

    with gr.Accordion("Więcej informacji", open = False):
      gr.Markdown("""Nasilenie się pożarów, potęgowane przez zmiany klimatyczne, stanowi poważne wyzwanie w szybkiej detekcji i reagowaniu.
      Tradycyjne metody wykrywania pożarów, takie jak obserwacja ludzka i raporty, często okazują się zbyt wolne, zwłaszcza na obszarach oddalonych.
      Automatyczne systemy wykrywania dymu oferują rozwiązanie, wykorzystując głębokie uczenie do szybkiego i dokładnego identyfikowania dymu na obrazach.
      Umiejętność odróżnienia dymu od zjawisk o podobnym wyglądzie, jak chmury, jest kluczowa. Ta różnica pozwala na szybsze lokalizowanie ognisk pożarów,
      co umożliwia szybsze czas reakcji i potencjalnie ratuje ogromne obszary regionów naturalnych i mieszkalnych przed zniszczeniem.
      Poprzez poprawę szybkości i dokładności wykrywania pożarów, możemy znacząco złagodzić ich wpływ na społeczności, gospodarki i ekosystemy.""")

    with gr.Column():
      image = gr.Image(label = "Obraz")

      gallery = gr.Gallery(value = examples, label="Przykładowe zdjęcia",columns=[4], rows=[1], height=200, object_fit = "scale-down")

      def get_select_index(evt: gr.SelectData):
        return examples[evt.index]

      gallery.select(get_select_index, None, image)
      action = gr.Button("Rozpoznaj")
      prediction = gr.Textbox(label = "Predykcja")

      action.click(fn=predict, inputs=image, outputs=prediction)

demo.launch(width = "75%", debug = True, allowed_paths=["/"])