from datasets import load_dataset
from functools import partial
from pandas import DataFrame
import earthview as ev
import gradio as gr 
import tqdm
import os

DEBUG = False      # False, "random", "samples"

if DEBUG == "random":
    import numpy as np

def open_dataset(dataset, subset, split, batch_size, shard, only_rgb, state):

    nshards = ev.get_nshards(subset)
    
    if shard == -1:
        shards = None
    else:
        shards = [shard]

    if DEBUG == "random":
        ds = range(batch_size)
    elif DEBUG == "samples":
        ds = ev.load_parquet(subset, batch_size=batch_size)
    elif not DEBUG:
        ds = ev.load_dataset(subset, dataset=dataset, split=split, shards=shards, cache_dir="dataset")
    
    dsi = iter(ds)

    state["subset"]  = subset
    state["dsi"] = dsi
    return (
        gr.update(label=f"Shard (max {nshards})", value=shard, maximum=nshards),
        *get_images(batch_size, only_rgb, state),
        state
    )
    
def get_images(batch_size, only_rgb, state):
    subset = state["subset"]

    images = []
    metadatas = []
    
    for i in tqdm.trange(batch_size, desc=f"Getting images"):
        if DEBUG == "random":
            images.append(np.random.randint(0,255,(384,384,3)))
            if not only_rgb:
                images.append(np.random.randint(0,255,(100,100,3)))

            metadatas.append({"bounds":[[1,1,4,4]], })
        else:
            try:
                item = next(state["dsi"])
            except StopIteration:
                break
            metadata = item["metadata"]
            item = ev.item_to_images(subset, item)

            if  subset == "satellogic":
                images.extend(item["rgb"])
                if not only_rgb:
                    images.extend(item["1m"])
            if  subset == "sentinel_1":
                images.extend(item["10m"])
            if  subset == "neon":
                images.extend(item["rgb"])
                if not only_rgb:
                    images.extend(item["chm"])
                    images.extend(item["1m"])
            
            metadatas.append(item["metadata"])

    return images, DataFrame(metadatas)

def update_shape(rows, columns):
    return gr.update(rows=rows, columns=columns)

def new_state():
    return gr.State({})

if __name__ == "__main__":
    with gr.Blocks(title="EarthView Viewer", fill_height = True) as demo:
        state = new_state()

        gr.Markdown(f"# Viewer for [{ev.DATASET}](https://huggingface.co/datasets/satellogic/EarthView) Dataset")
        batch_size = gr.Number(10, label = "Batch Size", render=False)
        shard = gr.Slider(label="Shard", minimum=0, maximum=10000, step=1, render=False)
        table = gr.DataFrame(render = False)
        # headers=["Index","TimeStamp","Bounds","CRS"], 

        gallery = gr.Gallery(
            label=ev.DATASET,
            interactive=False,
            object_fit="scale-down",
            columns=5, rows=2, render=False)

        with gr.Row():
            dataset = gr.Textbox(label="Dataset", value=ev.DATASET, interactive=False)
            subset = gr.Dropdown(choices=ev.get_subsets(), label="Subset", value="satellogic", )
            split = gr.Textbox(label="Split", value="train")
            initial_shard = gr.Number(label = "Initial shard", value=10, info="-1 for whole dataset")
            only_rgb = gr.Checkbox(label="Only RGB", value=True)

            gr.Button("Load (minutes)").click(
                open_dataset,
                inputs=[dataset, subset, split, batch_size, initial_shard, only_rgb, state],
                outputs=[shard, gallery, table, state])

        gallery.render()
        
        with gr.Row():
            batch_size.render()

            rows = gr.Number(2, label="Rows")
            columns = gr.Number(5, label="Coluns")

            rows.change(update_shape, [rows, columns], [gallery])
            columns.change(update_shape, [rows, columns], [gallery])

        with gr.Row():
            shard.render()
            shard.release(
                open_dataset,
                inputs=[dataset, subset, split, batch_size, shard, only_rgb, state],
                outputs=[shard, gallery, table, state])

            btn = gr.Button("Next Batch (same shard)", scale=0)
            btn.click(get_images, [batch_size, only_rgb, state], [gallery, table])
            btn.click()
        
        table.render()

    demo.launch(show_api=False)