import gradio as gr
import earthview as ev
import utils
import random
import pandas as pd
import os

# Load the Satellogic dataset
dataset = ev.load_dataset("satellogic", streaming=True).shuffle(seed=42)
data_iter = iter(dataset)

# File to store labels (will create if it doesn't exist)
label_file = "labels.csv"

# Initialize a DataFrame to hold labels (or load existing)
if os.path.exists(label_file):
    labels_df = pd.read_csv(label_file)
else:
    labels_df = pd.DataFrame(columns=["image_id", "bounds", "rating", "google_maps_link"])

def get_next_image():
    global data_iter, labels_df

    while True:  # Keep iterating until we find an unlabeled image
        try:
            sample = next(data_iter)
        except StopIteration:
            #refresh the dataset if we reach the end
            dataset = ev.load_dataset("satellogic", streaming=True).shuffle(seed=random.randint(0, 1000000))
            data_iter = iter(dataset)
            continue

        sample = ev.item_to_images("satellogic", sample)
        image = sample["rgb"][0]  # Get the first RGB image
        metadata = sample["metadata"]

        bounds = metadata["bounds"]
        google_maps_link = utils.get_google_map_link(sample, "satellogic")
        #generate a unique image ID:
        image_id = (str(bounds))

        # Check if image is already labeled
        if image_id not in labels_df["image_id"].values:
            return image, image_id, bounds, google_maps_link

def rate_image(image_id, bounds, rating, google_maps_link):
    global labels_df

    # Add the rating to the DataFrame
    new_row = pd.DataFrame({"image_id": [image_id], "bounds": [bounds], "rating": [rating], "google_maps_link": [google_maps_link]})
    labels_df = pd.concat([labels_df, new_row], ignore_index=True)

    # Save the DataFrame to CSV
    labels_df.to_csv(label_file, index=False)

    # Get the next image and its details
    next_image, next_image_id, next_bounds, next_google_maps_link = get_next_image()

    return next_image, next_image_id, next_bounds, next_google_maps_link

# Define the Gradio interface
iface = gr.Interface(
    fn=rate_image,
    inputs=[
        gr.Textbox(label="Image ID", visible=False),
        gr.Textbox(label="Bounds", visible=False),
        gr.Radio(["Cool", "Not Cool"], label="Rating"),
        gr.Textbox(label="Google Maps Link"),
    ],
    outputs=[
        gr.Image(label="Satellite Image"),
        gr.Textbox(label="Image ID", visible=False),
        gr.Textbox(label="Bounds", visible=False),
        gr.Textbox(label="Google Maps Link"),
    ],
    title="TerraNomaly - Satellite Image Labeling",
    description="Rate satellite images as 'Cool' or 'Not Cool'.",
    live=False,
)

# Get the first image and its details
initial_image, initial_image_id, initial_bounds, initial_google_maps_link = get_next_image()

# Set the initial values for the output components
iface.launch(
    share=True,
    initial_outputs=[
        initial_image,
        initial_image_id,
        initial_bounds,
        initial_google_maps_link,
    ],

)