import numpy as np
import gradio as gr
import os
import pandas as pd
from datasets import load_dataset
import geopandas
import plotly.express as px
import plotly.graph_objects as go
Secret_token = os.getenv('token')

dataset_s = load_dataset("FDSRashid/taraf_by_year", token = Secret_token,split="train")
dataset = load_dataset("FDSRashid/edge_geo", token = Secret_token,split="train")

taraf_s = dataset_s.to_pandas()
merged_geo = dataset.to_pandas()
merged_geo["Coordinates"] = geopandas.GeoSeries.from_wkt(merged_geo["full_geom"])
geodf = geopandas.GeoDataFrame(merged_geo, geometry= 'Coordinates')


taraf_s = taraf_s.sort_values(['City', 'Year'], ascending=True)
cities = taraf_s['City'].unique().tolist()
min_year = int(taraf_s['Year'].min())
max_year = int(taraf_s['Year'].max())


def plot_taraf_map(yaxis,min_year = 0, max_year = 400, heat=False):
    if min_year > max_year:
        raise gr.Error('Minimum Year cannot be bigger than final ear!')
    filtered = geodf[(geodf['Year'] >= min_year) & (geodf['Year'] <= max_year)]
    temp = filtered[['City', yaxis]].groupby('City').sum().join(filtered[['City', 'Coordinates']].set_index('City'))
    temp[f'{yaxis}_CubeRoot'] = temp[yaxis].apply(np.cbrt)
    filtered = geopandas.GeoDataFrame(temp, geometry= 'Coordinates').reset_index()
    if heat:
        fig = px.density_mapbox(data_frame = filtered, lat = filtered.geometry.y, lon = filtered.geometry.x,z = yaxis, title = f'Map of {yaxis}', opacity = .5, hover_data = 'City',center=dict(lat=21.4241, lon=39.826168), radius=17, zoom=2, height =500)
    else:
        fig = px.scatter_mapbox(data_frame = filtered, lat = filtered.geometry.y, lon = filtered.geometry.x,size = f'{yaxis}_CubeRoot',color = yaxis, title = f'Map of {yaxis}', opacity = .5, hover_data = 'City',center=dict(lat=21.4241, lon=39.826168), zoom=2, height =500 )
    fig.update_layout(title_font_color = 'red',margin=dict(l=0, r=0, b=0), title_x = .5, mapbox_style="white-bg",
                      mapbox_layers=[
                          {"below": 'traces',
                           "sourcetype": "raster",
                           "sourceattribution": "United States Geological Survey",
                           "source": ["https://basemap.nationalmap.gov/arcgis/rest/services/USGSImageryOnly/MapServer/tile/{z}/{y}/{x}"] }])
    return fig




with gr.Blocks() as demo:
    Yaxis = gr.Dropdown(choices = ['Taraf', 'Hadith', 'Isnad'], value = 'Taraf', label = 'Variable to Display', info = 'Choose the variable to visualize.')  
    First_Year = gr.Slider(min_year, max_year, value = 0, label = 'Begining', info = 'Choose the first year to display Tarafs')
    Last_Year = gr.Slider(min_year, max_year, value = 400, label = 'End', info = 'Choose the last year to display Tarafs')
    heat = gr.Checkbox(label="Heat Map", info="View as Heat Map")
    btn = gr.Button('Submit')
    btn.click(fn = plot_taraf_map, inputs = [Yaxis, First_Year, Last_Year, heat], outputs = gr.Plot())
demo.launch()