#!/usr/local/bin/python
# -*- coding: utf-8 -*-
import pickle
import os
import time
import re
import gc
import warnings
from itertools import product, chain
from functools import partial
from collections import OrderedDict
from multiprocessing import Pool, cpu_count

import numpy as np
import scipy as sc
import scipy.stats as stats
import pandas as pd
from peewee import fn
from IPython import embed

from qlknn.NNDB.model import Network, NetworkJSON, PostprocessSlice, PostprocessSlice_9D, db
from qlknn.NNDB.model import *
from qlknn.models.ffnn import QuaLiKizNDNN
from qlknn.training.datasets import shuffle_panda
from qlknn.plots.load_data import nameconvert
from qlknn.dataset.data_io import load_from_store
from qlknn.misc.analyse_names import split_parts, split_name
from qlknn.misc.tools import parse_dataset_name

if __name__ == '__main__':
    import matplotlib as mpl
    import matplotlib.pyplot as plt
    from matplotlib import gridspec, cycler

pretty = False

def mode_to_settings(mode):
    settings = {}
    if mode == 'debug':
        settings['plot']             = True
        settings['plot_pop']         = True
        settings['plot_nns']         = True
        settings['plot_slice']       = True
        settings['plot_poplines']    = True
        settings['plot_threshlines'] = True
        settings['plot_zerocolors']  = False
        settings['plot_thresh1line'] = False
        settings['calc_thresh1']     = False
        settings['hide_qualikiz']    = False
        settings['debug']            = True
        settings['parallel']         = False
        settings['plot_threshslope'] = False
    elif mode == 'quick':
        settings['plot']             = False
        settings['plot_pop']         = False
        settings['plot_nns']         = False
        settings['plot_slice']       = False
        settings['plot_poplines']    = False
        settings['plot_threshlines'] = False
        settings['plot_zerocolors']  = False
        settings['plot_thresh1line'] = False
        settings['calc_thresh1']     = False
        settings['hide_qualikiz']    = False
        settings['debug']            = False
        settings['parallel']         = True
        settings['plot_threshslope'] = False
    elif mode == 'pretty':
        settings['plot']             = True
        settings['plot_pop']         = False
        settings['plot_nns']         = True
        settings['plot_slice']       = False
        settings['plot_poplines']    = False
        settings['plot_threshlines'] = False
        settings['plot_zerocolors']  = False
        settings['plot_thresh1line'] = False
        settings['calc_thresh1']     = False
        settings['hide_qualikiz']    = False
        settings['debug']            = True
        settings['parallel']         = False
        settings['plot_threshslope'] = True
    return settings

def get_similar_not_in_table(table, max=20,
                             only_dim=None, only_sep=False,
                             no_particle=False, no_divsum=False,
                             no_mixed=True, target_names=None, no_mixedmode=True,
                             no_gam=True):
    non_sliced = (Network
                  .select()
                  .where(~fn.EXISTS(table.select().where(getattr(table, 'network') == Network.id)))
                 )
    if target_names is not None:
        non_sliced &= Network.select().where(Network.target_names == target_names)

    if only_dim is not None:
        non_sliced &= (Network.select()
                       .where(fn.array_length(Network.feature_names, 1) == only_dim)
                      )

    if no_mixed:
        non_sliced &= (Network.select()
                       .where(~(fn.array_to_string(Network.target_names, ',') % '%pf%' &
                                fn.array_to_string(Network.target_names, ',') % '%ef%'))
                       )
    if no_mixedmode:
        non_sliced &= (Network.select()
                       .where(~(
                           (fn.array_to_string(Network.target_names, ',') % '%ITG%' &
                            fn.array_to_string(Network.target_names, ',') % '%TEM%') |
                           (fn.array_to_string(Network.target_names, ',') % '%ITG%' &
                            fn.array_to_string(Network.target_names, ',') % '%ETG%') |
                           (fn.array_to_string(Network.target_names, ',') % '%TEM%' &
                            fn.array_to_string(Network.target_names, ',') % '%ETG%')
                               ))
                       )

    if no_gam:
        non_sliced &= (Network.select()
                       .where(~(fn.array_to_string(Network.target_names, ',') % '%gam%'))
                       )
    tags = []
    if no_divsum:
        tags.extend(["div", "plus"])
    if no_particle:
        tags.append('pf')
    if len(tags) != 0:
        filter = ~fn.array_to_string(Network.target_names, ',') % ('%' + tags[0] + '%')
        for tag in tags[1:]:
            filter &= ~(fn.array_to_string(Network.target_names, ',') % ('%' + tag + '%'))
        non_sliced &= (Network.select()
                       .where(filter)
                      )
    if only_sep:
        tags = ['TEM', 'ITG', 'ETG']
        filter = fn.array_to_string(Network.target_names, ',') % ('%' + tags[0] + '%')
        for tag in tags[1:]:
            filter |= fn.array_to_string(Network.target_names, ',') % ('%' + tag + '%')

    if non_sliced.count() > 0:
        network = non_sliced.get()
    else:
        raise Exception('No candidates found for slicing!')

    non_sliced &= (Network.select()
                  .where(Network.target_names == network.target_names)
                  .where(Network.feature_names == network.feature_names)
                  )
    non_sliced = non_sliced.limit(max)
    return non_sliced

def nns_from_NNDB(dim, max=20, only_dim=None):
    db.connect()
    if dim == 7:
        table = PostprocessSlice
    elif dim == 9:
        table = PostprocessSlice_9D
    non_sliced = get_similar_not_in_table(table, max=max, only_sep=True, no_particle=False, no_divsum=True, no_mixed=False, only_dim=only_dim)
    network = non_sliced.get()
    style = 'mono'
    if len(network.target_names) == 2:
        style = 'duo'
    elif len(network.target_names) == 3:
        style = 'triple'

    matches = []
    modes = []
    for target_name in network.target_names:
        matches.extend(re.compile('^.f.(ITG|ETG|TEM)_GB').findall(target_name))
        splitted = split_parts(target_name)
        if len(splitted) > 1:
            raise Exception('Error! Can only quickslice pure networks, not {!s}'.format(target_name))
        __, __, mode, __ = split_name(splitted[0])
        modes.append(mode)
    if modes[1:] == modes[:-1]:
        mode = modes[0]
        if mode == 'ITG':
            slicedim = 'Ati'
        elif mode == 'TEM' or mode == 'ETG':
            slicedim = 'Ate'
        else:
            raise Exception('Unknown mode {!s}'.format(mode))
    else:
        raise Exception('Unequal stability regime. Cannot determine slicedim')
    nn_list = {network.id: str(network.id) for network in non_sliced}
    print('Found {:d} {!s} with target {!s}'.format(len(non_sliced), network.__class__, network.target_names))

    nns = OrderedDict()
    for dbnn in non_sliced:
        nn = dbnn.to_QuaLiKizNN()
        nn.label = '_'.join([str(el) for el in [dbnn.__class__.__name__ , dbnn.id]])
        nns[nn.label] = nn

    db.close()
    return slicedim, style, nns

def nns_from_manual():
    nns = OrderedDict()

    dbnns = []
    labels = []

    dbnns.append(Network.get_by_id(1723))

    for ii, dbnn in enumerate(dbnns):
        nn = dbnn.to_QuaLiKizNN()
        if len(labels) == 0:
            nn.label = '_'.join([str(el) for el in [dbnn.__class__.__name__ , dbnn.id]])
        else:
            nn.label = labels[ii]
        nns[nn.label] = nn

    #nn = QuaLiKizNDNN.from_json('../../tests/gen3_test_files/Network_874_efiITG_GB/nn.json')
    #nn.label = 'manual'
    #nns[nn.label] = nn
    slicedim = 'Ati'
    if len(nn._target_names) == 1:
        style='mono'
    elif len(nn._target_names) == 2:
        style='duo'
    elif len(nn._target_names) == 3:
        style='triple'
    #from qlkANNk import QuaLiKiz4DNN
    #nns['4D'] = QuaLiKiz4DNN()
    #nns['4D'].label = '4D'
    #nns['4D']._target_names = ['efeITG_GB', 'efiITG_GB']
    db.close()
    return slicedim, style, nns

def prep_df(store, nns, unstack, filter_less=np.inf, filter_geq=-np.inf, shuffle=True, sort=False, calc_maxgam=False, clip=False, slice=None, frac=1):
    nn0 = list(nns.values())[0]
    target_names = nn0._target_names
    feature_names = nn0._feature_names

    input, data, const = load_from_store(store=store, columns=target_names)
    data.dropna(axis='index', how='all', inplace=True)

    try:
        input['logNustar'] = np.log10(input['Nustar'])
        del input['Nustar']
    except KeyError:
        print('No Nustar in dataset')

    if ('Zeff' == feature_names).any() and not ('Zeff' in input.columns):
        print('WARNING! creating Zeff. You should use a 9D dataset')
        input.insert(0, 'Zeff', np.full_like(input['Ati'], float(const['Zeff'])))
    if ('logNustar' == feature_names).any() and not ('logNustar' in input.columns):
        print('WARNING! creating logNustar. You should use a 9D dataset')
        input['logNustar']  = np.full_like(input['Ati'], np.log10(float(const['Nustar'])))

    if len(feature_names) == 4:
        print('WARNING! Slicing 7D to 4D dataset. You should use a 4D dataset')
        idx = input.index[(
            np.isclose(input['Ate'], 5.75,     atol=1e-5, rtol=1e-3) &
            np.isclose(input['An'], 2, atol=1e-5, rtol=1e-3) &
            np.isclose(input['x'], .45, atol=1e-5, rtol=1e-3)
        )]
        data = data.loc[idx]

    if any(input.columns != feature_names):
        print('WARNING! {!s} != {!s}, using 2* RAM to reorder.'.format(input.columns, feature_names))
        input = input[feature_names]

    #get_vars = target_names
    #data = pd.DataFrame()

    #dataset_vars = store.get_storer('/megarun1/flattened').non_index_axes[0][1]
    #for target_name in target_names:
    #    if target_name not in dataset_vars:
    #        print('WARNING! {!s} missing from dataset. Trying to reconstruct'.format(target_name))
    #        if target_name == 'efiTEM_GB_div_efeTEM_GB':
    #            parts = store.select('megarun1/flattened', columns=['efiTEM_GB', 'efeTEM_GB'])
    #        elif target_name == 'pfeTEM_GB_div_efeTEM_GB':
    #            parts = store.select('megarun1/flattened', columns=['pfeTEM_GB', 'efeTEM_GB'])
    #        elif target_name == 'efeITG_GB_div_efiITG_GB':
    #            parts = store.select('megarun1/flattened', columns=['efeITG_GB', 'efiITG_GB'])
    #        elif target_name == 'pfeITG_GB_div_efiITG_GB':
    #            parts = store.select('megarun1/flattened', columns=['pfeITG_GB', 'efiITG_GB'])
    #        else:
    #            raise Exception('Could not reconstruct {!s}'.format(target_name))
    #        se = parts.iloc[:,0] / parts.iloc[:,1]
    #        se.name = target_name
    #        data = data.append(se.to_frame())
    #        get_vars = get_vars[(get_vars != target_name)]

    input_cols = list(input.columns)
    print('Merging target and features')
    df = data = input.merge(data, left_index=True, right_index=True, copy=False)
    del input
    gc.collect()
    #df = input.join(data[target_names], how='inner')

    if calc_maxgam is True:
        df_gam = store.select('/megarun1/flattened', columns=['gam_leq_GB', 'gam_great_GB'])
        df_gam = (df_gam.max(axis=1)
                  .to_frame('maxgam')
        )
        df = df.join(df_gam)

    #itor = zip(['An', 'Ate', 'Ti_Te', 'q', 'smag', 'x'], ['0.00', '10.00', '1.00', '5.00', '0.40', '0.45'])
    #itor = zip(['Zeff', 'Ate', 'An', 'q', 'smag', 'x', 'Ti_Te', 'logNustar'], [1.0, 5.75, 2.5, 2.0, 0.10000000149011612, 0.33000001311302185, 1.0, -2.000217201545864])

    if slice is not None:
        for name, val in slice:
            df = df[np.isclose(df[name], float(val),     atol=1e-5, rtol=1e-3)]

    if clip is True:
        print('Clipping')
        df[target_names] = df[target_names].clip(filter_less, filter_geq, axis=1)
        #df = df[(df[target_names] < filter_less).all(axis=1)]
        #df = df[(df[target_names] >= filter_geq).all(axis=1)]

    print('Setting index')
    df.set_index(input_cols, inplace=True)
    if sort:
        print('Sorting')
        if sort and shuffle:
            print('WARNING! Sorting and shuffeling. Sort will be useless')
        df = df.sort_index(level=unstack)
    print('Unstacking slices')
    df = df.unstack(unstack)

    if shuffle:
        print("Every day I'm shuffling")
        df = shuffle_panda(df)

    if frac < 1:
        print('Taking {!s} fraction'.format(frac))
        if not shuffle:
            print('WARNING! Taking fraction without shuffle. You will always get the same slices!')
        idx = int(frac * len(df))
        df = df.iloc[:idx, :]
    #df = df.iloc[1040:2040,:]

    print('Converting to float64')
    df = df.astype('float64')
    print('dataset loaded!')
    return df, target_names

def is_unsafe(df, nns, slicedim):
    unsafe = True
    for nn in nns.values():
        slicedim_idx = nn._feature_names[nn._feature_names == slicedim].index[0]
        varlist = list(df.index.names)
        varlist.insert(slicedim_idx, slicedim)
        try:
            if ~np.all(varlist == nn._feature_names):
                unsafe = False
        except ValueError:
            raise Exception('Dataset has features {!s} but dataset has features {!s}'.format(varlist, list(nn._feature_names)))
    return unsafe


def calculate_thresh1(x, feature, target, debug=False):
    try:
        idx = target.index[target == 0][-1] #index of last zero
        slope, intercept, r_value, p_value, std_err = stats.linregress(feature[(target.index > idx) & ~target.isnull()], target[(target.index > idx) & ~target.isnull()])
        thresh_pred = x * slope + intercept
        thresh1 = x[thresh_pred < 0][-1]
    except (ValueError, IndexError):
        thresh1 = np.NaN
        if debug:
            print('No threshold1')
    return thresh1

def calculate_thresh2(feature, target, debug=False):
    if len(target.shape) > 1:
        raise NotImplementedError('2D threshold not implemented yet')
    try:
        idx = np.where(target == 0)[0][-1] #Only works for 1D
        idx2 = np.where(~np.isnan(target[idx+1:]))[0][0] + idx + 1
        #idx = np.arange(target.shape[0]),target.shape[1] - 1 - (target[:,::-1]==0).argmax(1) #Works for 2D
        thresh2 = (feature[idx] + feature[idx2]) / 2
    except IndexError:
        thresh2 = np.NaN
        if debug:
            print('No threshold2')

    return thresh2
#5.4 ms ± 115 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) total
def process_chunk(target_names, chunck, settings=None, unsafe=False):

    res = []
    for ii, row in enumerate(chunck.iterrows()):
        res.append(process_row(target_names, row, settings=settings, unsafe=unsafe))
    return res

def process_row(target_names, row, ax1=None, unsafe=False, settings=None):
    index, slice_ = row
    feature = slice_.index.levels[1]
    #target = slice.loc[target_names]
    target = slice_.values[:len(feature) * len(target_names)].reshape(len(target_names), len(feature))
    if np.all(np.logical_or(target == 0, np.isnan(target))):
        return (1,)
    else:
        # 156 µs ± 10.4 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each) (no zerocolors)
        thresh_nn = np.empty(len(target_names) * len(nns))
        thresh_nn_i = np.empty_like(thresh_nn, dtype='int64')
        popbacks = np.empty_like(thresh_nn)
        thresh1_misses = np.empty_like(thresh_nn)
        thresh2_misses = np.empty_like(thresh_nn)
        if settings['plot_zerocolors']:
            maxgam = slice_['maxgam']

        # Create slice, assume sorted
        # 14.8 µs ± 1.27 µs per loop (mean ± std. dev. of 7 runs, 100000 loops each)
        x = np.linspace(feature.values[0],
                        feature.values[-1],
                        200)
        #if plot:
        if not ax1 and settings['plot']:
            fig = plt.figure()
            if settings['plot_pop'] and settings['plot_slice']:
                gs = gridspec.GridSpec(2, 2, height_ratios=[10, 1], width_ratios=[5,1],
                                    left=0.05, right=0.95, wspace=0.05, hspace=0.05)
                ax2 = plt.subplot(gs[1,0])
                ax3 = plt.subplot(gs[0,1])
            if not settings['plot_pop'] and settings['plot_slice']:
                gs = gridspec.GridSpec(2, 1, height_ratios=[10, 2], width_ratios=[1],
                                    left=0.05, right=0.95, wspace=0.05, hspace=0.05)
                ax2 = plt.subplot(gs[1,0])
            if not settings['plot_pop'] and not settings['plot_slice']:
                gs = gridspec.GridSpec(1, 1, height_ratios=[1], width_ratios=[1],
                                    left=0.05, right=0.95, wspace=0.05, hspace=0.05)
            ax1 = plt.subplot(gs[0,0])
            #ax1.set_prop_cycle(cycler('color', ['#f1eef6','#d7b5d8','#df65b0','#dd1c77','#980043']))
            # http://tristen.ca/hcl-picker/#/clh/5/273/2A0A75/D59FEB
            #ax1.set_prop_cycle(cycler('color', ['#2A0A75','#6330B8','#9F63E2','#D59FEB']))
            if len(nns) == 1:
                color_range = np.array([.7])
            else:
                color_range = np.linspace(0, 0.9, len(nns))
            ax1.set_prop_cycle(cycler('color', plt.cm.plasma(color_range)))
            ax1.set_xlabel(nameconvert[slicedim])
            ax1.set_ylabel(nameconvert[list(nns.items())[0][1]._target_names[0]])
        if settings['calc_thresh1']:
            thresh1 = calculate_thresh1(x, feature, target, debug=settings['debug'])
            print('whyyy?')

        # 12.5 µs ± 970 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
        #if all(['ef' in name for name in target_names]):
        #    thresh2 = calculate_thresh2(feature.values, target[0,:], debug=settings['debug'])
        #elif all(['pf' in name for name in target_names]):
        thresh2 = calculate_thresh2(feature.values, np.abs(target[0,:]), debug=settings['debug'])
        #else:
        #    thresh2 = np.nan
        #    print('No thresh2!')
        #    embed()
        #    print('Weird stuff')

        if settings['plot'] and settings['plot_threshlines']:
            ax1.axvline(thresh2, c='black', linestyle='dashed')

        if settings['plot'] and settings['plot_threshslope']:
            if ~np.isnan(thresh2):
                pre_thresh =  x[x <= thresh2]
                ax1.plot(pre_thresh, np.zeros_like(pre_thresh), c='gray', linestyle='dashed')
                post_thresh =  x[x > thresh2]
                se = slice_.loc[target_names]
                se.index = se.index.droplevel()
                se = se.loc[se.index > thresh2].dropna()
                a = sc.optimize.curve_fit(lambda x, a: a * x, se.index-thresh2, se.values)[0][0]
                ax1.plot(post_thresh, a * (post_thresh-thresh2), c='gray', linestyle='dashed')

        # 13.7 µs ± 1.1 µs per loop (mean ± std. dev. of 7 runs, 100000 loops each)
        if unsafe:
            slice_list = [np.full_like(x, val) for val in index]
            slicedim_idx = np.nonzero(list(nns.values())[0]._feature_names.values == slicedim)[0][0]
            slice_list.insert(slicedim_idx, x)
        else:
            slice_dict = {name: np.full_like(x, val) for name, val in zip(df.index.names, index)}
            slice_dict[slicedim] = x



        # Plot target points
        if settings['plot'] and settings['plot_slice']:
            table = ax2.table(cellText=[[nameconvert[name] for name in df.index.names],
                                        ['{:.2f}'.format(xx) for xx in index]],cellLoc='center')
            table.auto_set_font_size(False)
            table.scale(1, 1.5)
            #table.set_fontsize(20)
            ax2.axis('tight')
            ax2.axis('off')
        #fig.subplots_adjust(bottom=0.2, transform=ax1.transAxes)


        # Plot nn lines
        nn_preds = np.ndarray([x.shape[0], 0])
        for ii, (nn_index, nn) in enumerate(nns.items()):
            clip_low = True
            if unsafe:
                low_bound = np.array([0 if ('ef' in name) and (not 'div' in name) else -np.inf for name in nn._target_names])
            else:
                low_bound = pd.Series({name: 0 if ('ef' in name) and (not 'div' in name) else -np.inf for name in nn._target_names})
            clip_high = False
            high_bound = None
            #if all(['ef' in name for name in nn._target_names]):
            #    clip_low = True
            #    low_bound = np.zeros((len(nn._target_names), 1))

            #    #high_bound = np.full((len(nn._target_names), 1), np.inf)
            #    clip_high = False
            #    high_bound = None
            #elif all(['pf' in name for name in nn._target_names]):
            #    #raise NotImplementedError('Particle bounds')
            #    clip_low = False
            #    low_bound = np.full((len(nn._target_names), 1), -80)
            #    clip_high = False
            #    high_bound = np.full((len(nn._target_names), 1), 80)
            #else:
            #    clip_low = False
            #    low_bound = None
            #    clip_high = False
            #    high_bound = None
            #    print('Mixed target!')
            #    #embed()
            #    print('Weird stuff')
            if unsafe:
                nn_pred = nn.get_output(np.array(slice_list).T, clip_low=clip_low, low_bound=low_bound, clip_high=clip_high, high_bound=high_bound, safe=not unsafe, output_pandas=False)
            else:
                nn_pred = nn.get_output(pd.DataFrame(slice_dict), clip_low=clip_low, low_bound=low_bound, clip_high=clip_high, high_bound=high_bound, safe=not unsafe, output_pandas=True).values
            nn_preds = np.concatenate([nn_preds, nn_pred], axis=1)

        if settings['plot'] and settings['plot_nns']:
            lines = []
            if style == 'triple':
                labels = np.repeat([nn.label for nn in nns.values()], 3)
                for ii in range(0, nn_preds.shape[1], 3):
                    lines.append(ax1.plot(x, nn_preds[:, ii], label=labels[ii])[0])
                    lines.append(ax1.plot(x, nn_preds[:, ii+1], label=labels[ii+1], c=lines[-1].get_color(), linestyle='dashed')[0])
                    lines.append(ax1.plot(x, nn_preds[:, ii+2], label=labels[ii+2], c=lines[-1].get_color(), linestyle='dotted')[0])
            elif style == 'duo':
                labels = np.repeat([nn.label for nn in nns.values()], 2)
                for ii in range(0, nn_preds.shape[1], 2):
                    lines.append(ax1.plot(x, nn_preds[:, ii], label=labels[ii])[0])
                    lines.append(ax1.plot(x, nn_preds[:, ii+1], label=labels[ii+1], c=lines[-1].get_color(), linestyle='dashed')[0])
            else:
                for ii, (nn, row) in enumerate(zip(nns.values(), nn_preds.T)):
                    pass
                    lines.append(ax1.plot(x, row, label=nn.label)[0])

        matrix_style = False
        if matrix_style:
            thresh_i = (np.arange(nn_preds.shape[1]),nn_preds.shape[0] - 1 - (nn_preds[::-1,:]==0).argmax(0))[1]
            thresh = x[thresh_i]
            thresh[thresh == x[-1]] = np.nan
        else:
            for ii, row in enumerate(nn_preds.T):
                if 'pf' in np.take(nn._target_names.values, ii, mode='wrap'):
                    thresh_nn_i[ii] = len(x)
                    thresh_nn[ii] = np.nan
                else:
                    try:
                        if row[-1] == 0:
                            thresh_nn[ii] = np.nan
                        else:
                            thresh_i = thresh_nn_i[ii] = np.where(np.diff(np.sign(row)))[0][-1]
                            thresh_nn[ii] = x[thresh_i]
                    except IndexError:
                        thresh_nn[ii] = np.nan

        if settings['plot'] and settings['plot_threshlines']:
            for ii, row in enumerate(thresh_nn):
                ax1.axvline(row, c=lines[ii].get_color(), linestyle='dotted')
                if settings['debug']:
                    print('network ', ii, 'threshold ', row)



        if matrix_style:
            masked = np.ma.masked_where(x[:, np.newaxis] > thresh, nn_preds)
            #popback_i = (masked.shape[0] - 1 - (masked[::1,:]!=0)).argmax(0)
            popback_i = masked.shape[0] - 1 - (masked.shape[0] - 1 - (masked[::-1,:]!=0)).argmin(0)
            popback = x[popback_i]
            popback[popback == x[-1]] = np.nan
        else:
            for ii, row in enumerate(nn_preds.T):
                if not np.isnan(thresh_nn[ii]):
                    try:
                        popback_i = np.flatnonzero(row[:thresh_nn_i[ii]])
                        popbacks[ii] = x[popback_i[-1]]
                    except (IndexError):
                        popbacks[ii] = np.nan
                else:
                    popbacks[ii] = np.nan

        # 5.16 µs ± 188 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

        wobble = np.abs(np.diff(nn_preds, n=2,axis=0))
        wobble_tot = np.mean(wobble, axis=0)
        with warnings.catch_warnings(): # col[ind:] is empty if no threshold
            warnings.simplefilter("ignore", category=RuntimeWarning)
            wobble_unstab = np.array([np.mean(col[ind:]) for ind, col in zip(thresh_nn_i + 1, wobble.T)])
        try:

            with warnings.catch_warnings(): # col[ind:] is empty if no threshold
                warnings.simplefilter("ignore", category=RuntimeWarning)
                thresh_2_i = np.where(np.abs(x - thresh2) == np.min(np.abs(x - thresh2)))[0][0]
            with warnings.catch_warnings(): # col[thresh_2_i:] is empty if no threshold
                warnings.simplefilter("ignore", category=RuntimeWarning)
                wobble_qlkunstab = np.array([np.mean(col[thresh_2_i:]) for col in wobble.T])
        except IndexError:
            thresh_2_i = np.nan
            wobble_qlkunstab = np.full_like(wobble_tot, np.nan)
        if settings['plot'] and settings['plot_pop']:
            thresh2_misses = thresh_nn - thresh2
            thresh2_popback = popbacks - thresh2
            slice_stats = np.array([thresh2_misses, thresh2_popback, np.log10(wobble_qlkunstab)]).T
            slice_strings = np.array(['{:.1f}'.format(xx) for xx in slice_stats.reshape(slice_stats.size)])
            slice_strings = slice_strings.reshape(slice_stats.shape)
            slice_strings = np.insert(slice_strings, 0, ['thre_mis', 'pop_mis', 'wobble_qlkunstab'], axis=0)
            table = ax3.table(cellText=slice_strings, loc='center')
            table.auto_set_font_size(False)
            ax3.axis('tight')
            ax3.axis('off')
            if settings['debug']:
                print(slice_stats.flatten())

        if settings['plot']:
            if settings['plot_zerocolors']:
                color = target.copy()
                color[(target == 0) & (maxgam == 0)] = 'green'
                color[(target != 0) & (maxgam == 0)] = 'red'
                color[(target == 0) & (maxgam != 0)] = 'magenta'
                color[(target != 0) & (maxgam != 0)] = 'blue'
            else:
                color='blue'
            if settings['hide_qualikiz']:
                color='white'
                zorder=1
                label=''
            else:
                zorder=1000
                #label = 'QuaLiKiz'
                #label = 'Turbulence model'
                #label=''
                labels=target_names
            markers = ['1', '2', '3']
            for label, column, marker in zip(labels, target, markers):
                ax1.scatter(feature[column != 0],
                            column[column != 0], c=color, label=label, marker=marker, zorder=zorder)
            ax1.scatter(feature[column==0],
                        column[column==0], edgecolors=color, marker='o', facecolors='none', zorder=zorder)

        # Plot regression
        if settings['plot'] and settings['plot_thresh1line'] and not np.isnan(thresh1):
            #plot_min = ax1.get_ylim()[0]
            plot_min = -0.1
            x_plot = x[(thresh_pred > plot_min) & (thresh_pred < ax1.get_ylim()[1])]
            y_plot = thresh_pred[(thresh_pred > plot_min) & (thresh_pred < ax1.get_ylim()[1])]
            ax1.plot(x_plot, y_plot, c='gray', linestyle='dotted')
            ax1.plot(x[x< thresh1], np.zeros_like(x[x< thresh1]), c='gray', linestyle='dotted')
            #ax1.axvline(thresh1, c='black', linestyle='dotted')

        slice_res = np.array([thresh_nn, popbacks, wobble_tot, wobble_unstab, wobble_qlkunstab]).T
        if settings['plot']:
            ax1.legend()
            ax1.set_ylim(bottom=min(ax1.get_ylim()[0], 0))
            plt.show()
            fig.savefig('slice.pdf', format='pdf', bbox_inches='tight')
            qlk_data = pd.DataFrame(target.T, columns=target_names, index=feature)
            cols = pd.MultiIndex.from_product([[nn.label for nn in nns.values()], target_names])
            nn_data = pd.DataFrame(nn_preds, columns=cols)
            nn_data.index = x
            nn_data.index.name = feature.name
            slice_data = pd.Series(dict(zip(df.index.names, index)))
            slice_latex = ('  {!s} &' * len(df.index.names)).format(*[nameconvert[name] for name in df.index.names]).strip(' &')
            slice_latex += ('\\\\\n' + ' {:.2f} &' * len(index)).format(*index).strip(' &')
            embed()
            plt.close(fig)
        return (0, thresh2, slice_res.flatten())
    #sliced += 1
    #if sliced % 1000 == 0:
    #    print(sliced, 'took ', time.time() - starttime, ' seconds')
def extract_stats(totstats, style):
    df = totstats.copy()
    df = df.reorder_levels([2,0,1], axis=1)

    results = pd.DataFrame()

    for relabs, measure in product(['rel', 'abs'], ['thresh', 'pop']):
        df2 = df[measure]
        qlk_data = df2['QLK']
        network_data = df2.drop('QLK', axis=1)
        if relabs == 'rel':
            mis = network_data.subtract(qlk_data, level=1).divide(qlk_data, level=1)
        elif relabs == 'abs':
            mis = network_data.subtract(qlk_data, level=1)

        quant1 = 0.025
        quant2 = 1 - quant1
        quant = mis.quantile([quant1, quant2])
        results['_'.join([measure, relabs, 'mis', 'median'])] = mis.median()
        results['_'.join([measure, relabs, 'mis', '95width'])] = quant.loc[quant2] - quant.loc[quant1]

        if relabs == 'abs':
            results['_'.join(['no', measure, 'frac'])] = mis.isnull().sum() / len(mis)
    results['wobble_unstab'] = df['wobble_unstab'].mean()
    results['wobble_qlkunstab'] = df['wobble_qlkunstab'].mean()
    results['wobble_tot'] = df['wobble_tot'].mean()

    if style == 'duo' or style == 'triple':
        duo_results = pd.DataFrame()
        measure = 'thresh'
        df2 = df[measure]
        network_data = df2.drop('QLK', axis=1)
        network_data = network_data.reorder_levels([1, 0], axis=1)
        efelike_name = network_data.columns[0][0]
        efilike_name = network_data.columns[1][0]
        if not (efelike_name.startswith('efe') and efilike_name.startswith('efi')):
            raise Exception('{!s} does not start with efe or {!s} does not start with efi'.format(efelike_name, efilike_name))
        mis = network_data[efilike_name] - network_data[efelike_name]
        quant = mis.quantile([quant1, quant2])
        duo_results['dual_thresh_mismatch_median'] = mis.median()
        duo_results['dual_thresh_mismatch_95width'] = quant.loc[quant2] - quant.loc[quant1]
        duo_results['no_dual_thresh_frac'] = mis.isnull().sum() / len(mis)
    else:
        duo_results = pd.DataFrame()
    return results, duo_results

def extract_nn_stats(results, duo_results, nns, frac, store_name, submit_to_nndb=False):
    db.connect()
    for network_name, res in results.unstack().iterrows():
        network_class, network_number = network_name.split('_')
        nn = nns[network_name]
        res_dict = {'network': network_number}
        res_dict['frac'] = frac
        res_dict['store_name'] = store_basename

        for stat, val in res.unstack(level=0).iteritems():
            res_dict[stat] = val.loc[nn._target_names].values

        try :
            duo_res = duo_results.loc[network_name]
            res_dict.update(duo_res)
        except KeyError:
            pass

        __, dim, __ = get_store_params(store_name)
        if dim == 7:
            postprocess_slice = PostprocessSlice(**res_dict)
        elif dim == 9:
            postprocess_slice = PostprocessSlice_9D(**res_dict)

        if submit_to_nndb is True:
            postprocess_slice.save()
    db.close()

def dump_results_to_disk(res, duo_res, frac, store_name, runname='slicestat'):
    res.to_csv(runname + '_results.csv')
    if len(duo_res) != 0:
        duo_res.to_csv(runname + '_duo_results.csv')
    meta = pd.Series({'frac': frac,
                      'store_name': store_name})
    meta.to_csv(runname + '_metadata.csv')

def get_store_params(store_name):
    unstable, set, gen, dim, label, filter = parse_dataset_name(store_name)
    if filter is not None:
        filter = int(filter)
    gen, dim = int(gen), int(dim)
    return gen, dim, filter

if __name__ == '__main__':
    nn_set = 'duo'
    nn_set = 'best'
    mode = 'pretty'
    mode = 'debug'
    dump_to_disk = False
    #dump_to_disk = True
    submit_to_nndb = False
    mode = 'quick'
    submit_to_nndb = True

    store_root = '../..'
    store_basename = 'gen3_7D_nions0_flat_filter8.h5.1'
    store_name = os.path.join(store_root, store_basename)
    store = pd.HDFStore(store_name, 'r')

    __, dim, __ = get_store_params(store_basename)
    slicedim, style, nns = nns_from_NNDB(dim, max=100)
    #slicedim, style, nns = nns_from_manual()
    #slicedim = 'An'

    if style != 'similar':
        labels=True
    else:
        labels=False

    if mode == 'quick':
        clip = False
        filter_geq = None
        filter_less = None
    else:
        clip = True
        filter_geq = -120
        filter_less = 120

    itor = None
    frac = 0.05
    df, target_names = prep_df(store, nns, slicedim, filter_less=filter_less, filter_geq=filter_geq, slice=itor, frac=frac)
    gc.collect()
    unsafe = is_unsafe(df, nns, slicedim)
    if not unsafe:
        print('Warning! Cannot use unsafe mode')

    settings = mode_to_settings(mode)
    if mode == 'pretty':
        plt.style.use('./thesis.mplstyle')
        mpl.rcParams.update({'font.size': 16})
    else:
        nameconvert = {name: name for name in df.columns.names +  df.index.names + target_names.tolist()}

    if settings['parallel']:
        num_processes = cpu_count()
        chunk_size = int(df.shape[0]/num_processes)
        chunks = [df.loc[df.index[i:i + chunk_size]] for i in range(0, df.shape[0], chunk_size)]
        pool = Pool(processes=num_processes)
        print('Using {:d} processes'.format(num_processes))

    print('Starting {:d} slices for {:d} networks'.format(len(df), len(nns)))
    starttime = time.time()
    #n=20
    #newind = np.hstack([np.repeat(np.array([*df.index]), n, axis=0), np.tile(np.linspace(df.columns.levels[1][0], df.columns.levels[1][-1], n), len(df))[:, None]])
    #embed()
    if not settings['parallel']:
        results = [process_chunk(target_names, df, settings=settings, unsafe=unsafe)]
    else:
        results = pool.map(partial(process_chunk, target_names, settings=settings, unsafe=unsafe), chunks)
    #for row in df.iterrows():
    #    process_row(row)
    print(len(df), 'took ', time.time() - starttime, ' seconds')

    zero_slices = 0
    totstats = []
    qlk_thresh = []
    for result in chain(*results):
        if result[0] == 1:
            zero_slices += 1
        else:
            totstats.append(result[2])
            qlk_thresh.append(result[1])

    stats = ['thresh', 'pop', 'wobble_tot', 'wobble_unstab', 'wobble_qlkunstab']
    totstats = pd.DataFrame(totstats, columns=pd.MultiIndex.from_tuples(list(product([nn.label for nn in nns.values()], target_names, stats))))

    qlk_columns = list(product(['QLK'], target_names, stats))
    qlk_data = np.full([len(totstats), len(qlk_columns)], np.nan)
    qlk_data[:, ::] = np.tile(qlk_thresh, np.array([len(qlk_columns),1])).T
    qlk_data = pd.DataFrame(qlk_data, columns=pd.MultiIndex.from_tuples(qlk_columns))

    totstats = totstats.join(qlk_data)
    res, duo_res = extract_stats(totstats, style)
    res.index.names = ['network_label', 'target_names']
    if len(duo_res) != 0:
        duo_res.index.names = ['network_label', 'target_names']

    # dump to disk
    if dump_to_disk:
        dump_results_to_disk(res, duo_res, frac, store_basename)
    extract_nn_stats(res, duo_res, nns, frac, store_basename, submit_to_nndb=submit_to_nndb)


    #print('WARNING! If you continue, you will overwrite ', 'totstats_' + style + '.pkl')
    #embed()
    #totstats._metadata = {'zero_slices': zero_slices}
    #with open('totstats_' + style + '.pkl', 'wb') as file_:
    #    pickle.dump(totstats, file_)
