# -*- coding: utf-8 -*-
""" TAB FORECAST """

from datetime import datetime as dt
from datetime import timedelta
import time
import math
from random import random
import orjson
import logging

import dash
import dash_bootstrap_components as dbc
from dash import dcc
from dash import html
from dash.dependencies import Output
from dash.dependencies import Input
from dash.dependencies import State
from dash.dependencies import ALL
from dash.exceptions import PreventUpdate
import dash_leaflet as dl

from ines_core_map_handler import MAP_LAYERS
from map_handler import DEFAULT_VAR
from map_handler import DEFAULT_MODEL
from map_handler import VARS 
from map_handler import WAS
from map_handler import MODELS
from map_handler import FREQ
from map_handler import DEBUG
from map_handler import END_DATE
from map_handler import PROB 
from map_handler import GRAPH_HEIGHT 
from map_handler import MODEBAR_CONFIG_TS
from map_handler import MODEBAR_LAYOUT_TS
from map_handler import cache, cache_timeout
from tabs.forecast import tab_forecast
from utils import calc_matrix


@dash.callback(
    [Output({'type':'collapse', 'index': 1}, 'is_open'),
     Output({'type':'collapse', 'index': 2}, 'is_open'),
     Output({'type':'collapse', 'index': 3}, 'is_open'),
     Output('group-2-toggle', 'disabled'),
     Output('group-3-toggle', 'disabled'),
     Output('variable-dropdown-forecast','value')],
    [Input('group-1-toggle', 'n_clicks'),
     Input('group-2-toggle', 'n_clicks'),
     Input('group-3-toggle', 'n_clicks'),
     Input('variable-dropdown-forecast', 'value')],
    [State({'type':'collapse', 'index': 1}, 'is_open'),
     State({'type':'collapse', 'index': 2}, 'is_open'),
     State({'type':'collapse', 'index': 3}, 'is_open'),],
    prevent_initial_call=True
)
def render_forecast_tab(modbutton, probbutton, wasbutton, var, modopen,
                        probopen, wasopen):
    """ Function rendering requested tab """
    ctx = dash.callback_context

    if ctx.triggered:
        button_id = ctx.triggered[0]["prop_id"].split(".")[0]

        if button_id == "group-1-toggle" and modbutton:
            if modopen is True:
                return not modopen, False, False, dash.no_update, dash.no_update, dash.no_update
            return not modopen, False, False, dash.no_update, dash.no_update, dash.no_update
        elif button_id == "group-2-toggle" and probbutton:
            if probopen is True:
                return False, not probopen, False, dash.no_update, dash.no_update, dash.no_update
            return False, not probopen, False, dash.no_update, dash.no_update, dash.no_update
        elif button_id == "group-3-toggle" and wasbutton:
            if wasopen is True:
                return False, False, not wasopen, dash.no_update, dash.no_update, 'SCONC_DUST'
            return False, False, not wasopen, dash.no_update, dash.no_update, 'SCONC_DUST'

    if var == 'SCONC_DUST':
        # raise PreventUpdate
        return modopen, probopen, wasopen, False, False, dash.no_update
    elif var == 'OD550_DUST':
        # only models and prob can be opened
        if wasopen:
            return True, False, False, False, True, dash.no_update
        return modopen, probopen, wasopen, False, True, dash.no_update

    if modopen:
        return True, False, False, True, True, dash.no_update
    return True, False, False, True, True, dash.no_update


@dash.callback(
    [Output('prob-dropdown', 'options'),
     Output('prob-dropdown', 'value')],
    [Input('variable-dropdown-forecast', 'value')],
    prevent_initial_call=False
)
def update_prob_dropdown(var):
    """ Update Prob maps dropdown """
    if var in ['OD550_DUST','SCONC_DUST']:
        opt_list = PROB[var]['prob_thresh']
        units = PROB[var]['units']
        return [
                {'label': '> {} {}'.format(prob, units),
                 'value': 'prob_{}'.format(prob)}
                for prob in opt_list], 'prob_{}'.format(opt_list[0])

    raise PreventUpdate

@dash.callback(
    [  # Output('alert-models-auto', 'is_open'),
     Output('model-dropdown', 'options'),
     Output('model-dropdown', 'value'),
     Output('btn-anim-download', 'style')],
    [Input('variable-dropdown-forecast', 'value')],
    [Input('model-dropdown', 'value'),],
    prevent_initial_call=True
    )
def update_models_dropdown(variable, checked):
    """ Update Models maps dropdown """

    btn_style = { 'display' : 'block' }
    models = VARS[variable]['models']
    if models == 'all':
        models = list(MODELS.keys())
    else:
        models = eval(models)

    options = [{
        'label': MODELS[model]['name'],
        'value': model,
        'disabled': model not in models,
        } for model in MODELS]

    checked = [c for c in models if c in checked or len(models)==1]
    if len(checked) > 1:
        btn_style = { 'display' : 'none' }
    logging.debug('MODELS %s OPTS %s %s', models, type(options), options)

    return options, checked, btn_style

#MOVED TO GENERIC CALLBACKS
# def rotate_models_caret(collapse_open):
# def rotate_prob_caret(collapse_open):
# def rotate_was_caret(collapse_open):

@dash.callback(
    [Output('info-collapse', 'is_open'),
     Output('download-collapse', 'is_open')],
    [Input('info-button', 'n_clicks'),
     Input('download-button', 'n_clicks')],
    [State('info-collapse', 'is_open'),
     State('download-collapse', 'is_open')],
    prevent_initial_call=True
)
def sidebar_bottom(n_info, n_download, open_info, open_download):
    """ Returns bottom sidebar with info and download menus """
    ctx = dash.callback_context
    if ctx.triggered:
        button_id = ctx.triggered[0]["prop_id"].split(".")[0]

        if button_id == 'info-button':
            logging.debug('clicked INFO %s, %s', not open_info, False)
            return not open_info, False
        elif button_id == 'download-button':
            logging.debug('clicked DOWN %s %s', False, not open_download)
            return False, not open_download

    logging.debug('clicked NONE %s %s', False, False)
    raise PreventUpdate

@dash.callback(
    Output('btn-anim-download', 'href'),
    [Input('model-dropdown', 'value'),
     Input('variable-dropdown-forecast', 'value'),
     Input('model-date-picker', 'date'),
     Input('model-slider-graph', 'value')],
     prevent_initial_call=False,
)
def download_anim_link(models, variable, date, tstep):
    """ Download PNG frame """
    if variable is None or date is None or tstep is None:
        raise PreventUpdate

    # from callback_tools import get_models_figure
    from callback_tools import download_image_link

    logging.debug('GIF %s %s %s', models, variable, date)
    try:
        curdate = dt.strptime(date, '%Y-%m-%d').strftime('%Y%m%d')
    except:
        curdate = date
    anim = download_image_link(models, variable, curdate, anim=True)
    logging.debug('DOWNLOAD LINK %s', anim)
    return anim

@dash.callback(
    Output('was-slider-container', 'children'),
    [Input('was-date-picker', 'date')],
)
def update_was_timeslider(date):
    """ Update time slider number of days """
    from tabs.forecast import gen_time_slider
    try:
        date = dt.strptime(date, "%Y-%m-%d").strftime("%Y%m%d")
    except:
        pass
    _, time_slider = gen_time_slider('was', date)
    return time_slider

@dash.callback(
    [Output('was-graph', 'children'),
     Output('was-previous', 'data')],
    [Input('was-apply', 'n_clicks'),
     Input('was-date-picker', 'date'),
     Input('was-slider-graph', 'value')],
    [State('was-dropdown', 'value'),
     State('variable-dropdown-forecast', 'value'),
     State('was-previous', 'data'),
     State({'tag': 'view-style', 'index': ALL}, 'active'),
     State({'tag': 'was-map', 'index': ALL}, 'zoom'),
     State({'tag': 'was-map', 'index': ALL}, 'center')],
    prevent_initial_call=False
    )
def update_was_figure(n_clicks, date, day, was, var, previous, selected_tiles, zoom, center):
    """ Update Warning Advisory Systems maps """
    logging.debug('WAS %s %s %s %s %s %s %s %s', n_clicks, date, day, was, previous, selected_tiles, zoom, center)
    ctx = dash.callback_context
    if ctx.triggered:
        button_id = ctx.triggered[0]["prop_id"].split(".")[0]
        if button_id != 'was-apply' and date is None and day is None:
            raise PreventUpdate
    from callback_tools import get_was_figure
    from callback_tools import get_figure
    if not zoom or previous != was:
        zoom = WAS[was]['zoom']
    else:
        zoom = zoom[0]
    if not center or previous != was:
        center = WAS[was]['center']
    else:
        center = center[0]

    previous = was

    logging.debug('WAS %s', was)
    if date is not None:
        date = date.split(' ')[0]
        try:
            date = dt.strptime(
                date, "%Y-%m-%d").strftime("%Y%m%d")
        except:
            pass
        logging.debug('SERVER: callback date %s', date)
    else:
        date = END_DATE

    logging.debug("WAS figure %s %s %s", date, was, day)
    if was:
        selected_tiles = list(MAP_LAYERS.keys())[selected_tiles.index(True)]
        layers = get_was_figure(was, day, selected_date=date)
        fig = get_figure(selected_tiles=selected_tiles, zoom=zoom, center=center, tag='was', layers=layers)
        return fig, previous

    raise PreventUpdate 

@dash.callback(
    Output('prob-slider-container', 'children'),
    [Input('prob-date-picker', 'date')],
)
def update_prob_timeslider(date):
    """ Update time slider number of days """
    from tabs.forecast import gen_time_slider
    try:
        date = dt.strptime(date, "%Y-%m-%d").strftime("%Y%m%d")
    except:
        pass
    _, time_slider = gen_time_slider('prob', date)
    return time_slider

@dash.callback(
    Output('prob-graph', 'children'),
    [Input('prob-apply', 'n_clicks'),
     Input('prob-date-picker', 'date'),
     Input('prob-slider-graph', 'value')],
    [State('prob-dropdown', 'value'),
     State('variable-dropdown-forecast', 'value'),
     State({'tag': 'view-style', 'index': ALL}, 'active'),
     State({'tag': 'prob-map', 'index': ALL}, 'zoom'),
     State({'tag': 'prob-map', 'index': ALL}, 'center')],
    prevent_initial_call=False
)
@cache.memoize(timeout=cache_timeout)
def update_prob_figure(n_clicks, date, day, prob, var, selected_tiles, zoom, center):
    """ Update Probability maps """
    # if not prob in case user navigates to section via URL
    if not prob:
        prob = 'prob_0.1'
    if not zoom:
        zoom = None
    else:
        zoom = zoom[0]
    if not center:
        center = None
    else:
        center = center[0]
    logging.debug('PROB %s %s %s %s %s %s %s %s', date, day, var, prob, var, selected_tiles, zoom, center)
    ctx = dash.callback_context
    if ctx.triggered:
        button_id = ctx.triggered[0]["prop_id"].split(".")[0]
        if button_id != 'prob-apply' and var is None and day is None:
            raise PreventUpdate

    from callback_tools import get_prob_figure
    from callback_tools import get_figure

    if date is not None:
        date = date.split(' ')[0]
        try:
            date = dt.strptime(
                date, "%Y-%m-%d").strftime("%Y%m%d")
        except:
            pass
        logging.debug('SERVER: callback date %s', date)
    else:
        date = END_DATE

    if prob:
        prob = prob.replace('prob_', '')
        selected_tiles = list(MAP_LAYERS.keys())[selected_tiles.index(True)]
        layers = get_prob_figure(var, prob, day, selected_date=date)
        fig = get_figure(selected_tiles=selected_tiles, zoom=zoom, center=center, tag='prob', layers=layers)

        logging.debug("FIG %s", fig)
        return fig
    raise PreventUpdate

def get_view(index, url, attr):
    """ Helper function to assign selected map with selected view """
    mod_url = mod_attr = []
    prob_url = prob_attr = []
    was_url = was_attr = []
    if index == 1:
        mod_url = url
        mod_attr = attr
    if index == 2:
        prob_url = url 
        prob_attr = attr
    if index == 3:
        was_url = url 
        was_attr = attr
    return mod_url, mod_attr, prob_url, prob_attr, was_url, was_attr

@dash.callback(
    [Output({'tag': 'model-tile-layer', 'index': ALL}, 'url'),
     Output({'tag': 'model-tile-layer', 'index': ALL}, 'attribution'),
     Output({'tag': 'prob-tile-layer', 'index': ALL}, 'url'),
     Output({'tag': 'prob-tile-layer', 'index': ALL}, 'attribution'),
     Output({'tag': 'was-tile-layer', 'index': ALL}, 'url'),
     Output({'tag': 'was-tile-layer', 'index': ALL}, 'attribution'),
     Output({'tag': 'view-style', 'index': ALL}, 'active')],
    [Input({'tag': 'view-style', 'index': ALL}, 'n_clicks')],
    [State({'tag': 'model-tile-layer', 'index': ALL}, 'url'),
     State({'tag': 'prob-tile-layer', 'index': ALL}, 'url'),
     State({'tag': 'was-tile-layer', 'index': ALL}, 'url'),
     State({'tag': 'view-style', 'index': ALL}, 'active')],
    prevent_initial_call=True
)
@cache.memoize(timeout=cache_timeout)
def update_styles_button(*args):
    """ Function updating styles button """
    #get the index of the map which view was triggered
    for i in range(1, 4):
        if len(args[i]) >= 1:
            index = i
    graphs = args[index]
    num_graphs = len(graphs)
    ctx = dash.callback_context
    if ctx.triggered and num_graphs > 0:
        button_id = orjson.loads(ctx.triggered[0]["prop_id"].split(".")[0])
        logging.debug("BUTTON ID %s %s", str(button_id), type(button_id))
        if button_id['index'] in MAP_LAYERS:
            logging.debug("CURRENT ARGS %s", str(args))
            active = args[-1]
            logging.debug("NUM GRAPHS %s", num_graphs)
            url = [MAP_LAYERS[button_id['index']]['url'] for _ in range(num_graphs)]
            attr = [MAP_LAYERS[button_id['index']]['attribution'] for _ in range(num_graphs)]
            res = [False for _ in active]
            st_idx = list(MAP_LAYERS.keys()).index(button_id['index'])
            if active[st_idx] is False:
                res[st_idx] = True
            logging.debug('***** %s %s %s', url, attr, res)
            mod_url, mod_attr, prob_url, prob_attr, was_url, was_attr = get_view(index, url, attr)
            return mod_url, mod_attr, prob_url, prob_attr, was_url, was_attr, res 

    logging.debug('NOTHING TO DO')
    raise PreventUpdate

@dash.callback(
    [Output('model-clicked-coords', 'data'),
     Output('current-popups-stored', 'data'),
     Output(dict(tag='model-map', index=ALL, n_clicks=ALL), 'children')],
    [Input(dict(tag='model-map', index=ALL, n_clicks=ALL), 'click_lat_lng')],
    [State(dict(tag='model-map', index=ALL, n_clicks=ALL), 'id'),
     State(dict(tag='model-map', index=ALL, n_clicks=ALL), 'children'),
     State('model-date-picker', 'date'),
     State('model-slider-graph', 'value'),
     State('variable-dropdown-forecast', 'value'),
     State('model-clicked-coords', 'data'),
     State('current-popups-stored', 'data')],
    prevent_initial_call=False
)
def models_popup(click_data, map_ids, res_list, date, tstep, var, coords, popups):
    """ Manages models popups for timeseries """
    from callback_tools import get_single_point
    logging.debug("CLICK: %s", str(click_data))
    if click_data.count(None) == len(click_data):
        raise PreventUpdate

    logging.debug("MAPID: %s %s", str(map_ids), type(map_ids))
    # logging.debug("RESLIST: %s %s", str(res_list), type(res_list))

    ctxt = dash.callback_context.triggered[0]["prop_id"].split(".")[0]
    logging.debug("CTXT %s %s", ctxt, type(ctxt))
    if not ctxt or ctxt is None:
        raise PreventUpdate

    if popups is None:
        popups = {}

    trigger = orjson.loads(ctxt)
    logging.debug('TRIGGER %s %s', trigger, type(trigger))

    # curr_models = [m['index'] for m in map_ids]
    res = res_list
    if trigger in map_ids:
        model = trigger['index']
        mod_idx = map_ids.index(trigger)
        click = click_data[mod_idx]
        if tstep is None:
            tstep = 0
        else:
            tstep = int(tstep/FREQ)

        logging.debug("MODEL %s CLICK %s DATE %s TSTEP %s MODIDX %s", model, click, date, tstep, mod_idx)

        if click is not None and model is not None:
            lat, lon = click
            try:
                selected_date = dt.strptime(date, '%Y%m%d')
            except:
                selected_date = dt.strptime(date, '%Y-%m-%d')

            if model in MODELS and MODELS[model]['start'] == 12:
                if tstep < 4:
                    date = (selected_date - timedelta(days=1)).strftime("%Y%m%d")
                    tstep = 4 + int(tstep)
                else:
                    date = selected_date.strftime("%Y%m%d")
                    tstep = int(tstep) - 4
                valid_dt = dt.strptime(date, '%Y%m%d') + timedelta(hours=(tstep+4)*FREQ)
            else:
                date = selected_date.strftime("%Y%m%d")
                valid_dt = dt.strptime(date, '%Y%m%d') + timedelta(hours=tstep*FREQ)

            logging.debug("MODEL %s CLICK %s DATE %s TSTEP %s", model, click, date, tstep, mod_idx)
            value = get_single_point(model, date, int(tstep), var, lat, lon)
            logging.debug("VALUE: %s", str(value))

            marker = dl.Popup(
                children=[
                    html.Div([
                        html.Span(html.P(
                            '{:.2f}'.format(value*VARS[var]['mul'])),
                            className='popup-map-value',
                        ),
                        html.Span([
                            html.B("Lat {:.2f}, Lon {:.2f}".format(lat, lon)), html.Br(),
                            "DATE {:02d} {} {} {:02d}UTC".format(
                                valid_dt.day,
                                dt.strftime(valid_dt, '%b'),
                                valid_dt.year,
                                valid_dt.hour
                            ),
                            html.Br(),
                            html.Button("EXPLORE TIMESERIES",
                                id=dict(
                                    tag='ts-button',
                                    index=model,
                                    random=int(random()*100)
                                ),
                                n_clicks=0,
                                className='popup-ts-button'
                            )],
                            className='popup-map-body',
                        )],
                    )
                ],
                position=[lat, lon],
                autoClose=True, 
                closeOnEscapeKey=True,
                closeOnClick=False,
                closeButton=True,
                className='popup-map-point',
            )

            logging.debug("|||| %s %s", res[mod_idx][-1], type(res[mod_idx][-1]))
            if not res[mod_idx] or res[mod_idx] is None:
                logging.debug("*** 1 ***")
                res[mod_idx] = marker
            elif res[mod_idx][-1] is None:
                logging.debug("*** 2 ***")
                res[mod_idx][-1] = marker
            else:
                logging.debug("*** 3 ***")
                res[mod_idx].append(marker)
            if not coords or coords is None:
                coords = { model: [lat, lon] }
            else:
                coords[model] = [lat, lon]

            popups[model] = sum([res[mod_idx][i]['type'] == 'Popup' and 1 or 0
                    for i in range(len(res[mod_idx])) if 'type' in
                    res[mod_idx][i]]) + 1

            logging.debug("COORDS: %s", str(coords))
            logging.debug("POPUPS: %s", str(popups))
            return coords, popups, res

    return {}, {}, dash.no_update


@dash.callback(
    [Output('ts-modal', 'children'),
     Output('ts-modal', 'is_open'),
     Output({ 'tag': 'ts-button', 'index': ALL, 'random': ALL }, 'n_clicks')],
    [Input({ 'tag': 'ts-button', 'index': ALL, 'random': ALL }, 'n_clicks')],
    [State('model-dropdown', 'value'),
     State('model-date-picker', 'date'),
     State('variable-dropdown-forecast', 'value'),
     State('model-clicked-coords', 'data'),
     State('current-popups-stored', 'data')],
     prevent_initial_call=True
)
@cache.memoize(timeout=cache_timeout)
def show_timeseries(ts_button, mod, date, variable, coords, popups):
    """ Renders model comparison timeseries, retrieve timeseries according to
    coordinates selected. """
    from callback_tools import get_timeseries

    ctx = dash.callback_context
    if ctx.triggered:
        button_id = orjson.loads(ctx.triggered[0]["prop_id"].split(".")[0])
        logging.debug("BUTTONID: %s %s %s", mod, str(button_id), str(ts_button))
        logging.debug('SHOW TS COORDS: %s %s %s', coords, type(coords), mod)
        ts_popups = 0
        for m in popups:
            logging.debug("___%s___", m)
            ts_popups += popups[m]
            if m == button_id['index']:
                break

        if ts_popups == 0:
            ts_popups = 1

        if button_id['tag'] == 'ts-button' and ts_button[ts_popups-1] >= 1:
            lat, lon = coords[button_id['index']]

            if lat is None or lon is None:
                raise PreventUpdate

            logging.debug('SHOW TS """"" %s %s %s', mod, lat, lon)
            figure = get_timeseries(mod, date, variable, lat, lon, forecast=True)
            figure.update_layout(MODEBAR_LAYOUT_TS)
            ts_body = dbc.ModalBody(
                dcc.Graph(
                    id='timeseries-modal',
                    figure=figure,
                    config=MODEBAR_CONFIG_TS
                )
            )

            return ts_body, True, [0 for _ in ts_button]
        return dash.no_update, False, [0 for _ in ts_button]
    return dash.no_update, False, [0 for _ in ts_button]

@dash.callback(
    [Output({'tag': 'model-map', 'index': ALL, "n_clicks": ALL}, 'zoom'),
     Output({'tag': 'model-map', 'index': ALL, "n_clicks": ALL}, 'center')],
     Input('country-focus', 'n_clicks'),
    [State('model-dropdown', 'value'),
     State('country-zoom', 'value'),
     State('country-lat', 'value'),
     State('country-lon', 'value'),
     ],
    prevent_initial_call=True
)
def zoom_country(n_clicks, model, zoom, lat, lon):
    """ Set zoom and center over coordinates entered """
    count = len(model)
    if count == 0 or zoom is None:
        raise PreventUpdate
    zoom = [float(zoom)]
    center = [[float(lat), float(lon)]]
    return zoom*count, center*count

@dash.callback(
        Output('country-focus', 'n_clicks'),
        Output('country-zoom', 'value'),
        Output('country-lat', 'value'),
        Output('country-lon', 'value'),
        Input({'tag': 'model-map', 'index': ALL, "n_clicks": ALL}, 'viewport'),
        Input({'tag': 'model-map', 'index': ALL, "n_clicks": ALL}, 'id'),
        prevent_initial_call=True
        )
def zooms(viewport, models):
    """ Syncronize all maps to have same center and zoom in mosaic """
    ctx = dash.callback_context
    if ctx.triggered:
        changed = dict(ctx.triggered_id) 
        index = models.index(changed)
        if viewport[index] is not None:
            return 1, \
                    viewport[index]['zoom'], \
                    viewport[index]['center'][0], \
                    viewport[index]['center'][1]
    raise PreventUpdate

@dash.callback(
    [Output('slider-interval', 'disabled'),
     Output('slider-interval', 'n_intervals'),
     Output('open-timeseries', 'style'),
     Output('btn-play', 'className'),
     ],
    Input('btn-play', 'n_clicks'),
    [State('slider-interval', 'disabled'),
     State('model-slider-graph', 'value')],
    prevent_initial_call=True
    )
def start_stop_autoslider(n_play, disabled, value):
    """ Play/Pause map animation """
    ctx = dash.callback_context
    logging.debug("VALUE %s", value)
    if not value:
        value = 0

    if ctx.triggered:
        button_id = ctx.triggered[0]["prop_id"].split(".")[0]
        if button_id == 'btn-play' and disabled:
            ts_style = { 'display': 'none' }
            return not disabled, int(value/FREQ), ts_style, 'fa fa-pause text-center'
        elif button_id == 'btn-play' and not disabled:
            ts_style = { 'display': 'block' }
            return not disabled, int(value/FREQ), ts_style, 'fa fa-play text-center'

    raise PreventUpdate

@dash.callback(
    Output('model-slider-graph', 'value'),
    [Input('slider-interval', 'n_intervals')],
    prevent_initial_call=True
)
@cache.memoize(timeout=cache_timeout)
def update_slider(n):
    """ Update slider value according to the number of intervals """
    logging.debug('SERVER: updating model-slider-graph %s' + str(n))
    if not n:
        return
    if n >= 24:
        tstep = int(round(24*math.modf(n/24)[0], 0))
    else:
        tstep = int(n)
    logging.debug('SERVER: updating model-slider-graph %s' + str(tstep*FREQ))
    return tstep*FREQ

@dash.callback(
    Output('forecast-tab', 'children'),
    [Input('models-apply', 'n_clicks'),
     Input('prob-apply', 'n_clicks'),
     Input('was-apply', 'n_clicks')],
    [State({'tag': 'tab-name', 'index': ALL}, 'id')],
    prevent_initial_call=True
    )
@cache.memoize(timeout=cache_timeout)
def update_tab_content(models_clicks, prob_clicks, was_clicks, curtab):
    """ Updates tab content """
    ctx = dash.callback_context

    if ctx.triggered:
        button_id = ctx.triggered[0]["prop_id"].split(".")[0]
        if button_id not in ('models-apply', 'prob-apply', 'was-apply'):
            raise PreventUpdate

        logging.debug(":::::::::::: %s %s", len(curtab), curtab[0]['index'])
        curtab_name = curtab[0]['index']

        nexttab_name = button_id.replace('-apply', '')
        logging.debug(":::::::::::: %s %s", curtab_name, nexttab_name)
        if curtab_name != nexttab_name:
            return tab_forecast(nexttab_name)

        raise PreventUpdate

    raise PreventUpdate

@dash.callback(
    Output('graph-collection', 'children'),
    [Input('models-apply', 'n_clicks'),
     Input('model-slider-graph', 'value'),
     Input('model-date-picker', 'date')],
    [State('model-dropdown', 'value'),
     State('variable-dropdown-forecast', 'value'),
     State({'tag': 'view-style', 'index': ALL}, 'active'),
     State({'tag': 'model-map', 'index': ALL, "n_clicks": ALL}, 'zoom'),
     State({'tag': 'model-map', 'index': ALL, "n_clicks": ALL}, 'center'),
     ],
    prevent_initial_call=False
    )
@cache.memoize(timeout=cache_timeout)
def update_models_figure(n_clicks, tstep, date, model, variable, selected_tiles, zoom, center):
    """ Update mosaic of maps figures according to all parameters """
    ctx = dash.callback_context

    if ctx.triggered:
        button_id = ctx.triggered[0]["prop_id"].split(".")[0]
        if button_id != 'models-apply' and tstep is None and date is None:
            raise PreventUpdate
    elif tstep is None and date is None:
        raise PreventUpdate

    from callback_tools import get_model_figure
    from callback_tools import get_figure

    logging.debug('SERVER: calling figure from picker callback')

    st_time = time.time()

    logging.debug('SERVER: tstep %s', str(tstep))

    if date is not None:
        date = date.split(' ')[0]
        try:
            date = dt.strptime(
                date, "%Y-%m-%d").strftime("%Y%m%d")
        except:
            pass
        logging.debug('SERVER: callback date %s', date)
    else:
        date = END_DATE

    if model is None:
        model = DEFAULT_MODEL

    if variable is None:
        variable = DEFAULT_VAR

    if tstep is not None:
        tstep = int(tstep/FREQ)
    else:
        tstep = 0

    logging.debug('SERVER: tstep calc %s', tstep)
    # logging.debug('#### IDS, ZOOM, CENTER: %s %s %s', ids, zoom, center)

    # curr_mods = { i['index']: { 'zoom': z, 'center': c }
    #        for i, z, c in zip(ids, zoom, center) }

    figures = []

    ncols, nrows = calc_matrix(len(model))

    if not zoom:
        zoom = [None]
    if not center:
        center = [None]

    if len(model) != len(zoom):
        logging.debug("############## %s %s **********", len(model), len(zoom))
        zoom = [None for _ in model]
        center = [None for _ in model]

    logging.debug('#### ZOOM, CENTER: %s %s %s %s %s', zoom, center, model, ncols, nrows)
    selected_tiles = list(MAP_LAYERS.keys())[selected_tiles.index(True)]
    # for mod in model:
    #     mod_zoom = mod in curr_mods and curr_mods[mod]['zoom'] or None
    #     mod_center = mod in curr_mods and curr_mods[mod]['center'] or None
    for mod, mod_zoom, mod_center in zip(model, zoom, center):
        logging.debug("MOD %s ZOOM %s CENTER %s VIEW %s", mod, mod_zoom, mod_center, selected_tiles)

        # Get current tag
        if mod in MODELS:
            tag = 'model'
        else:
            tag = isinstance(tag, str) and tag or str(tag)

        # Get current index
        if isinstance(mod, str):
            index = mod
        else:
            index = str(mod)
    
        layers = get_model_figure(var=variable, model=mod, tstep=tstep, selected_date=date, 
                                  aspect=(nrows, ncols))
        figure = get_figure(aspect=(nrows, ncols), selected_tiles=selected_tiles, center=mod_center, 
                            zoom=mod_zoom, tag=tag, index=index, layers=layers)

        figure.style = {
            'height': '{}vh'.format(int(GRAPH_HEIGHT/nrows)),
        }

        # add n_clicks to id dict to reset the state of map, for centering
        figure.id['n_clicks'] = n_clicks
        # logging.debug('FIGURE KEYS 2 %s', figure.id)
        figures.append(figure)

    if DEBUG:
        for fig in figures:
            logging.debug("************ %s %s**************", fig.id, fig.style, )

    res = [
        dbc.Row(
            [
                dbc.Col(
                    figures[row+col+(row*(ncols-1))],
                    width=int(12/ncols),
                )
                for col in range(ncols)
                if len(figures) > row+col+(row*(ncols-1))
            ],
            align="start",
            # no_gutters is replaced with the className because of bootstrap update
            # no_gutters=True,
            className='g-0',
        ) for row in range(nrows)
    ]
    logging.debug("**** REQUEST TIME %s", str(time.time() - st_time))
    return res
