import os
import pandas as pd
from PIL import Image
from itertools import islice
import earthview as ev  # Your custom library
import gradio as gr

# Path to the first shard of the Satellogic dataset
filename = os.path.join("dataset", "satellogic", "train-00000-of-00065.parquet")

# Check if the file exists
if not os.path.isfile(filename):
    raise FileNotFoundError(f"{filename} not found")

# Loads a dataset with pandas, this loads a single file
data = pd.read_parquet(filename)

# Global variables to manage state
current_image = None
current_metadata_index = 0
ratings = []
image_id = 0
bounds = []
timestamps = []

# Function to transform a metadata sample to bounds and timestamps
def item_to_bounds_timestamps(sample):
    bounds_list = sample["metadata"]["bounds"]
    timestamp_list = sample["metadata"]["timestamp"]
    bounds = []
    timestamps = []
    for b, t in zip(bounds_list, timestamp_list):
        bounds.append(b)
        timestamps.append(t)
    return bounds, timestamps

# Function to load and display the next image
def load_next_image():
    global current_image, image_id, current_metadata_index, bounds, timestamps

    if current_metadata_index >= len(data):
        return None, "No more images", None

    sample = data.iloc[current_metadata_index]
    bounds_sample, timestamps_sample = item_to_bounds_timestamps(sample)

    # Use earthview library to convert arrays to PIL images
    sample = ev.item_to_images("satellogic", sample.to_dict())
    current_image = sample["rgb"][0]  # Get the first image
    image_id += 1

    return current_image, f"Image ID: {image_id}", bounds_sample[0]

# Function to handle rating submission
def submit_rating(rating, bounds_str):
    global image_id, current_metadata_index, ratings, bounds, timestamps

    ratings.append(rating)
    bounds.append(bounds_str)
    timestamps.append("timestamp")  # Use a valid timestamp if available
    
    current_metadata_index += 1

    return load_next_image()

# Gradio Interface Layout
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            image_display = gr.Image(label="Satellite Image", type="pil")
            image_id_display = gr.Textbox(label="Image ID")
            bounds_display = gr.Textbox(label="Bounds")
            load_button = gr.Button("Load Next Image")
            load_button.click(fn=load_next_image, outputs=[image_display, image_id_display, bounds_display])
        with gr.Column():
            rating_radio = gr.Radio(["0", "1"], label="Rating (0 = No, 1 = Yes)")
            submit_button = gr.Button("Submit Rating")
            submit_button.click(fn=lambda rating, bounds_str: submit_rating(rating, bounds_str), inputs=[rating_radio, bounds_display], outputs=[image_display, image_id_display, bounds_display])

# Launch the Gradio interface
demo.launch()