# -*- coding: utf-8 -*-
""" Data Handler """

import plotly.graph_objs as go
from dash import html
import dash_leaflet as dl
import dash_leaflet.express as dlx
from dash_extensions.javascript import arrow_function
from dash_extensions.javascript import Namespace
from matplotlib.colors import ListedColormap
import numpy as np
from netCDF4 import Dataset as nc_file
import pandas as pd
import json
import orjson
from datetime import datetime
from dateutil.relativedelta import relativedelta
from collections import OrderedDict
import calendar
import time
import os
from datetime import datetime as dt
from datetime import timedelta

from utils import concat_dataframes
from utils import retrieve_timeseries
from utils import retrieve_single_point
from utils import get_colorscale

from pathlib import Path
from flask_caching import Cache
import uuid
import socket

#SETUP CACHE
cache_dir = "/dev/shm/{}".format(str(uuid.uuid1()))
Path(cache_dir).mkdir(parents=True, exist_ok=True)

cache_config = {
    "DEBUG": True,
    "CACHE_TYPE": "FileSystemCache",
    "CACHE_DIR": cache_dir,
}

cache = Cache(config=cache_config)
cache_timeout = 86400

#SETUP BASE URL 
HOSTNAME = socket.gethostbyname_ex(socket.gethostname())[0]

if HOSTNAME in ('bscesdust03.bsc.es', 'dust.hqads.aemet.es'):
    PATHNAME = '/daily_dashboard/'
else:
    PATHNAME = '/dashboard/'

DIR_PATH = os.path.dirname(os.path.realpath(__file__))

DEBUG = True  # False

VARS = json.load(open(os.path.join(DIR_PATH, 'conf/vars.json')))
MODELS = json.load(open(os.path.join(DIR_PATH, 'conf/models.json')))
OBS = json.load(open(os.path.join(DIR_PATH, 'conf/obs.json')))
WAS = json.load(open(os.path.join(DIR_PATH, 'conf/was.json')))
PROB = json.load(open(os.path.join(DIR_PATH, 'conf/prob.json')))
DATES = json.load(open(os.path.join(DIR_PATH, 'conf/dates.json')))
ALIASES = json.load(open(os.path.join(DIR_PATH, 'conf/aliases.json')))
STATS_CONF = json.load(open(os.path.join(DIR_PATH, 'conf/stats.json')))
STYLES = json.load(open(os.path.join(DIR_PATH, 'conf/styles.json')))
ALL_COLORS = json.load(open(os.path.join(DIR_PATH, 'conf/colors.json')))
MODEBARS = json.load(open(os.path.join(DIR_PATH, 'conf/modebars.json')))

MODEBAR_CONFIG = MODEBARS['config']
MODEBAR_CONFIG_TS = MODEBARS['config_ts']
MODEBAR_LAYOUT = MODEBARS['layout']
MODEBAR_LAYOUT_TS = MODEBARS['layout_ts']
INFO_STYLE = MODEBARS['info_style']

COLORS = ALL_COLORS['std']
COLORS_NORGB = ALL_COLORS['std_norgb']
COLORS_PROB = ALL_COLORS['prob']

COLORMAP = ListedColormap(COLORS_NORGB)
COLORMAP_PROB = ListedColormap(COLORS_PROB)

START_DATE = DATES['start_date']
DELAY = DATES['delay']['delayed']
DELAY_DATE = DATES['delay']['start_date']
END_DATE = DATES['end_date']
END_DATE = END_DATE or (DELAY and (dt.now() -
    timedelta(days=1)).strftime("%Y%m%d") or dt.now().strftime("%Y%m%d"))

ROUTE_DEFAULTS = {
        'tab':['forecast'],
        'var': ['OD550_DUST'],
        'model': ['median'],
        'for_option': ['models'],
        'eval_option': ['nrt'],
        'obs_option': ['rgb'],
        'country': ['burkinafaso'],
        'download': [None],
        'date': [END_DATE]
        }

STATS = OrderedDict({
    'bias': 'MBE',
    'corr': 'r',
    'rmse': 'RMSE',
    'frge': 'FGE',
    'totn': 'TOTAL CASES'
})

GRAPH_HEIGHT = 92.8

# Frequency = 3 Hourly
FREQ = 3
OPACITY = 0.7

DEFAULT_VAR = 'OD550_DUST'
DEFAULT_MODEL = 'median'

DISCLAIMER_NO_FORECAST = [html.Span(html.P("""Dust data ©2023 WMO Barcelona Dust Regional Center."""), id='forecast-disclaimer')]

DISCLAIMER_MODELS = [html.Span(html.P("""Dust data ©2023 WMO Barcelona Dust Regional Center."""), id='forecast-disclaimer')]

DISCLAIMER_OBS = html.P("""Aerosol data ©2023 WMO Barcelona Dust Regional Center, NASA.""")

GEOJSON_TEMPLATE = "{}/geojson/{}/{:02d}_{}_{}.geojson"
NETCDF_TEMPLATE = "{}/netcdf/{}{}.nc"


class Observations1dHandler(object):
    """ Class which handles 1D obs data """

    def __init__(self, sdate, edate, obs):
        fday = sdate[:-2] + '01'
        lday = edate[:-2] + str(calendar.monthrange(int(edate[:4]), int(edate[4:6]))[1])
        date_range = pd.date_range(fday, lday, freq='M')
        months = [d.strftime("%Y%m") for d in date_range.to_pydatetime()]
        filepaths = ["{}.nc".format(os.path.join(OBS[obs]['path'], 'netcdf',
            OBS[obs]['template'].format(OBS[obs]['obs_var'], month))) for month
            in months]
        input_files = [nc_file(filepath) for filepath in filepaths]
        if 'longitude' in input_files[0].variables:
            lon_var = 'longitude'
            lat_var = 'latitude'
        else:
            lon_var = 'lon'
            lat_var = 'lat'

        lon = input_files[0].variables[lon_var][:]
        lat = input_files[0].variables[lat_var][:]
        self.varname = [var for var in input_files[0].variables if var == OBS[obs]['obs_var']][0]
        if DEBUG: print('VARNAME', self.varname)

        sites = pd.read_csv(os.path.join('./conf/',
            OBS[obs]['sites']))
#         sites = pd.read_table(os.path.join('./conf/',
#             OBS[obs]['sites']), delimiter=r"\s+", engine='python')

        idxs, self.station_names = np.array([[idx, st_name[~st_name.mask].tobytes().decode('utf-8')]
                              for idx, st_name in
                              enumerate(input_files[0].variables['station_name'][:])
                              if st_name[~st_name.mask].tobytes().decode('utf-8').upper() in map(str.upper, sites['SITE'])]
                              ).T

        if DEBUG: print('IDXS', idxs)
        if DEBUG: print('ST_NAMES', self.station_names)

        self.clon = lon[idxs.astype(int)]
        self.clat = lat[idxs.astype(int)]

        self.values = {
            self.varname: np.concatenate([input_file.variables[self.varname][:, idxs.astype(int)] for input_file in input_files])
        }


    def generate_obs1d_tstep_trace(self, var):
        """ Generate trace to be added to data, per variable and timestep """
        varname = self.varname
        val = np.ma.masked_where(self.values[varname]<0., self.values[varname])
        notnan = (np.array([i for i in range(val.shape[1]) if not val[:, i].mask.all()]),)
        clon = self.clon[notnan]
        clat = self.clat[notnan]
        cstations = self.station_names[notnan]
        df = pd.DataFrame({
            'lon': clon.round(2),
            'lat': clat.round(2),
            'stations': cstations
            })  # .T, columns=['lon', 'lat', 'station'])
        dicts = df.to_dict('rows')
        geojson = dlx.dicts_to_geojson(dicts, lon="lon")

        # Geojson rendering logic, must be JavaScript as it is executed in clientside.
        ns = Namespace("evaluationTab", "evaluationMaps")
        point_to_layer = ns("pointToLayer")
        # Create geojson.
        return df, dl.GeoJSON(data=geojson,
                options=dict(
                    pointToLayer=point_to_layer,
                ),
                hideout=dict(
                    circleOptions=dict(
                        fillOpacity=OPACITY,
                        stroke=False,
                        fillColor='#f0b450',
                        radius=8),
                )
            )


class ObsTimeSeriesHandler(object):
    """ Class to handle evaluation time series """

    def __init__(self, obs, start_date, end_date, variable, models=None):
        self.obs = obs
        if models is None:
            models = list(MODELS.keys())
        self.model = models
        self.variable = variable
        self.dataframe = []
        if DEBUG: print("ObsTimeSeries", start_date, end_date)
        self.date_range = pd.date_range(start_date, end_date, freq='D')

        fname_tpl = os.path.join(OBS[obs]['path'],
                                'feather',
                                '{{}}-{dat}-{{}}_interp.ft')

        months = np.unique([d.strftime("%Y%m") for d in self.date_range.to_pydatetime()])

        self.date_index = pd.date_range('{}01'.format(months[0]),
                '{}01'.format((datetime.strptime(months[-1], '%Y%m') +
                    relativedelta(days=31)).strftime('%Y%m')), freq='3H')

        if DEBUG:
            print('MONTHS', months)
        if DEBUG:
            print('DATE_INDEX', self.date_index)

        fname_obs = fname_tpl.format(dat=obs)
        notnans, obs_df = concat_dataframes(fname_obs, months, variable,
            rename_from=OBS[obs]['obs_var'])
        self.dataframe.append(obs_df)

        for mod in self.model:
            fname_mod = fname_tpl.format(dat=mod)
            _, mod_df = concat_dataframes(fname_mod, months, variable,
                    rename_from=None, notnans=notnans)
            self.dataframe.append(mod_df)


    def retrieve_timeseries(self, idx, st_name, model):

        old_indexes = self.dataframe[0]['station'].unique()
        new_indexes = np.arange(old_indexes.size)
        dict_idx = dict(zip(new_indexes, old_indexes))
        if DEBUG: print("RETRIEVE TS", idx, dict_idx[idx], st_name, model)
        fig = go.Figure()
        for mod, df in zip([self.obs]+self.model, self.dataframe):
            if df is None:
                continue
            if DEBUG: print("MOD", mod, "COLS", df.columns)
            if df.columns[-1].upper() == self.variable:
                df = df.rename(columns = { df.columns[-1]: self.variable })

            if DEBUG: print("BUILDING TIME-SERIES")
            try:
                tmp_df = df[df['station']==dict_idx[idx]].set_index('time')
                tmp_df.drop_duplicates(keep='first')
                timeseries = \
                   tmp_df.reindex(self.date_index)
                if timeseries[self.variable].isnull().all():
                    continue
            except Exception as e:
                if DEBUG: print("ERROR timeseries", str(e))
                continue

            if DEBUG: print("SELECTING COORDS")
            if 'lat' in df.columns:
                lat_col = 'lat'
                lon_col = 'lon'
            else:
                lat_col = 'latitude'
                lon_col = 'longitude'

            if mod == self.obs:
                if DEBUG: print("OBSERVATION")
                sc_mode = 'markers'
                marker = {'size': 10, 'symbol': "triangle-up-dot", 'color': '#f0b450'}
                line = {}
                visible = True
                name = mod.upper()
            else:
                if DEBUG: print("MODEL", mod)
                try:
                    sc_mode = 'lines'
                    marker = {}
                    line = { 'color': MODELS[mod]['color'] }
                    visible = (mod == model) and True or 'legendonly'
                    name = "{}".format(MODELS[mod]['name'])
                    cur_lat = round(timeseries[lat_col].dropna()[0], 2)
                    cur_lon = round(timeseries[lon_col].dropna()[0], 2)
                except Exception as e:
                    if DEBUG: print("ERROR MODEL", mod, str(e), timeseries[lat_col].dropna())
                    continue

            if mod == 'median':
                line['dash'] = 'dash'

            if DEBUG: print("ADD TRACE")
            fig.add_trace(dict(
                type='scatter',
                name=name,
                x=timeseries.index,
                y=timeseries[self.variable].round(2),
                mode=sc_mode,
                marker=marker,
                line=line,
                visible=visible,
                connectgaps=False
                )
            )

        title = "{} @ {} (lat = {:.2f}, lon = {:.2f})".format(
            VARS[self.variable]['name'], st_name, cur_lat, cur_lon,
        )

        fig.update_layout(
            title=dict(text=title, x=0.45, y=0.99),
            # uirevision=True,
            autosize=True,
            showlegend=True,
            plot_bgcolor='#F9F9F9',
            font_size=12,
            hovermode="x",        # highlight closest point on hover
            margin={"r": 10, "t": 35, "l": 10, "b": 10},
        )
        fig.update_xaxes(
            range=[self.date_range[0], self.date_range[-1]],
            rangeslider_visible=True,
            rangeselector=dict(
                buttons=list([
                    dict(step="all", label="all"),
                    dict(count=14, label="2w",
                         step="day", stepmode="backward"),
                    dict(count=7, # label="1w",
                         step="day", stepmode="backward"),
                ])
            )
        )

        if DEBUG: print('FIG TYPE', type(fig))
        return fig


class TimeSeriesHandler(object):
    """ Class to handle forecast time series """

    def __init__(self, model, date, variable):
        if isinstance(model, str):
            model = [model]
        self.model = model
        self.variable = variable
        self.fpaths = []
        try:
            self.month = datetime.strptime(date, "%Y%m%d").strftime("%Y%m")
            self.currdate = datetime.strptime(date, "%Y%m%d").strftime("%Y%m%d")
        except:
            self.month = datetime.strptime(date, "%Y-%m-%d").strftime("%Y%m")
            self.currdate = datetime.strptime(date, "%Y-%m-%d").strftime("%Y%m%d")

    def retrieve_single_point(self, tstep, lat, lon, model=None):

        if not model:
            model = self.model[0]

        if DEBUG: print("----------", model)

        method = 'netcdf'
        if (MODELS[model]['start'] == 12 and not DELAY and DELAY_DATE and (datetime.strptime(self.currdate, "%Y%m%d") >= datetime.strptime(DELAY_DATE, "%Y%m%d"))) or \
                (MODELS[model]['start'] == 12 and not DELAY and not DELAY_DATE):
            mod_date = (datetime.strptime(self.currdate, "%Y%m%d") -
                    timedelta(days=1)).strftime("%Y%m%d")
        else:
            mod_date = self.currdate
        path_template = '{}{}.nc'.format(mod_date,
                                         MODELS[model]['template'],
                                         self.variable)

        fpath = os.path.join(MODELS[model]['path'], method, path_template)

        return retrieve_single_point(fpath, tstep, lat, lon,
                                     self.variable)

    def retrieve_timeseries(self, lat, lon, model=None, method='netcdf', forecast=False):

        if not model:
            model = self.model

        # if DEBUG: print("----------", model)

        obs_eval = model[0] not in MODELS and model[0] in OBS
        if obs_eval:
            all_models = [model[0]] + list(MODELS.keys())
        else:
            all_models = list(MODELS.keys())

        for mod in all_models:
            if obs_eval:
                filedir = OBS[model[0]]['path']
                path_tpl = '{}-{}-{}_interp.ft'  # 202010-median-OD550_DUST_interp.ft
            else:   # if mod in MODELS:
                filedir = MODELS[mod]['path']
                path_tpl = '{}-{}-{}.ft'  # 202010-median-OD550_DUST_interp.ft

            if method == 'feather':
                path_template = path_tpl.format(self.month, mod, self.variable)
            elif method == 'netcdf':
                path_template = '{}*{}.nc'.format(self.month, MODELS[mod]['template'], self.variable)

            if forecast:
                method = 'netcdf'
                if (MODELS[mod]['start'] == 12 and not DELAY and DELAY_DATE and (datetime.strptime(self.currdate, "%Y%m%d") >= datetime.strptime(DELAY_DATE, "%Y%m%d"))) or \
                        (MODELS[mod]['start'] == 12 and not DELAY and not DELAY_DATE):
                    mod_date = (datetime.strptime(self.currdate, "%Y%m%d") -
                            timedelta(days=1)).strftime("%Y%m%d")
                else:
                    mod_date = self.currdate
                path_template = '{}{}.nc'.format(mod_date, MODELS[mod]['template'], self.variable)

            fpath = os.path.join(filedir,
                                 method,
                                 path_template)
            self.fpaths.append(fpath)

        title = "{} @ lat = {} and lon = {}".format(
            VARS[self.variable]['name'], round(lat, 2), round(lon, 2)
        )

        mul = VARS[self.variable]['mul']

        fig = go.Figure()

        for mod, fpath in zip(all_models, self.fpaths):
            # print(mod, fpath)
            if mod not in MODELS and mod in OBS:
                variable = OBS[mod]['obs_var']
            else:
                variable = self.variable

            if not os.path.exists(fpath):
                if DEBUG: print("NOT retrieving", fpath, "File doesn't exist.")
                continue

            if DEBUG: print('Retrieving *** FPATH ***', fpath)
            try:
                ts_lat, ts_lon, ts_index, ts_values = retrieve_timeseries(
                        fpath, lat, lon, variable, method=method, forecast=forecast)
            except Exception as e:
                if DEBUG: print("NOT retrieving", fpath, "ERROR:", str(e))
                continue

            if forecast is True or mod in ('cams', 'ema-regcm4'):
                ts = pd.Series(index=ts_index, data=ts_values)
            else:
                date_index = pd.date_range('{}01'.format(self.month),
                        '{}01'.format((datetime.strptime(self.month, '%Y%m') +
                            relativedelta(days=31)).strftime('%Y%m')), freq='3H')
                ts = pd.Series(index=ts_index, data=ts_values).reindex(date_index)

            if isinstance(ts_lat, np.ndarray):
                ts_lat = float(ts_lat)
                ts_lon = float(ts_lon)
            if isinstance(ts.values, np.ndarray):
                ts_values = (ts.values*mul).round(2)
            else:
                ts_values = round((ts.values*mul), 2)

            if obs_eval and mod == model[0]:
                sc_mode = 'markers'
                marker = {'size': 12, 'symbol': "triangle-up-dot", 'color': '#f0b450'}
                line = {}
                visible = True
                name = "{}".format(mod.upper())
            elif obs_eval:
                sc_mode = 'lines'
                marker = {}
                line = { 'color': MODELS[mod]['color'] }
                if mod in model:
                    visible = True
                else:
                    visible = 'legendonly'
                name = "{}".format(MODELS[mod]['name'])
            else:
                sc_mode = 'lines'
                marker = {}
                line = { 'color': MODELS[mod]['color'] }
                if mod in model:
                    visible = True
                else:
                    visible = 'legendonly'
                name = "{} ({}, {})".format(
                        MODELS[mod]['name'], round(ts_lat, 2), round(ts_lon, 2))

            if mod == 'median':
                line['dash'] = 'dash'
            else:
                line['dash'] = 'solid'

            fig.add_trace(dict(
                    type='scatter',
                    name=name,
                    x=ts.index,
                    y=ts_values,
                    mode=sc_mode,
                    marker=marker,
                    line=line,
                    visible=visible,
                )
            )

        fig.update_layout(
            title=dict(text=title, x=0.45, y=.99),
            yaxis=dict(exponentformat="none"),
            # uirevision=True,
            autosize=True,
            showlegend=True,
            plot_bgcolor='#F9F9F9',
            font_size=12,
            # hovermode="closest",        # highlight closest point on hover
            hovermode="x",        # highlight closest point on hover
            margin={"r": 10, "t": 35, "l": 10, "b": 10},
        )

        return fig


class FigureHandler(object):
    """ Class to manage the figure creation """

    def __init__(self, model=None, selected_date=None):
        """ FigureHandler init """

        self.st_time = time.time()
        if isinstance(model, list):
            model = model[0]

        self.model = model

        if self.model not in MODELS and self.model in OBS:
            self.filedir = OBS[self.model]['path']
            self.filevars = [OBS[self.model]['obs_var']]
            self.confvars = [OBS[self.model]['mod_var']]
            filetpl = OBS[self.model]['template'].format(OBS[self.model]['obs_var'], selected_date) + '.nc'
            filepath = os.path.join(self.filedir, 'netcdf', filetpl)

        elif self.model in MODELS:
            if DEBUG: print("MODEL", model)
            self.filedir = MODELS[self.model]['path']
            self.filevars = VARS
            self.confvars = None

            filepath = NETCDF_TEMPLATE.format(
                self.filedir,
                selected_date,
                MODELS[self.model]['template']
            )
        else:
            self.model = None
            self.filedir = None
            self.filevars = None
            self.confvars = None
            filepath = None
            self.bounds = None

        if filepath is None or not os.path.exists(filepath):
            self.filedir = None
            self.filevars = None
            self.confvars = None
            filepath = None
            self.bounds = None
            self.varlist = None
            self.rdatetime = None
            self.tim = [0]
        elif filepath is not None:
            self.input_file = nc_file(filepath)
            if 'lon' in self.input_file.variables:
                lon = self.input_file.variables['lon'][:]
                lat = self.input_file.variables['lat'][:]
            else:
                lon = self.input_file.variables['longitude'][:]
                lat = self.input_file.variables['latitude'][:]
            time_obj = self.input_file.variables['time']
            self.tim = time_obj[:]
            tim_units = time_obj.units.split()
            if len(tim_units) == 3:
                self.what, _, rdate = tim_units
                rtime = "00:00"
            elif len(tim_units) >= 4:
                self.what, _, rdate, rtime = tim_units[:4]
            if len(rtime) > 5:
                rtime = rtime[:5]
            self.rdatetime = datetime.strptime("{} {}".format(rdate, rtime),
                                               "%Y-%m-%d %H:%M")
            varlist = [var for var in self.input_file.variables if (var.upper() in self.filevars) or (var in self.filevars)]
            self.varlist = varlist
            if DEBUG: print('VARLIST', varlist)
            self.xlon, self.ylat = np.meshgrid(lon, lat)

        if not self.varlist:
            self.varlist = VARS.keys()

        if DEBUG: print('VARLIST', self.varlist, 'CONFVAR', self.confvars)

        if self.confvars is not None:
            self.bounds = {
                varname.upper(): np.array(VARS[confvar]['bounds']).astype('float32')
                for varname, confvar in zip(self.filevars, self.confvars) if (varname.upper() in self.varlist) or (varname in self.varlist)
            }
        else:
            self.bounds = {
                varname.upper(): np.array(VARS[varname.upper()]['bounds']).astype('float32')
                for varname in self.varlist
            }

        self.colormaps = {
            varname.upper(): get_colorscale(self.bounds[varname.upper()], COLORMAP)
            for varname in self.varlist
        }
        # print(varlist, self.confvars, self.filevars, self.bounds, self.colormaps)

        if selected_date:
            self.selected_date_plain = selected_date

            self.selected_date = datetime.strptime(
                selected_date, "%Y%m%d").strftime("%Y-%m-%d")

            if not self.rdatetime:
                self.rdatetime = datetime.strptime(selected_date, "%Y%m%d")
                self.what = 'hours'

        self.fig = None
        if DEBUG: print("FILEPATH", filepath)
        if DEBUG: print("FILEDIR", self.filedir)

    def get_center(self, center=None):
        """ Returns center of map """
        if center is None and hasattr(self, 'ylat'):
            center =  [
                    (self.ylat.max()-self.ylat.min())/2 + self.ylat.min() + (self.ylat.max()-self.ylat.min())/6,
                    (self.xlon.max()-self.xlon.min())/2 + self.xlon.min(),
            ]
        elif center is None:
            center = [ 30, 15 ]

        return center

    def set_data(self, varname, tstep=0):
        """ Set time dependent data """
        if self.model in OBS:
            mul = VARS[OBS[self.model]['mod_var']]['mul']
        else:
            mul = VARS[varname]['mul']

        realvar = [var for var in self.varlist if var.upper()==varname.upper()][0]
        # if DEBUG: print("***", mul, realvar, "***")
        var = self.input_file.variables[realvar][tstep]*mul
        idx = np.where((~var.ravel().mask) & (var.ravel() >= self.bounds[varname.upper()][0]))  # !=-9.e+33)
        xlon = self.xlon.ravel()[idx]
        ylat = self.ylat.ravel()[idx]
        var = var.ravel()[idx]
        # if DEBUG: print("***", xlon.shape, ylat.shape, var.shape, "***")
        return xlon.data, ylat.data, var.data

    def retrieve_cdatetime(self, tstep=0):
        """ Retrieve data from NetCDF file """
        tstep = int(tstep)
        tim = int(self.tim[tstep])
        if self.what == 'days':
            cdatetime = self.rdatetime + relativedelta(days=tim)
        elif self.what == 'hours':
            cdatetime = self.rdatetime + relativedelta(hours=tim)
        elif self.what == 'minutes':
            cdatetime = self.rdatetime + relativedelta(minutes=tim)
        elif self.what == 'seconds':
            cdatetime = self.rdatetime + relativedelta(seconds=tim)

        return cdatetime

    def generate_contour_tstep_trace_leaflet(self, varname, tstep=0):
        """ Generate trace to be added to data, per variable and timestep """
        from dash_server import app

        if varname not in VARS and self.model in OBS:
            name = VARS[OBS[self.model]['mod_var']]['name']
        else:
            name = VARS[varname]['name']
        if self.bounds:
            bounds = self.bounds[varname.upper()]
        else:
            bounds = [0, 1]
        if self.model in OBS:
            data_path = os.path.basename(OBS[self.model]['path'][:-1])
        else:
            data_path = os.path.basename(MODELS[self.model]['path'])

        if DEBUG: print(data_path)
        colorscale = COLORS

        geojson_url = app.get_asset_url(os.path.join('geojsons',
            GEOJSON_TEMPLATE.format(data_path,
                self.selected_date_plain, tstep, self.selected_date_plain,
                varname)))
        if DEBUG: print("MODEL", self.model, "GEOJSON_URL", geojson_url)

        style = dict(weight=0, opacity=0, color='white', dashArray='', fillOpacity=OPACITY)

        # Create colorbar.
        ctg = ["{:d}".format(int(cls)) if cls.as_integer_ratio()[1] == 1 else "{:.1f}".format(cls) for i, cls in enumerate(bounds[1:-1])]
        if DEBUG: print("BOUNDS", bounds)
        if DEBUG: print("CTG", ctg)
        indices = list(range(len(ctg) + 2))
        colorbar = dl.Colorbar(
                min=0, max=len(ctg)+1,
                classes=indices,
                colorscale=colorscale,
                tickValues=indices[1:-1],
                tickText=ctg,
                position='topleft',
                width=270,
                height=15,
                style={ 'top': '55px' }
                )

        # Geojson rendering logic, must be JavaScript as it is executed in clientside.
        ns = Namespace("forecastTab", "forecastMaps")
        style_handle = ns("styleHandle")

        geojson = dl.GeoJSON(
                url=geojson_url,
                options=dict(style=style_handle),
                hideout=dict(colorscale=colorscale, bounds=bounds, style=style, colorProp="value")
                )  # url to geojson file

        return geojson, colorbar


    def generate_contour_tstep_trace(self, varname, tstep=0):
        """ Generate trace to be added to data, per variable and timestep """
        from dash_server import app

        geojson_file = GEOJSON_TEMPLATE.format(self.filedir,
                self.selected_date_plain, tstep, self.selected_date_plain, varname)

        if os.path.exists(geojson_file):
            geojson = orjson.loads(open(geojson_file).read())
        else:
            if DEBUG: print('ERROR', geojson_file, 'not available')
            geojson = {
                    "type": "FeatureCollection",
                    "features": []
                    }

        if varname not in VARS and self.model in OBS:
            name = VARS[OBS[self.model]['mod_var']]['name']
        else:
            name = VARS[varname]['name']
        # if DEBUG: print(self.bounds)
        if self.bounds:
            bounds = self.bounds[varname.upper()]
        else:
            bounds = [0, 1]
        loc_val = [
            (
                feature['id'],
                np.around(feature['properties']['value'], 2),
            )
            for feature in geojson['features']
            if feature['geometry']['coordinates']
        ]
        locations, values = np.array(loc_val).T if loc_val else ([], [])
        # if DEBUG: print(varname, self.colormaps[varname], values)
        return dict(
            type='choroplethmapbox',
            name=name+'_contours',
            geojson=geojson,
            z=values,
            ids=locations,
            locations=locations,
            zmin=bounds[0],
            zmax=bounds[-1],
            colorscale=self.colormaps[varname.upper()],
            showscale=False,
            showlegend=False,
            hoverinfo='none',
            marker=dict(
                opacity=0.6,
                line_width=0,
            ),
            colorbar=None,
        )

    def generate_var_tstep_trace_leaflet(self, varname=None, tstep=0):
        """ Generate trace to be added to data, per variable and timestep """
        colorscale = COLORS
        xlon, ylat, val = self.set_data(varname, tstep)
        df = pd.DataFrame(np.array([xlon, ylat, val]).T.round(2), columns=['lon', 'lat', 'value'])
        dicts = df.to_dict('rows')
        for item in dicts:
            item["tooltip"] = \
                    "Lat {:.2f} Lon {:.2f} Val {:.2f}".format(item['lat'], item['lon'], item['value'])
        geojson = dlx.dicts_to_geojson(dicts, lon="lon")
        # geobuf = dlx.geojson_to_geobuf(geojson)

        if DEBUG: print("GEOBUF CREATED ***********")
        # Geojson rendering logic, must be JavaScript as it is executed in clientside.
        ns = Namespace("forecastTab", "forecastMaps")
        point_to_layer = ns("pointToLayer")
        bind_tooltip = ns("bindTooltip")
        if DEBUG: print("BIND", str(bind_tooltip))
        # Create geojson.
        # return dl.GeoJSON(data=geobuf, format="geobuf",
        return dl.GeoJSON(data=geojson,
                options=dict(
                    pointToLayer=point_to_layer,
                    # onEachFeature=bind_tooltip,
                ),  # how to draw points
                hideout=dict(
                    colorProp='value',
                    circleOptions=dict(
                        fillOpacity=0,
                        stroke=False,
                        radius=0),
                    min=0,
                    max=val.max(),
                    colorscale=colorscale)
                )
        

    def generate_var_tstep_trace(self, varname=None, tstep=0):
        """ Generate trace to be added to data, per variable and timestep """
        if not varname:
            return dict(
                type='scattermapbox',
                # below='',
                lon=[15],
                lat=[30],
                hoverinfo='none',
                opacity=0,
                showlegend=False,
                marker=dict(
                    showscale=False,
                    size=0,
                    colorbar=None,
                ),
            )
        xlon, ylat, val = self.set_data(varname, tstep)
        if self.model in OBS:
            name = OBS[self.model]['name']
        else:
            name = MODELS[self.model]['name']
        if DEBUG: print("***", name, "***")
        return dict(
            type='scattermapbox',
            # below='',
            lon=xlon,
            lat=ylat,
            text=val,
            name=name,
            hovertemplate="lon: %{lon:.2f}<br>lat: %{lat:.2f}<br>" +
            "value: %{text:.2f}",
            opacity=OPACITY,
            showlegend=False,
            marker=dict(
                # autocolorscale=True,
                showscale=False,
                color=val,
                # opacity=0.6,
                size=0,
                colorscale=self.colormaps[varname.upper()],
                cmin=self.bounds[varname.upper()][0],
                cmax=self.bounds[varname.upper()][-1],
                colorbar=None,
            ),
        )

    def get_title(self, varname, tstep=0):
        """ return title according to the date """
        if self.model in OBS:
            name = OBS[self.model]['name']
            title = OBS[self.model]['title']
        else:
            name = MODELS[self.model]['name']
            title = VARS[varname]['title']
        rdatetime = self.retrieve_cdatetime(tstep=0)
        cdatetime = self.retrieve_cdatetime(tstep)
        return r'{} {}'.format(name, title % {
            'rhour':  rdatetime.strftime("%H"),
            'rday':   rdatetime.strftime("%d"),
            'rmonth': rdatetime.strftime("%b"),
            'ryear':  rdatetime.strftime("%Y"),
            'shour':  cdatetime.strftime("%H"),
            'sday':   cdatetime.strftime("%d"),
            'smonth': cdatetime.strftime("%b"),
            'syear':  cdatetime.strftime("%Y"),
            'step':   "{:02d}".format(tstep*FREQ),
        })

    def hour_to_step(self, hour):
        """ Convert hour to relative tstep """
        cdatetime = self.rdatetime.date() + relativedelta(hours=hour)

        for step in range(len(self.tim)):
            if self.retrieve_cdatetime(step) == cdatetime:
                return step

        return 0

    def retrieve_var_tstep(self, varname=None, tstep=0, hour=None, static=True, aspect=(1,1), center=None, selected_tiles='carto-positron', zoom=None, layer=None, tag='empty'):
        """ run plot """

        if hour is not None:
            tstep = int(self.hour_to_step(hour))
        else:
            tstep = int(tstep)

        if varname is not None and self.model in OBS:
            varname = OBS[self.model]['obs_var']

        if DEBUG: print('VARNAME', varname)

        if varname and self.filedir:
            if DEBUG: print('Adding contours ...')
            try:
                cont_time = time.time()
                geojson_contours, colorbar = self.generate_contour_tstep_trace_leaflet(varname, tstep)
                if DEBUG: print("****** CONTOURS EXEC TIME", str(time.time() - cont_time))
            except Exception as err:
                if DEBUG: print("----------- ERROR:", str(err))
                self.filedir = None
                # if DEBUG: print('ERROR: geojson {}'.format(geojson_url))
                geojson_contours = None
                colorbar = None
        else:
            if DEBUG: print('Adding one point ...')
            geojson_contours = None
            colorbar = None

        if DEBUG: print("ASPECT", aspect)
        if center is None:
            center = self.get_center(center)
        if zoom is None:
            zoom = 3.5 -(aspect[0]-aspect[0]*0.4)
        if colorbar is not None:
            colorbar.width = 320 - 15 * aspect[0]
            colorbar.style = {'overflow':'hidden', 'white-space':'nowrap'}
            if aspect == (3,4):
                colorbar.width = 320 - 38 * aspect[0]
            colorbar.height = 8
        if DEBUG: print("ZOOM", zoom)
        if DEBUG: print("CENTER", center)


        if DEBUG: print('Update layout ...')
        if not varname:
            fig_title=html.P("")
        elif varname and not self.filedir:
            if self.model in MODELS:
                curr_name = MODELS[self.model]['name']
            elif self.model in OBS:
                curr_name = OBS[self.model]['name']
            else:
                curr_name = ''

            fig_title = html.P(html.B("{} - DATA NOT AVAILABLE".format(curr_name)))
        else:
            fig_title = html.P(html.B(
                [
                    item for sublist in self.get_title(varname, tstep).split('<br>') for item in [sublist, html.Br()]
                ][:-1]
            ))
        if self.model is not None:
            CUR_INFO_STYLE = INFO_STYLE.copy()
            if colorbar is not None:
                CUR_INFO_STYLE['width'] = str(colorbar.width) + "px"
            if aspect[0] > 2:
                CUR_INFO_STYLE['fontSize'] = "{}px".format(int(INFO_STYLE['fontSize'][:-2])-aspect[0]+ 0.3)
            info = html.Div(
                children=fig_title,
                id="{}-info".format(self.model),
                className="info",
                style=CUR_INFO_STYLE
            )
        else:
            info = None
            
        if isinstance(self.model, str):
            curr_index = self.model
        else:
            curr_index = str(self.model)

        tag_template_tile = "{}-tile-layer"
        tag_template_map = "{}-map"

        if self.model in MODELS:
            curr_tag = 'model'
        else:
            curr_tag = isinstance(tag, str) and tag or str(tag)

        if not isinstance(layer, list):
            layer = [layer]

        if DEBUG: print("********* TAG", tag, "CURR_TAG", curr_tag)
        fig = dl.Map(children=[
            dl.TileLayer(
                id=dict(
                    tag=tag_template_tile.format(curr_tag),
                    index=curr_index
                ),
                url=STYLES[selected_tiles]['url'],
                attribution=STYLES[selected_tiles]['attribution']
            ),
            dl.FullscreenControl(
                position='topright',
            ),
            geojson_contours,
            colorbar,
            info
            ] + layer,
            zoomSnap=0.1,
            zoom=zoom,
            wheelPxPerZoomLevel=120,
            wheelDebounceTime=80,
            center=center,
            id=dict(
                tag=tag_template_map.format(curr_tag),
                index=curr_index
                ),
            inertia=True,
            preferCanvas=True,
            animate=False,
            minZoom=2,
            #className="graph-with-slider",
        )
        # if DEBUG: print("---", fig)
        if DEBUG: print("*** FIGURE EXECUTION TIME: {} ***".format(str(time.time() - self.st_time)))
        return fig


class ScoresFigureHandler(object):
    """ Class to manage the figure creation """

    def __init__(self, network, statistic, selection=None):

        if network == 'aeronet':
            self.sites = pd.read_csv(os.path.join('./conf/',
                OBS[network]['sites']))
#             self.sites = pd.read_table(os.path.join('./conf/',
#                 OBS[network]['sites']), delimiter=r"\s+", engine='python')
            self.size = 15
        else:
            self.sites = None
            self.size = 7

        network_name = OBS[network]['name']
        filedir = OBS[network]['path']
        filename = "{}_{}.h5".format(selection, statistic)
        tab_name = "{}_{}".format(statistic, selection)
        filepath = os.path.join(filedir, "h5", filename)
        if DEBUG: print('SCORES filepath', filepath, 'SELECTION', selection, 'TAB', tab_name)
        self.dframe = pd.read_hdf(filepath, tab_name)  # .replace('_', ' ', regex=True)

        months = ' - '.join([datetime.strptime(sel, '%Y%m').strftime("%B %Y") for sel in selection.split('-')])

        self.title = """{model} - {network_name} {score}<br>{selection}""".format(
                score=STATS_CONF[statistic]['name'], network_name=network_name, model='{model}', selection=months)
        self.xlon = np.array([-27, 60])
        self.ylat = np.array([0, 67])
        self.stat = statistic

    def get_mapbox(self, style='carto-positron', relayout=False, zoom=2.8, center=None):
        """ Returns mapbox layout """
        if center is None and hasattr(self, 'ylat'):
            center = go.layout.mapbox.Center(
                lat=(self.ylat.max()-self.ylat.min())/2 +
                self.ylat.min(),
                lon=(self.xlon.max()-self.xlon.min())/2 +
                self.xlon.min(),
            )
        elif center is not None:
            center = go.layout.mapbox.Center(center)
        else:
            center = go.layout.mapbox.Center({'lat': 30, 'lon': 15})
        mapbox_dict = dict(
            uirevision=True,
            style=style,
            bearing=0,
            center=center,
            pitch=0,
            zoom=zoom
        )

        if not relayout:
            return mapbox_dict

        return dict(
            args=["mapbox", mapbox_dict],
            label=STYLES[style].capitalize(),
            method="relayout"
        )

    def generate_trace(self, xlon, ylat, stats, vals):
        """ Generate trace to be added to data, per variable and timestep """
        if stats is None:
            hovertemplate="lon: %{lon:.2f}<br>lat: %{lat:.2f}<br>value: %{text}"
        else:
            hovertemplate="lon: %{lon:.2f}<br>lat: %{lat:.2f}<br>value: %{text}<br>station: %{customdata}"
        name = '{} score'.format(STATS_CONF[self.stat]['name'])
        return dict(
            type='scattermapbox',
            lon=xlon,
            lat=ylat,
            text=vals,
            customdata=stats,
            name=name,
            hovertemplate=hovertemplate,
            opacity=0.8,
            mode='markers',
            showlegend=False,
            marker=dict(
                showscale=True,
                # colorscale=STATS_CONF[self.stat]['cmap'],
                colorscale=get_colorscale(np.array(STATS_CONF[self.stat]['bounds']), STATS_CONF[self.stat]['cmap']),
                cmax=STATS_CONF[self.stat]['max'],
                cmin=STATS_CONF[self.stat]['min'],
                cmid=STATS_CONF[self.stat]['mid'],
                color=vals,
                size=self.size,
                colorbar=dict(
                    x=0.94,
                    y=0.45,
                    len=0.9,
                    tickmode='array',
                    tickvals=np.array(STATS_CONF[self.stat]['bounds']),
                    thickness=20,
                )
            ),
        )

    def retrieve_scores(self, model, aspect=(1,1), center=None):
        """ run plot """
        
        data_available=''
        if self.sites is not None:
            stations = [st for st in self.sites['SITE']]
            for site in stations:
                self.dframe.loc[self.dframe['station'] == site, 'lon'] = \
                    str(self.sites.loc[self.sites['SITE'] == site, 'LONGITUDE'].values[0].round(2))
                self.dframe.loc[self.dframe['station'] == site, 'lat'] = \
                    str(self.sites.loc[self.sites['SITE'] == site, 'LATITUDE'].values[0].round(2))
        if DEBUG and self.sites is not None: print('BEFORE', model, '\n', self.dframe[['lon', 'lat', 'station', model]].to_markdown())
        self.dframe = self.dframe.replace('-', np.nan)
        if DEBUG and self.sites is not None: print('AFTER', model, '\n', self.dframe[['lon', 'lat', 'station', model]].to_markdown())
        # self.dframe.dropna(inplace=True)
        if DEBUG: print('DATAFRAME', self.dframe, model)
        self.fig = go.Figure(go.Scattermapbox())
        if self.sites is not None and model in self.dframe.columns:
            if DEBUG: print('DATA FOR MODEL', model, '\n', self.dframe[['lon', 'lat', 'station', model]].to_markdown())
            if DEBUG: print('DATA FOR MODEL', model, '\n', self.dframe[['lon', 'lat', 'station', model]].dropna().to_markdown())
            xlon, ylat, stats, vals = self.dframe[['lon', 'lat', 'station', model]].dropna().values.T
            if DEBUG: print('Adding SCORES points ...', xlon, ylat, vals)
            self.fig.add_trace(self.generate_trace(xlon.astype(float), ylat.astype(float), stats, vals.astype(float)))
            if not vals.any(): data_available='<br>NO DATA AVAILABLE'

        elif model in self.dframe.columns:
            xlon, ylat, vals = self.dframe[['lon', 'lat', model]].dropna().values.T
            stats = None
            if DEBUG: print('Adding SCORES points ...', xlon, ylat, vals)
            self.fig.add_trace(self.generate_trace(xlon.astype(float), ylat.astype(float), stats, vals.astype(float)))
            if not vals.any(): data_available='<br>NO DATA AVAILABLE'

        else:
            data_available='<br>NO DATA AVAILABLE'

        if DEBUG: print('Update layout ...', self.title.format(model=MODELS[model]['name']))
        fig_title=dict(text='{} {}'.format(self.title.format(model=MODELS[model]['name']), data_available),
                       xanchor='left',
                       yanchor='top',
                       x=0.01, y=0.95)

        self.fig.update_layout(
            title=fig_title,
            uirevision=True,
            autosize=True,
            hovermode="closest",        # highlight closest point on hover
            mapbox=self.get_mapbox(zoom=2.8-(0.5*aspect[0]), center=center),
            font_size=12-(0.5*aspect[0]),
            # width="100%",
            legend=dict(
                x=0.01,
                y=0.9,
                bgcolor="rgba(0,0,0,0)"
            ),
            margin={"r": 0, "t": 0, "l": 0, "b": 0},
        )

        # if DEBUG: print('Returning fig of size {}'.format(sys.getsizeof(self.fig)))
        return self.fig


class VisFigureHandler(object):
    """ Class to manage the figure creation """

    def __init__(self, selected_date=None):

        self.path_tpl = '/data/daily_dashboard/obs/visibility/{year}/{month}/{year}{month}{day}{tstep0:02d}{tstep1:02d}_visibility.csv'
        self.title = ["Visibility reduced by airborne dust", html.Br(), "{date} {tstep0:02d}-{tstep1:02d} UTC"]
        self.xlon = np.array([-25, 60])
        self.ylat = np.array([0, 65])
        self.ec = 'none'
        self.size = 80
        self.freq = 6
        self.colors = ('#714921', '#da7230', '#fcd775', 'CadetBlue')
        self.labels = ("<1 km", "1 - 2 km", "2 - 5 km", "Haze")
        self.values = (0, 1, 2, 3)
        self.markers = ('o', 'o', 'o', '^')

        if selected_date:
            self.selected_date_plain = selected_date

            self.selected_date = datetime.strptime(
                selected_date, "%Y%m%d").strftime("%Y-%m-%d")
        else:
            self.selected_date_plain = None
            self.selected_date = None

    def set_data(self, tstep=0):
        """ Set time dependent data """
        tstep0 = tstep
        tstep1 = tstep + self.freq

        year = datetime.strptime(self.selected_date_plain, '%Y%m%d').strftime('%Y')
        month = datetime.strptime(self.selected_date_plain, '%Y%m%d').strftime('%m')
        day = datetime.strptime(self.selected_date_plain, '%Y%m%d').strftime('%d')

        filepath = self.path_tpl.format(year=year, month=month, day=day, tstep0=tstep0, tstep1=tstep1)
        if DEBUG: print("VIS FILEPATH", filepath)
        if not os.path.exists(filepath):
            return [], [], [], [], [], ()

        data = pd.read_table(filepath)

        # uncertain
        cx = np.where((data['WW'].astype(str) == "HZ") | (data['WW'].astype(str) == "5")| (data['WW'].astype(str) == "05"))

        # vis <= 1km
        c0t = np.where((data['VV'] <= 1000))[0]
        c0x = np.where([c0t == i for i in cx[0]])[-1]
        c0 = (np.delete(c0t, c0x),)

        # vis 1km <= 2km
        c1t = np.where((data['VV'] > 1000) & (data['VV'] <= 2000))[0]
        c1x = np.where([c1t == i for i in cx[0]])[-1]
        c1 = (np.delete(c1t, c1x),)

        # vis 2km <= 5km
        c2t = np.where((data['VV'] > 2000) & (data['VV'] <= 5000))[0]
        c2x = np.where([c2t == i for i in cx[0]])[-1]
        c2 = (np.delete(c2t, c2x),)

        xlon = data['LON'].values
        ylat = data['LAT'].values
        stats = data['STATION'].values

        visibility = 'VV' in data and data['VV']
        humidity = 'HUMIDITY' in data and data['HUMIDITY']

        if DEBUG: print("VIS DATA", xlon, ylat, stats, (c0, c1, c2, cx))
        return xlon, ylat, stats, visibility, humidity, (c0, c1, c2, cx)

    def generate_var_tstep_trace(self, xlon, ylat, stats, visibility, humidity, value, color, label, marker, tstep=0):
        """ Generate trace to be added to data, per variable and timestep """
        name = 'visibility {}'.format(label)

        legend = []
        res = np.zeros((len(xlon)))
        for i, (val, lab) in enumerate(zip(value, label)):
            # assign a value from 0 to 3 to the different
            res[val] = i
            leg = html.Div([
                    html.Span(
                        "",
                        className="vis-legend-point",
                        style={
                            'backgroundColor': self.colors[i],
                            }
                        ),
                    html.Span(
                        lab,
                        className="vis-legend-label",
                        )
                    ],
                    style={ 'display': 'block' })
            legend.append(leg)

        legend = html.Div(
                legend,
                className="vis-legend"
            )

        if DEBUG:
            print('VIS ___', xlon, ylat, '\n*********', visibility, '\n*********', humidity)
        df = pd.DataFrame({
            'lon': np.array(xlon).round(2),
            'lat': np.array(ylat).round(2),
            'stat': stats,
            'visibility': visibility,
            'humidity': humidity,
            'value': res
            })
        dicts = df.to_dict('rows')
        for item in dicts:
            item["tooltip"] = \
                    """<span class='popup-map-body'>
                    <b>{station}</b>
                    <br>LAT: <b>{lat}</b> LON: <b>{lon}</b>
                    <br>VISIBILITY: <b>{visibility:.1f} km</b>
                    {humidity}
                    </span>""".format(
                            lat=item['lat'],
                            lon=item['lon'],
                            station=item['stat'],
                            visibility=float(item['visibility']/1e3),
                            humidity=(item['humidity'] is not False and "<br>RELATIVE HUMIDITY: <b>{} %</b>".format(round(item['humidity'])) or '')
                            )

        geojson = dlx.dicts_to_geojson(dicts, lon="lon")
        ns = Namespace("observationsTab", "observationsMaps")
        point_to_layer = ns("pointToLayer")
        geojson = dl.GeoJSON(data=geojson,
                options=dict(
                    pointToLayer=point_to_layer,
                ),
                hideout=dict(
                    colorProp='value',
                    colorscale=self.colors,
                    circleOptions=dict(
                        fillOpacity=OPACITY,
                        stroke=False,
                        #fillColor='#f0b450',
                        radius=8),
                )
            )
        
        if list(visibility):
            cur_title = self.get_title(tstep)
        else:
            cur_title = self.get_title()

        info = html.Div(
            children=cur_title,
            id="vis-info",
            className="info",
            style=INFO_STYLE
        )
        return df, [geojson, info, legend]

    def get_title(self, tstep=None):
        """ return title according to the date """
        if tstep is None:
            return ["NO DATA AVAILABLE"]
        tstep0 = tstep
        tstep1 = tstep + self.freq
        fdate = datetime.strptime(self.selected_date_plain, '%Y%m%d').strftime('%d %B %Y')
        return [tit.format(date=fdate, tstep0=tstep0, tstep1=tstep1)
                if isinstance(tit, str) else tit for tit in self.title]

    def retrieve_var_tstep(self, tstep=0, hour=None, static=True, aspect=(1,1), center=None):
        """ run plot """

        tstep = int(tstep)

        xlon, ylat, stats, visibility, humidity, vals = self.set_data(tstep)
        if tstep is not None:
            return self.generate_var_tstep_trace(xlon, ylat, stats, visibility, humidity, vals, self.colors, self.labels, self.markers, tstep)
        if DEBUG: print('Adding one point ...')
        return None


class ProbFigureHandler(object):
    """ Class to manage the figure creation """

    def __init__(self, var=None, prob=None, selected_date=None):
        """ Initialization with variable, prob and date """

        if var is None:
            var = DEFAULT_VAR

        self.varname = var

        probs = PROB[var]['prob_thresh']
        if prob is None:
            prob = probs[0]

        self.bounds = np.arange(0, 110, 10)

        self.prob = prob

        geojson_path = PROB[var]['geojson_path']
        geojson_file = PROB[var]['geojson_template']
        netcdf_path = PROB[var]['netcdf_path']
        netcdf_file = PROB[var]['netcdf_template']

        self.geojson = os.path.join(geojson_path, geojson_file).format(prob=prob, date=selected_date, var=var)
        self.filepath = os.path.join(netcdf_path, netcdf_file).format(prob=prob, date=selected_date, var=var)

        if os.path.exists(self.filepath):
            self.input_file = nc_file(self.filepath)

            if 'lon' in self.input_file.variables:
                lon = self.input_file.variables['lon'][:]
                lat = self.input_file.variables['lat'][:]
            else:
                lon = self.input_file.variables['longitude'][:]
                lat = self.input_file.variables['latitude'][:]
            time_obj = self.input_file.variables['time']
            self.tim = time_obj[:]
            tim_units = time_obj.units.split()
            if len(tim_units) == 3:
                self.what, _, rdate = tim_units
                rtime = "00:00"
            elif len(tim_units) >= 4:
                self.what, _, rdate, rtime = tim_units[:4]
            if len(rtime) > 5:
                rtime = rtime[:5]
            self.rdatetime = datetime.strptime("{} {}".format(rdate, rtime),
                                               "%Y-%m-%d %H:%M")
            varlist = [var for var in self.input_file.variables if var in VARS]
            self.xlon, self.ylat = np.meshgrid(lon, lat)

        if selected_date:
            self.selected_date_plain = selected_date

            self.selected_date = datetime.strptime(
                selected_date, "%Y%m%d").strftime("%Y-%m-%d")

        self.fig = None

    def set_data(self, varname, tstep=0):
        """ Set time dependent data """
        mul = 1  # VARS[varname]['mul']
        var = self.input_file.variables[varname][tstep]*mul
        idx = np.where(var.ravel() >= VARS[varname]['bounds'][0])  # !=-9.e+33)
        xlon = self.xlon.ravel()[idx]
        ylat = self.ylat.ravel()[idx]
        var = var.ravel()[idx]

        return xlon, ylat, var

    def retrieve_cdatetime(self, tstep=0):
        """ Retrieve data from NetCDF file """
        if DEBUG: print("----------", tstep, self.tim)
        tim = int(self.tim[tstep])
        if self.what == 'days':
            cdatetime = self.rdatetime + relativedelta(days=tim)
        elif self.what == 'hours':
            cdatetime = self.rdatetime + relativedelta(hours=tim)
        elif self.what == 'minutes':
            cdatetime = self.rdatetime + relativedelta(minutes=tim)
        elif self.what == 'seconds':
            cdatetime = self.rdatetime + relativedelta(seconds=tim)

        return cdatetime

    def generate_contour_tstep_trace(self, varname, tstep=0):
        """ Generate trace to be added to data, per variable and timestep """
        from dash_server import app
        if varname is None:
            varname = self.varname

        if DEBUG: print("##############", tstep)

        geojson_file = self.geojson.format(step=tstep)
        geojson_url = app.get_asset_url(geojson_file.replace('/data/daily_dashboard', 'geojsons'))

        name = VARS[varname]['name']
        bounds = self.bounds
        colorscale = COLORS_PROB

        if DEBUG: print("GEOJSON_URL", geojson_url)

        style = dict(weight=0, opacity=0, color='white', dashArray='', fillOpacity=OPACITY)

        # Create colorbar.
        ctg = ["{:.1f}".format(cls) if '.' in str(cls) else "{:d}".format(cls)
                for i, cls in enumerate(bounds)]
        indices = list(range(len(ctg)))
        colorbar = dl.Colorbar(
                min=-.1, max=len(ctg)-.7,
                classes=indices,
                colorscale=colorscale,
                tickValues=indices,
                tickText=ctg,
                position='topleft',
                width = 330,
                height=8,
                style={ 'top': '65px','overflow':'hidden', 'white-space':'nowrap'}
                )

        # Geojson rendering logic, must be JavaScript as it is executed in clientside.
        ns = Namespace("forecastTab", "forecastMaps")
        style_handle = ns("styleHandle")

        geojson = dl.GeoJSON(
                url=geojson_url,
                options=dict(style=style_handle),
                hideout=dict(colorscale=colorscale, bounds=bounds, style=style, colorProp="value")
                )  # url to geojson file

        return geojson, colorbar

    def get_title(self, varname, tstep=0):
        """ return title according to the date """
        # tstep += 1
        rdatetime = self.rdatetime
        cdatetime = self.retrieve_cdatetime(tstep)

        mod_avail = 0
        for model in MODELS:
            if model == 'median':
                continue
            path_template = '{}{}.nc'.format(rdatetime.strftime("%Y%m%d"), MODELS[model]['template'])
            fpath = os.path.join(MODELS[model]['path'], 'netcdf', path_template)
            if os.path.exists(fpath):
                mod_avail += 1

        return PROB[varname]['title'].format(
            prob=self.prob,
            members=mod_avail,
            rday=rdatetime.strftime("%d"),
            rmonth=rdatetime.strftime("%b"),
            ryear=rdatetime.strftime("%Y"),
            sday=cdatetime.strftime("%d"),
            smonth=cdatetime.strftime("%b"),
            syear=cdatetime.strftime("%Y"),
        )

    def retrieve_var_tstep(self, varname=None, day=0, static=True, aspect=(1,1)):
        """ run plot """
        tstep = int(day)
        if varname is None:
            varname = self.varname
        if DEBUG: print("***", varname, day, static, self.geojson, self.filepath)

        geojson = colorbar = None

        if varname and os.path.exists(self.geojson.format(step=day)):
            if DEBUG: print('Adding contours ...')
            geojson, colorbar = self.generate_contour_tstep_trace(varname, tstep)
        if varname and os.path.exists(self.filepath) and static:
            if DEBUG: print('Adding points ...')
        elif varname is None or not os.path.exists(self.filepath):
            if DEBUG: print('Adding one point ...')

        if DEBUG: print('Update layout ...')
        if not varname:
            if DEBUG: print('ONE')
            fig_title = ''
        elif varname and not os.path.exists(self.filepath):
            if DEBUG: print('TWO')
            fig_title = html.P(html.B("DATA NOT AVAILABLE"))
        else:
            if DEBUG: print('THREE')
            fig_title = html.P(html.B(
                [
                    item for sublist in self.get_title(varname, tstep).split('<br>') for item in [sublist, html.Br()]
                ][:-1]
            ))

        info = html.Div(
            children=fig_title,
            id="prob-info",
            className="info",
            style=INFO_STYLE
        )
        return geojson, colorbar, info


class WasFigureHandler(object):
    """ Class to manage the figure creation """

    def __init__(self, was='burkinafaso', model='median', variable='SCONC_DUST', selected_date=None):
        """ Initialize WasFigureHandler with shapefile and netCDF data """
        self.model = model
        self.was = was
        self.variable = variable

        if self.model and selected_date:
            # read nc file
            if DEBUG: print("MODEL", model)
            filepath = NETCDF_TEMPLATE.format(
                MODELS[self.model]['path'],
                selected_date,
                MODELS[self.model]['template']
            )

            if os.path.exists(filepath):
                self.input_file = nc_file(filepath)
                time_obj = self.input_file.variables['time']
                self.tim = time_obj[:]
                tim_units = time_obj.units.split()
                if len(tim_units) == 3:
                    self.what, _, rdate = tim_units
                    rtime = "00:00"
                elif len(tim_units) >= 4:
                    self.what, _, rdate, rtime = tim_units[:4]
                if len(rtime) > 5:
                    rtime = rtime[:5]
                self.rdatetime = datetime.strptime("{} {}".format(rdate, rtime),
                                                   "%Y-%m-%d %H:%M")
        else:
            self.input_file = None

        if selected_date:
            self.selected_date_plain = selected_date

            self.selected_date = datetime.strptime(
                selected_date, "%Y%m%d").strftime("%Y-%m-%d")

        self.fig = None

    def get_regions_data(self, day=0):
        input_dir = WAS[self.was]['path'].format(was=self.was, format='h5', date=self.selected_date_plain)
        input_file = WAS[self.was]['template'].format(date=self.selected_date_plain, var=self.variable, format='h5')

        input_path = os.path.join(input_dir, input_file)

        if DEBUG: print("INPUT PATH", input_path, self.selected_date_plain)
        if not os.path.exists(input_path):
            return [], [], []

        df = pd.read_hdf(input_path, 'was_{}'.format(self.selected_date_plain)).set_index('day')

        names, colors, definitions = df.loc['Day{}'.format(day)].values.T
        return names, colors, definitions

    def get_geojson_url(self, day=0):
        from dash_server import app
        geojsons_dir = WAS[self.was]['path'].format(was=self.was, format='geojson', date=self.selected_date_plain)
        geojson_file = WAS[self.was]['template'].format(date=self.selected_date_plain, var=self.variable, format='geojson').replace('.geojson', '_{}.geojson'.format(day))
        geojson_filepath = os.path.join(geojsons_dir, geojson_file)
        if not os.path.exists(geojson_filepath):
            if DEBUG: print("ERROR: WAS GEOJSON URL", geojson_filepath)
            return None

        pathlist = os.path.normpath(geojsons_dir).split(os.sep)
        geojsons_dir_cleaned = os.sep.join(pathlist[pathlist.index('was'):])

        geojson_path = os.path.join(geojsons_dir_cleaned, geojson_file)
        geojson_url = app.get_asset_url(os.path.join('geojsons', geojson_path))

        if DEBUG: print("WAS GEOJSON URL", geojson_url)

        return geojson_url

    def retrieve_cdatetime(self, tstep=0):
        tim = int(self.tim[tstep])
        """ Retrieve data from NetCDF file """
        if self.what == 'days':
            cdatetime = self.rdatetime + relativedelta(days=tim)
        elif self.what == 'hours':
            cdatetime = self.rdatetime + relativedelta(hours=tim)
        elif self.what == 'minutes':
            cdatetime = self.rdatetime + relativedelta(minutes=tim)
        elif self.what == 'seconds':
            cdatetime = self.rdatetime + relativedelta(seconds=tim)

        return cdatetime

    def generate_contour_tstep_trace(self, day=0):
        """ Generate trace to be added to data, per variable and timestep """

        if not hasattr(self, 'was'):
            return None, None

        names, colors, definitions = self.get_regions_data(day=day)
        colors = np.array(colors)
        if DEBUG: print("::::::::::", names, colors, definitions)
        style = dict(weight=1, opacity=1, color='white', dashArray='3', fillOpacity=OPACITY)

        # Create colorbar.
        colormap = WAS[self.was]['colors']

        legend = []
        # res = np.zeros(len(colormap))
        # if DEBUG: print("^^^^^", colormap)
        for color_idx, (color, definition) in enumerate(colormap.items()):
            # assign a value from 0 to 3 to the different
            leg = html.Div([
                    html.Span(
                        "",
                        className="was-legend-point",
                        style={
                            'backgroundColor': color,
                            }
                        ),
                    html.Span(
                        definition, 
                        className="was-legend-label",
                        )
                    ],
                    style={ 'display': 'block' })
            legend.append(leg)

        legend = html.Div(
                legend,
                className="was-legend"
            )

        # Geojson rendering logic, must be JavaScript as it is executed in clientside.
        ns = Namespace("forecastTab", "wasMaps")
        style_handle = ns("styleHandle")

        geojson = dl.GeoJSON(
                url=self.get_geojson_url(day=day),
                # data=geojson_data,
                # zoomToBounds=True,
                hoverStyle=arrow_function(dict(weight=2, color='white', dashArray='', fillOpacity=OPACITY)),
                options=dict(style=style_handle),
                hideout=dict(colorscale=list(colormap.keys()),
                    bounds=[i for i in range(len(colormap.keys()))],
                    style=style,
                    colorProp="value")
                )  # url to geojson file

        return geojson, legend

    def get_title(self, day=0):
        """ return title according to the date """
        rdatetime = self.retrieve_cdatetime(tstep=0)
        cdatetime = rdatetime + relativedelta(days=day)
        return r'{}'.format(WAS[self.was]['title'] % {
            'rhour':  rdatetime.strftime("%H"),
            'rday':   rdatetime.strftime("%d"),
            'rmonth': rdatetime.strftime("%b"),
            'ryear':  rdatetime.strftime("%Y"),
            'shour':  cdatetime.strftime("%H"),
            'sday':   cdatetime.strftime("%d"),
            'smonth': cdatetime.strftime("%b"),
            'syear':  cdatetime.strftime("%Y"),
            #'step':   "{:02d}".format(tstep*FREQ),
        })

    def retrieve_var_tstep(self, day=0, static=True, aspect=(1,1)):
        """ run plot """
        self.fig = go.Figure()
        day = int(day)
        if DEBUG: print('Adding contours ...')
        geojson, legend = self.generate_contour_tstep_trace(day)
        if DEBUG: print('Update layout ...')
        if self.input_file is not None:
            fig_title = html.P(html.B(
                [
                    item
                    for sublist in self.get_title(day).split('<br>')
                    for item in [sublist, html.Br()]
                ][:-1]
            ))
        else:
            fig_title = html.P(html.B("DATA NOT AVAILABLE"))

        info = html.Div(
            children=fig_title,
            id="was-info",
            className="info",
            style=INFO_STYLE
        )
        return geojson, legend, info
