import os
from fast_dash import FastDash, UploadImage, Image, Upload
import PIL
import numpy as np

from app import func
import tempfile
import base64

examples = ["India_900498_S2Hand", "Spain_7370579_S2Hand", "USA_430764_S2Hand"]

def detect_water(select_an_example: str = examples, input_tiff_image: Upload = None) -> (Image, Image):
    """
    NASA and IBM recently uploaded their foundation model, Prithvi, on Hugging Face at https://huggingface.co/ibm-nasa-geospatial. 
    This demo, built with Fast Dash, showcases a version of Prithvi that they finetuned to detect water from satellite images.
    Select an example or upload your own TIFF image.
    Before uploading your own image, read the format details at https://huggingface.co/ibm-nasa-geospatial/Prithvi-100M-sen1floods11.
    """

    # If example is selected
    if input_tiff_image is None:
        input_satellite_image = PIL.Image.open(os.path.join("examples", f"{select_an_example}.png"))
        water_prediction_mask = PIL.Image.open(os.path.join("examples", f"{select_an_example}_result.png"))

    # If file is uploaded, run inference
    else:
        with tempfile.NamedTemporaryFile(delete=False) as f:
            contents = input_tiff_image.encode("utf8").split(b";base64,")[1]
            f.write(base64.decodebytes(contents))

        input_satellite_image, water_prediction_mask = func(f)

        input_satellite_image = PIL.Image.fromarray(np.uint8(input_satellite_image)).convert('RGB')
        water_prediction_mask = PIL.Image.fromarray(np.uint8(water_prediction_mask)).convert('RGB')

    return input_satellite_image, water_prediction_mask

app = FastDash(detect_water, theme="cosmo", port=7860)
server = app.server

if __name__ == "__main__":
    app.run()