import json
import pandas as pd
from urllib.request import urlopen
from urllib.parse import quote                                                                                                                                                                
import plotly.graph_objects as go
import plotly.express as px
import ipywidgets as widgets
from IPython.display import display
import numpy as np

"""
Explanation: Climate Projections
https://data.planetos.com/datasets/climate_risk_monthly_version1
Notebook for showing how to fetch Climate Risk data from Planet OS Datahub and make simple plots from the raw data. 
Note that this dataset is the first version of climate risk data. Right now, data is country level, but we plan to have county level data for some countries as well. Please, let us know all your requirements and what you would like to see in the future.
End of explanation
"""


apikey = open('APIKEY').read()
server = 'https://api.planetos.com/v1'
dataset = 'climate_risk_monthly_v1'

def get_data(server,dataset,model,scenarios,location,apikey):
    if " " in location:
        location = location.replace(' ', '%20')
    if type(scenarios) == list:
        scenarios = ','.join(scenarios)
    url = f'{server}/data/dataset_physical_values/{dataset}?node:station={quote(location)}&classifier:scenario={scenarios}&classifier:model={model}&count=100000&apikey={apikey}'
    data = json.loads(urlopen(url).read())
    return data
      
    
def data_to_pd(data,scenarios):
    variables = data['entries'][0]['data'].keys()
    data_out = {}
    for scenario in scenarios:
        #print (scenario)
        scenario_data = {}
        for var in variables:
            vardata = [t['data'][var] for t in data['entries'] if t['classifiers']['classifier:scenario'] == scenario]
            time = [pd.to_datetime(t['axes']['time']) for t in data['entries'] if t['classifiers']['classifier:scenario'] == scenario]
            if not 'time' in scenario_data:
                scenario_data['time'] = time
            scenario_data[var] = vardata
        data_pd = pd.DataFrame(scenario_data)
        data_out[scenario] = data_pd
            
    return data_out

def get_available_model_scenarios(server,dataset,model,location):
    if " " in location:
        location = location.replace(' ', '%20')
    url = f'{server}/data/dataset_physical_values/{dataset}?node:station={quote(location)}&classifier:model={model}&count=1&apikey={apikey}'
    #print (url)    
    data = json.loads(urlopen(url).read())
    scenarios = [s['classifiers']['classifier:scenario'] for s in data['entries']]
    return scenarios

def get_available_locations(server,dataset):
    url = f'{server}/datasets/{dataset}/stations?apikey={apikey}'
    data = json.loads(urlopen(url).read())
    return list(data['station'].keys())

def get_available_models(server, dataset, location):
    if " " in location:
        location = location.replace(' ', '%20')
    
    url = f'{server}/data/dataset_physical_values/{dataset}?node:station={quote(location)}&count=1&apikey={apikey}'
    data = json.loads(urlopen(url).read())
    models = np.unique([m['classifiers']['classifier:model'] for m in data['entries']])
    return models

"""
Explanation: Please put your datahub API key into a file called APIKEY and place it to the notebook folder or assign your API key directly to the variable API_key!
End of explanation
"""


locations = get_available_locations(server,dataset)

drop_down = widgets.Dropdown(
    options=sorted(locations),
    description='Available locations:',
    disabled=False
)

display(drop_down)

"""
Explanation: Getting available locations and making dropdown list to choose desired one
End of explanation
"""


location = drop_down.value
print ('Getting models for location: ', location)
models = get_available_models(server, dataset, location)
drop_down2 = widgets.Dropdown(
    options=models,
    description='Available models:',
    disabled=False
)
display(drop_down2)

for mo in models:
    scenarios = get_available_model_scenarios(server,dataset,mo,location)
    print (mo, scenarios)

"""
Explanation: Dropdown where you can choose a model you want to use. Also, in addition to models, you can choose between quantiles 5-75.
End of explanation
"""


model = drop_down2.value
print (model)
scenarios = get_available_model_scenarios(server,dataset,model,location)
print (scenarios)
data = get_data(server,dataset,model,scenarios,location,apikey)
scenarios_data_pd_list = data_to_pd(data,scenarios)

"""
Explanation: Getting avaiable scenarios of the model and reading it to pandas dataframes.
End of explanation
"""


fig = go.Figure()
for scenario in scenarios:
    data_pd = scenarios_data_pd_list[scenario]
    fig.add_traces(go.Scatter(
        x=data_pd['time'], 
        y=data_pd['tas'].values, 
        mode='lines', 
        name = f'{location} {scenario}',
        legendgroup = scenario
        ))
fig.update_layout(title = f'{location} Monthly Temperature [C]')
fig.show()

"""
Explanation: Making a plot of Monthly Temperatures from chosen location and model.
End of explanation
"""


fig = go.Figure()
for scenario in scenarios:
    data_pd = scenarios_data_pd_list[scenario]
    fig.add_traces(go.Scatter(
        x=data_pd.time, 
        y=data_pd['pr'], 
        mode='lines', 
        name = f'{location} {scenario}',
        legendgroup = scenario
        ))
fig.update_layout(title = f'{location} Monthly Precipitation [kg m-2 s-1]')
fig.show()

"""
Explanation: Plot precipitation data
End of explanation
"""


# quantile_keys = [f for f in models if 'quantile' in f]
# scenarios = ['historical']#get_available_model_scenarios(server,dataset,model,location)
# var = 'tas'
# for sc in scenarios:
#     q_data = {}
#     for q_key in sorted(quantile_keys):   
#         data = get_data(server,dataset,q_key,sc,location,apikey)
#         scenarios_data_pd_list = data_to_pd(data,[sc])
        
#         q_data[q_key] = scenarios_data_pd_list[sc][var].values
        
#         if not 'time' in q_data:
#             q_data['time'] = scenarios_data_pd_list[sc]['time'].values
    
#     q_data_pd = pd.DataFrame(q_data)
#     q_data_pd = q_data_pd.set_index('time')
#     q_data_resampled = q_data_pd.resample('10Y').mean().transpose()
#     fig = px.box(q_data_resampled)
#     fig.update_layout(title = f'{location} {var} Quantiles Scenario: {sc}')

#     fig.show()

"""
Explanation: Make graph about quantiles.
Commented out right now as quantiles are work on progress with subarea data.
End of explanation
"""


