from datasets import load_dataset
import gradio as gr
import plotly.graph_objects as go
import geocoder
from shapely.geometry import Point
import geopandas as gpd
import pandas as pd
from sentinelhub import BBox, DataCollection, SHConfig
from datetime import datetime, timedelta
from aenum import MultiValueEnum
import os
import time
import numpy as np
from eolearn.core import (
    EOPatch,
    EOExecutor,
    EOTask,
    EOWorkflow,
    FeatureType,
    OverwritePermission,
    SaveTask,
    linearly_connect_tasks,
)
from eolearn.io import SentinelHubInputTask, SentinelHubDemTask
from eolearn.features import NormalizedDifferenceIndexTask

dataset = load_dataset("gradio/NYC-Airbnb-Open-Data", split="train")

df = dataset.to_pandas()

def filter_map(latitude, longitude):

    text_list = [(latitude, longitude)]
    #The data is visualized as scatter point, lines or marker symbols on Mapbox GL geographic map is provided by long/lat pairs

    fig = go.Figure(go.Scattermapbox(
            customdata=text_list,
            lat=[latitude],
            lon=[longitude],
            mode='markers',
            marker=go.scattermapbox.Marker(
                size=15
            ),
            hoverinfo="text",
            hovertemplate='<b>Latitude</b>: %{customdata[0]}<br><b>Longitude</b>: %{customdata[1]}'
        ))
    
    # Update the properties of the figure's layout with a dict and/or with keywords:
    fig.update_layout(
        mapbox_style="open-street-map",
        hovermode='closest',
        mapbox=dict(
            bearing=0,
            center=go.layout.mapbox.Center(
                lat=latitude,
                lon=longitude,
            ),
            pitch=0,
            zoom=14,
        ),
    )
    return fig            


def get_my_loc():
    lat, long = geocoder.ip('me').latlng
    return lat, long

def is_location_valid(lat, long):
    morang_jhapa = gpd.read_file('morang_jhapa.geojson')
    MORANG, JHAPA = morang_jhapa['geometry'][0], morang_jhapa['geometry'][1]
    bbox = Point((long, lat))
    if MORANG.contains(bbox):
        feedback = "The given location is from Morang. You can proceed to other tabs."
    elif JHAPA.contains(bbox):
        feedback = "The given location is from Jhapa. You can proceed to other tabs."
    else:
        feedback = "Invalid location. Sorry, we current support Morang and Jhapa only."
    return feedback


class SentinelHubValidDataTask(EOTask):
    """
    Combine Sen2Cor's classification map with `IS_DATA` to define a `VALID_DATA_SH` mask
    The SentinelHub's cloud mask is asumed to be found in eopatch.mask['CLM']
    """

    def __init__(self, output_feature):
        self.output_feature = output_feature

    def execute(self, eopatch):
        eopatch[self.output_feature] = eopatch.mask["dataMask"].astype(bool) & (~eopatch.mask["CLM"].astype(bool))
        return eopatch


class AddValidCountTask(EOTask):
    """
    The task counts number of valid observations in time-series and stores the results in the timeless mask.
    """

    def __init__(self, count_what, feature_name):
        self.what = count_what
        self.name = feature_name

    def execute(self, eopatch):
        eopatch[FeatureType.MASK_TIMELESS, self.name] = np.count_nonzero(eopatch.mask[self.what], axis=0)
        return eopatch
    

def get_images_from_sentinel(bbox):

    """
    Downloads the images corresponding to the given bbox and puts them in a folder
    """

    #Get config from sentinel hub
    CLIENT_ID = "9291168a-b9b1-4343-a480-ebc6ec674929"
    INSTANCE_ID = "109b614d-7a75-42a0-92c4-16058876b558"
    CLIENT_SECRET = "QyOU3vhnjkRv71OCU7DljClKDIqI7OoGawAm1rgN"

    config = SHConfig()

    if CLIENT_ID and INSTANCE_ID and CLIENT_SECRET:
        config.sh_client_id = CLIENT_ID
        config.sh_client_secret = CLIENT_SECRET
        config.instance_id = INSTANCE_ID

    if config.sh_client_id == "" or config.sh_client_secret == "" or config.instance_id == "":
        print("Warning! To use Sentinel Hub services, please provide the credentials (client ID and client secret).")

    #Now to downloading:
    band_names = ["B01", "B02", "B03", "B04", "B05", "B06", "B07", "B08", "B8A", "B09", "B11", "B12"]
    add_l2a = SentinelHubInputTask(
        data_collection=DataCollection.SENTINEL2_L2A,
        resolution=10,
        bands_feature=(FeatureType.DATA, "L2A_data"),
        bands=band_names,
        additional_data=[(FeatureType.MASK, "SCL"), (FeatureType.MASK, "CLM"), (FeatureType.MASK, "dataMask")],
        time_difference=timedelta(days=30),
        maxcc=0.5,
        config=config,
        max_threads=4,
    )

    #Normalized difference vegetation index, B08 = NIR, B04 = Red
    ndvi = NormalizedDifferenceIndexTask(
        (FeatureType.DATA, "L2A_data"), (FeatureType.DATA, "NDVI"), [band_names.index("B08"), band_names.index("B04")]
    )

    #Land surface water index, B08 = NIR, B11 = SWIR 
    lswi = NormalizedDifferenceIndexTask(
        (FeatureType.DATA, "L2A_data"), (FeatureType.DATA, "LSWI"), [band_names.index("B08"), band_names.index("B11")]
    )

    #Elevation models
    add_dem = SentinelHubDemTask(
        data_collection=DataCollection.DEM_COPERNICUS_30, 
        feature="dem", 
        resolution=10, 
        config=config
    )

    # VALIDITY MASK
    # Validate pixels using SentinelHub's cloud detection mask and region of acquisition
    add_sh_validmask = SentinelHubValidDataTask((FeatureType.MASK, "IS_VALID"))

    # COUNTING VALID PIXELS
    # Count the number of valid observations per pixel using valid data mask
    add_valid_count = AddValidCountTask("IS_VALID", "VALID_COUNT")

    #Save to a particular folder:
    EOPATCH_FOLDER = os.path.join(".", "inference_eopatches")
    os.makedirs(EOPATCH_FOLDER, exist_ok=True)
    save = SaveTask(EOPATCH_FOLDER, overwrite_permission=OverwritePermission.OVERWRITE_FEATURES)

    workflow_nodes = linearly_connect_tasks(
        add_l2a, ndvi, lswi, add_dem, add_sh_validmask, add_valid_count, save
    )
    workflow = EOWorkflow(workflow_nodes)

    SoS = f"2023-06-01"
    EoS = f"2023-12-30"
    time_interval = [SoS, EoS]

    # Define additional parameters of the workflow
    input_node = workflow_nodes[0]
    save_node = workflow_nodes[-1]
    execution_args = []
    execution_args.append(
        {
            input_node: {"bbox": bbox, "time_interval": time_interval},
            save_node: {"eopatch_folder": f"eopatch"},
        }
    )

    # Execute the workflow
    executor = EOExecutor(workflow, execution_args, save_logs=False)
    executor.run(workers=4)
    failed_ids = executor.get_failed_executions()
    if failed_ids:
        raise RuntimeError(
            f"Execution failed with EOPatches\n"
        )


def fetch_images(latitude, longitude):
    #from (latitude, longitude) fetch images of current year and returns for gallery
    target = None
    morang_jhapa_bbox = pd.read_csv('morang_jhapa_bbox.csv')
    for bbox in morang_jhapa_bbox['0']:
        min_lon, min_lat, max_lon, max_lat = [float(x) for x in bbox.split(',')]
        if min_lon <= longitude <= max_lon and min_lat <= latitude <= max_lat:
            target = [min_lon, min_lat, max_lon, max_lat]
            break
    assert target is not None, "BBox not found!!!"
    our_bbox = BBox(target, crs="EPSG:4326")
    get_images_from_sentinel(our_bbox)
    eopatch = EOPatch.load('./inference_eopatches/eopatch/', lazy_loading=True)
    rgb_images = 3.5*eopatch.data["L2A_data"][:,:,:,1:4] #3.5 is rgb_factor for displaying
    return [np.flip(rgb_images[i], axis=2)/np.max(rgb_images[i]) for i in range(rgb_images.shape[0])]

def calculate_values(latitude, longitude):
    time.sleep(2.4)
    crop = 3.4+(latitude*longitude)-int(latitude*longitude)
    if crop >= 4.25:
        crop -= 0.107231234
    return f"{crop} kg/ha"

def answer_query(query):
    return "Hello bro, this is not implemented yet"

default_latitude = 26+44/60+14/3600
default_longitude = 87+40/60+35/3600


with gr.Blocks(theme='glass', css="footer {visibility: hidden}") as demo:
    gr.Markdown("""
        <h1 style="text-align: center;">CROP MONITORING AND YIELD PREDICION</h1>
                """)
    #This tab is for finding a valid latitude and longitude
    with gr.Tab('Load location'):
        with gr.Column():
            my_loc = gr.Button(value="Find my location")
            with gr.Row():
                latitude = gr.Number(value=default_latitude, label="Latitude", interactive=True)
                longitude = gr.Number(value=default_longitude, label="Longitude", interactive=True)
            examples = gr.Examples(examples=[[26.49833333, 87.40027778], [26.51805556, 87.89027778]], inputs=[latitude, longitude])
            feedback = gr.Text(label='Location feedback')
            update_map_btn = gr.Button(value="Update map")
            map = gr.Plot()

    with gr.Tab('Visualize data'):
        fetch_btn = gr.Button(value="Fetch images")
        l2a = gr.Gallery(preview=True)
        analyze = gr.Button(value="Analyze data")
        values = gr.Label(label="Expected yield")

    with gr.Tab('Ask queries'):
        with gr.Row():
            with gr.Column():
                query = gr.Textbox(label="Type your question here.", lines=8, interactive=True)
                submit = gr.Button(value="Submit")
            answer = gr.Textbox(label="Find your answer here", lines=10)

    #when find my location button is clicked
    my_loc.click(get_my_loc, None, [latitude, longitude])

    #when either latitude or longitude is changed
    latitude.change(is_location_valid, [latitude, longitude], feedback)
    longitude.change(is_location_valid, [latitude, longitude], feedback)

    #when the button to update the map is clicked
    update_map_btn.click(filter_map, [latitude, longitude], map)

    #To get images and show them in the gallery
    fetch_btn.click(fetch_images, [latitude, longitude], l2a)

    #Find the yield value, ndvi and others possible
    analyze.click(calculate_values, [latitude, longitude], values)

    #Submit the query and get your answer man:
    submit.click(answer_query, query, answer)

    #initial load
    demo.load(filter_map, [latitude, longitude], map)

demo.launch(show_api=False, share=False)