import dash
import dash_core_components as dcc
import dash_html_components as html
from dash.dependencies import Input, Output
import plotly.graph_objs as go
import pandas as pd
import numpy as np
from math import ceil
from math import floor

df = pd.DataFrame(np.random.randn(10000, 10000))

selected_x0 = 0
selected_x1 = len(df.columns)
selected_y0 = 0
selected_y1 = len(df.index)

def select_df(selection):
    global selected_x0
    global selected_x1
    global selected_y0
    global selected_y1

    if 'xaxis.range[0]' in selection and 'xaxis.range[1]' in selection:
        x0 = int(floor((selection['xaxis.range[0]'])))
        x1 = int(ceil((selection['xaxis.range[1]'])))
        selected_x0 = max(x0, 0)
        selected_x1 = min(x1, len(df.columns))
    elif not selection or 'xaxis.autorange' in selection:
        selected_x0 = 0
        selected_x1 = len(df.columns)

    if 'yaxis.range[0]' in selection and 'yaxis.range[1]' in selection:
        y0 = int(floor((selection['yaxis.range[0]'])))
        y1 = int(ceil((selection['yaxis.range[1]'])))
        selected_y0 = max(y0, 0)
        selected_y1 = min(y1, len(df.columns))
    elif not selection or 'yaxis.autorange' in selection:
        selected_y0 = 0
        selected_y1 = len(df.index)

    if selected_x1 - selected_x0 <= 11:
        x_scale = 1
    else:
        x_scale = int(ceil(selected_x1 - selected_x0) / 10)
        x_scale = int(2 * round(x_scale / 2))
        selected_x0 = int(floor(selected_x0 / x_scale))
        selected_x1 = int(ceil(selected_x1 / x_scale))
        selected_x0 *= x_scale
        selected_x1 *= x_scale
        if selected_x1 > len(df.columns):
            selected_x1 = len(df.columns)
            selected_x0 = selected_x1 - x_scale * 10
            assert selected_x0 >= 0

    if selected_y1 - selected_y0 <= 11:
        y_scale = 1
    else:
        y_scale = int(ceil(selected_y1 - selected_y0) / 10)
        y_scale = int(2 * round(y_scale / 2))
        selected_y0 = int(floor(selected_y0 / y_scale))
        selected_y1 = int(ceil(selected_y1 / y_scale))
        selected_y0 *= y_scale
        selected_y1 *= y_scale
        if selected_y1 > len(df.columns):
            selected_y1 = len(df.columns)
            selected_y0 = selected_y1 - y_scale * 10
            assert selected_y0 >= 0

    selected_df = df.iloc[range(selected_y0, selected_y1), range(selected_x0, selected_x1)]

    if x_scale > 1:
        x_groups = np.array([x_scale * (column // x_scale)
                             for column in range(len(selected_df.columns))])
        selected_df = selected_df.groupby(by=x_groups, axis=1, sort=False).mean()
        selected_df.set_axis([int(ceil(selected_x0 + x_scale * (column + 0.5)))
                              for column in range(len(selected_df.columns))],
                             axis=1, inplace=True)

    if y_scale > 1:
        y_groups = np.array([y_scale * (index // y_scale)
                             for index in range(len(selected_df.index))])
        selected_df = selected_df.groupby(by=y_groups, axis=0, sort=False).mean()
        selected_df.set_axis([int(ceil(selected_y0 + y_scale * (index + 0.5)))
                              for index in range(len(selected_df.index))],
                             axis=0, inplace=True)

    return selected_df

app = dash.Dash(__name__)

selected_df = select_df({})
figure = {
    'data': [{
        'x': selected_df.columns,
        'y': selected_df.index,
        'z': selected_df.values,
        'type': 'heatmap',
    }],
    'layout': {
        'xaxis': {
            'scaleanchor': 'y',
        }
    },
}

app.layout = html.Div(
    style={
        'background-color': 'red',
        'height': '100vmin',
        'width': '100vmin',
        'overflow': 'hidden',
        'position': 'relative',
    },
    children=[
        # dcc.Loading(children=[  # WILL NOT WORK
            dcc.Graph(
                id='graph',
                figure=figure,
                style={
                    'position': 'absolute',
                    'top': 0,
                    'left': 0,
                    'background-color': 'blue',
                    'width': '100%',
                    'height': '100%',
                    'overflow': 'hidden',
                },
            ),
        # ]),
    ],
)

@app.callback(Output('graph', 'figure'),
              [Input('graph', 'relayoutData')])
def selected_df_figure(selection):
    selected_df = select_df(selection or {})
    figure['data'][0]['x'] = selected_df.columns
    figure['data'][0]['y'] = selected_df.index
    figure['data'][0]['z'] = selected_df.values
    return figure

if __name__ == '__main__':
    app.run_server(debug=True)