from __future__ import division
import re
from itertools import product
import gc
import os
import warnings
from collections import OrderedDict
import copy

from IPython import embed
import pandas as pd
import numpy as np
import dask.dataframe as dd

from qlknn.dataset.data_io import put_to_store_or_df, save_to_store, load_from_store, sep_prefix
from qlknn.misc.analyse_names import heat_vars, particle_vars, particle_diffusion_vars, momentum_vars, is_partial_diffusion, is_partial_particle, is_pure_heat, split_name, is_transport
from qlknn.misc.tools import profile

@profile
def regime_filter(data, geq, less):
    """ Filter the dataset based on the total ion/electron heat flux
    This filter is used to constain the dataset to experimentally relevant
    fluxes. We have an estimate for the expected total heat flux, and
    filter out the full datapoint.

    Args:
        data: The dataset to filter. Usually `pandas.DataFrame`. Needs to
              contain 'efi_GB' and 'efe_GB'.
        geq:  Lower bound on ef[i|e] heat flux. Inclusive bound.
        less: Upper bound on ef[i|e] heat flux. Exclusive bound.
    """
    within = pd.Series(np.full(len(data), True, dtype='bool'), index=data.index)
    within &= (data['efe_GB'] < less) & (data['efi_GB'] < less)
    within &= (data['efe_GB'] >= geq) & (data['efi_GB'] >= geq)
    data = data.loc[within]
    return data

@profile
def div_filter(store, filter_bounds=None):
    """ Filter flux_div_flux variables based on bounds
    We know from experience the maximum relative difference in flux between
    the ions and electrons of different modes. As such we remove the fluxpoint
    if it falls outside the given bounds.

    For heat fluxes:     low_bound < flux_div_flux < high_bound.
    For particle fluxes: low_bound < abs(flux_div_flux) < high_bound

    Note: Technically, as heat fluxes are non-negative anyway, we could use the same
    bounds.

    Args:
        store:         The store name or `pd.HDFStore` to apply the filter to

    Kwargs:
        filter_bounds: A dictionary with as keys the flux_div_flux variable to filter
                       and values a tuple with (low_bound, high_bound). For defaults,
                       see `filter_defaults['div']`
    """
    all_filter_bounds = copy.deepcopy(filter_defaults['div'])
    if filter_bounds is None:
        filter_bounds = {}
    all_filter_bounds.update(filter_bounds)


    for group in store:
        if isinstance(store, pd.HDFStore):
            name = group.lstrip(sep_prefix)
        else:
            name = group

        # Read filter bound. If not in dict, skip this variable
        if name in all_filter_bounds:
            low, high = all_filter_bounds[name]
        else:
            continue

        # Load the variable from store, save the name
        se = store[group]
        se.name = name
        # And save the pre-filter instances for our debugging print
        pre = np.sum(~se.isnull())

        if is_partial_particle(name):
            se = se.abs()

        # Apply the filter and save to store/dataframe
        put_to_store_or_df(store, se.name, store[group].loc[(low < se) & (se < high)])
        print('{:5.2f}% of sane unstable {!s:<9} points inside div bounds'.format(np.sum(~store[group].isnull()) / pre * 100, group))
    return store

@profile
def stability_filter(data):
    """ Filter out the stable points based on growth rate

    QuaLiKiz gives us growth rate information, so we can use this to filter
    out all stable points from the dataset. Of course, we only check the
    stability of the relevant mode: TEM for TEM and ITG for ITG, 'elec' for
    ETG (legacy, sorry). multi-scale (e.g. electron and ion-scale) for
    electron-heat-flux like vars (efe, dfe and family) and ion-scale
    otherwise (e.g. efi, pfe, pfi). TEM and ITG-stabilty are defined in
    `hypercube_to_pandas`. We define electron-unstable if we have a nonzero
    growthrate for kthetarhos <= 2, ion-unstable if we have a nonzero
    growthrate for kthetarhos > 2, and multiscale-unstable if electron-unstable
    or ion-unstable.

    Args:
        data: `pd.DataFrame` containing the data to be filtered, and `TEM`, `ITG`,
              `gam_leq_GB`, `gam_great_GB`

    """
    for col in data.columns:
        splitted = re.compile('(?=.*)(.)(|ITG|ETG|TEM)_(GB|SI|cm)').split(col)
        try:
            transp, species, mode, norm = split_name(col)
        except ValueError:
            print('skipping {!s}'.format(col))

        if not is_transport(col):
            print('skipping {!s}'.format(col))
            continue
        # First check for which regime this variable should be filtered
        if mode == 'TEM':
            gam_filter = 'tem'
        elif mode == 'ITG':
            gam_filter = 'itg'
        elif mode == 'ETG':
            gam_filter = 'elec'
        elif transp in heat_vars and species == 'e':
            gam_filter = 'multi'
        else:
            gam_filter = 'ion'

        pre = np.sum(~data[col].isnull())
        print(col, gam_filter)
        # Now apply a filter based on the regime of the variable
        if gam_filter == 'ion':
            data[col] = data[col].loc[data['gam_leq_GB'] != 0]
        elif gam_filter == 'elec':
            data[col] = data[col].loc[data['gam_great_GB'] != 0]
        elif gam_filter == 'multi':
            if 'gam_great_GB' in data.columns:
                data[col] = data[col].loc[(data['gam_leq_GB'] != 0) | (data['gam_great_GB'] != 0)]
            else:
                # If gam_great_GB is not there, the QuaLiKiz run was not with electron scale
                data[col] = data[col].loc[(data['gam_leq_GB'] != 0)]
        elif gam_filter == 'tem':
            data[col] = data[col].loc[data['TEM'] == True]
        elif gam_filter == 'itg':
            data[col] = data[col].loc[data['ITG'] == True]
        print('{:5.2f}% of sane {!s:<9} points unstable at {!s:<5} scale'.format(np.sum(~data[col].isnull()) / pre * 100, col, gam_filter))
    return data

@profile
def negative_filter(data):
    """ Check if none of the heat-flux variables is negative

    Only checks on `heat_vars` e.g. efe_GB, efiTEM_GB etc.

    Args:
        data to perform the 'negative check' on

    Returns:
        Per-element `True` if none of the checked heat-flux variables is negative
    """
    filter_functions = {'negative': negative_filter,
                    'ck': ck_filter,
                    'septot': septot_filter,
                    'ambipolar': ambipolar_filter,
                    'femtoflux': femtoflux_filter}
    #anyisneg = pd.Series(np.full(len(data), True, dtype='bool'), index=data.index)
    heat_cols = [col for col in data.columns if is_pure_heat(col)]
    nonnegative = (data[heat_cols] >= 0).all(axis=1)
    return nonnegative

@profile
def ck_filter(data, bound):
    """ Check if convergence checks cki and cki are within bounds"""
    ck = (data[['cki', 'cke']].abs() < bound).all(axis=1)
    return ck

@profile
def septot_filter(data, septot_factor, startlen=None):
    """ Check if ITG/TEM/ETG heat flux !>> total_flux"""
    if startlen is None and not isinstance(data, dd.DataFrame):
        startlen = len(data)
    difference_okay = pd.Series(np.full(len(data), True, dtype='bool'), index=data.index)
    sepnames = []
    for type, spec in product(heat_vars, ['i', 'e']):
        totname = type + spec + '_GB'
        if spec == 'i': # no ETG
            seps = ['ITG', 'TEM']
        else: # All modes
            seps = ['ETG', 'ITG', 'TEM']
        sepnames = [type + spec + sep + '_GB' for sep in seps if type + spec + sep + '_GB' in data.columns]
        print('Checking {!s}'.format(sepnames))
        difference_okay &= data[sepnames].abs().le(septot_factor * data[totname].abs(), axis=0).all(axis=1)
        if startlen is not None:
            print('After filter {!s:<6} {!s:<6} {:.2f}% left'.format('septot', totname, 100*difference_okay.sum()/startlen))
    return difference_okay

@profile
def ambipolar_filter(data, bound):
    """ Check if ambipolarity is conserved """
    return (data['absambi'] < bound) & (data['absambi'] > 1/bound)

@profile
def femtoflux_filter(data, bound):
    """ Check if flux is no 'femto_flux', a very small non-zero flux"""
    fluxes = [col for col in data if len(re.compile('(?=.*)(.)(|ITG|ETG|TEM)_(GB|SI|cm)').split(col)) == 5 if re.compile('(?=.*)(.)(|ITG|ETG|TEM)_(GB|SI|cm)').split(col)[0] in particle_vars + heat_vars + momentum_vars]
    no_femto = pd.Series(np.full(len(data), True, dtype='bool'), index=data.index)
    absflux = data[fluxes].abs()
    no_femto &= ~((absflux < bound) & (absflux != 0)).any(axis=1)
    return no_femto

def sanity_filter(data, ck_bound, septot_factor, ambi_bound, femto_bound,
                  stored_negative_filter=None,
                  stored_ck_filter=None,
                  stored_ambipolar_filter=None,
                  stored_septot_filter=None,
                  stored_femtoflux_filter=None,
                  startlen=None):
    """ Filter out insane points

    There are points where we do not trust QuaLiKiz. These points are
    filtered out by functions defined in this module. Currently:

        negative_filter:  Points with negative heat flux
        ck_filter:        Points with too high convergence errors
        septot_factor:    Points where sep flux >> total flux
        ambipolar_filter: Points that don't conserve ambipolarity
        femtoflux_filter: Point with very tiny fluxes

    Optionally one can provide a earlier-stored filter in the form
    of a list of indices to remove from the dataset
    Args:
        ck_bound:        Maximum ck[i/e]
        septot_factor:   Maximum factor between tot flux and sep flux
        ambi_bound:      Maximum factor between dq_i/dt and dq_e/dt
        femto_bound:     Maximum value of what is defined as femtoflux

    Kwargs:
        stored_[filter]: List of indices contained in filter [Default: None]
        starlen:         Total amount of points at start of function. By
                         default all points in dataset.
    """
    if startlen is None and not isinstance(data, dd.DataFrame):
        startlen = len(data)
    # Throw away point if negative heat flux
    if stored_negative_filter is None:
        data = data.loc[negative_filter(data)]
    else:
        data = data.reindex(index=data.index.difference(stored_negative_filter), copy=False)
    if startlen is not None:
        print('After filter {!s:<13} {:.2f}% left'.format('negative', 100*len(data)/startlen))
    else:
        print('filter {!s:<13} done'.format('negative'))
    gc.collect()

    # Throw away point if cke or cki too high
    if stored_ck_filter is None:
        data = data.loc[ck_filter(data, ck_bound)]
    else:
        data = data.reindex(index=data.index.difference(stored_ck_filter), copy=False)
    if startlen is not None:
        print('After filter {!s:<13} {:.2f}% left'.format('ck', 100*len(data)/startlen))
    else:
        print('filter {!s:<13} done'.format('ck'))
    gc.collect()

    # Throw away point if sep flux is way higher than tot flux
    if stored_septot_filter is None:
        data = data.loc[septot_filter(data, septot_factor, startlen=startlen)]
    else:
        data = data.reindex(index=data.index.difference(stored_septot_filter), copy=False)
    if startlen is not None:
        print('After filter {!s:<13} {:.2f}% left'.format('septot', 100*len(data)/startlen))
    else:
        print('filter {!s:<13} done'.format('septot'))
    gc.collect()

    # Throw away point if ambipolarity is not conserved
    if stored_ambipolar_filter is None:
        data = data.loc[ambipolar_filter(data, ambi_bound)]
    else:
        data = data.reindex(index=data.index.difference(stored_ambipolar_filter), copy=False)
    if startlen is not None:
        print('After filter {!s:<13} {:.2f}% left'.format('ambipolar', 100*len(data)/startlen))
    else: