""" Decoding tools"""

import numpy as np
from neurosynth.base.mask import Masker
from neurosynth.base import imageutils
from neurosynth.analysis.reduce import average_within_regions
from os import path
import pandas as pd
import matplotlib.pyplot as plt
from six import string_types

class Decoder:

    def __init__(self, dataset=None, method='pearson', features=None,
                 mask=None, image_type='association-test_z', threshold=0.001):
        """ Initialize a new Decoder instance.

        Args:
            dataset: An optional Dataset instance containing features to use in
                decoding.
            method: The decoding method to use (optional). By default, Pearson
                correlation.
            features: Optional list of features to use in decoding. If None,
                use all features found in dataset. If features is a list of
                strings, use only the subset of features in the Dataset that
                are named in the list. If features is a list of filenames,
                ignore the dataset entirely and use only the features passed as
                image files in decoding.
            mask: An optional mask to apply to features and input images. If
                None, will use the one in the current Dataset.
            image_type: An optional string indicating the type of image to use
                when constructing feature-based images. See
                meta.analyze_features() for details. By default, uses reverse
                inference z-score images.
            threshold: If decoding from a Dataset instance, this is the feature
                threshold to use to generate the feature maps used in the
                decoding.

        """

        self.dataset = dataset

        if dataset is not None:
            self.masker = self.dataset.masker
            if features is None:
                features = dataset.get_feature_names()
            if mask is not None:
                self.masker.add(mask)
        elif mask is not None:
            self.masker = Masker(mask)
        else:
            self.masker = None

        self.method = method.lower()

        if self.method == 'roi':
            self.feature_names = features
        else:
            self.load_features(features, image_type=image_type,
                               threshold=threshold)

    def decode(self, images, save=None, round=4, names=None, **kwargs):
        """ Decodes a set of images.

        Args:
          images: The images to decode. Can be:
            - A single String specifying the filename of the image to decode
            - A list of filenames
            - A single NumPy array containing the image data
          save: Optional filename to save results to. If None (default), returns
            all results as an array.
          round: Optional integer indicating number of decimals to round result
            to. Defaults to 4.
          names: Optional list of names corresponding to the images in filenames.
            If passed, must be of same length and in same order as filenames.
            By default, the columns in the output will be named using the image
            filenames.

        Returns:
          An n_features x n_files numpy array, where each feature is a row and
          each image is a column. The meaning of the values depends on the
          decoding method used. """

        if isinstance(images, string_types):
            images = [images]

        if isinstance(images, list):
            imgs_to_decode = imageutils.load_imgs(images, self.masker)
        else:
            imgs_to_decode = images

        methods = {
            'pearson': self._pearson_correlation,
            'dot': self._dot_product,
            'roi': self._roi_association
        }

        result = np.around(
            methods[self.method](imgs_to_decode, **kwargs), round)

        # if save is not None:

        if names is None:
            if type(images).__module__ == np.__name__:
                names = ['image_%d' % i for i in range(images.shape[1])]
            elif self.method == 'roi':
                names = ['cluster_%d' % i for i in range(result.shape[1])]
            else:
                names = images

        result = pd.DataFrame(result, columns=names, index=self.feature_names)

        if save is not None:
            result.to_csv(save, index_label='Feature')
        return result

    def set_method(self, method):
        """ Set decoding method. """
        self.method = method

    def load_features(self, features, image_type=None, from_array=False,
                      threshold=0.001):
        """ Load features from current Dataset instance or a list of files.
        Args:
            features: List containing paths to, or names of, features to
                extract. Each element in the list must be a string containing
                either a path to an image, or the name of a feature (as named
                in the current Dataset). Mixing of paths and feature names
                within the list is not allowed.
            image_type: Optional suffix indicating which kind of image to use
                for analysis. Only used if features are taken from the Dataset;
                if features is a list of filenames, image_type is ignored.
            from_array: If True, the features argument is interpreted as a
                string pointing to the location of a 2D ndarray on disk
                containing feature data, where rows are voxels and columns are
                individual features.
            threshold: If features are taken from the dataset, this is the
                threshold passed to the meta-analysis module to generate fresh
                images.

        """
        if from_array:
            if isinstance(features, list):
                features = features[0]
            self._load_features_from_array(features)
        elif path.exists(features[0]):
            self._load_features_from_images(features)
        else:
            self._load_features_from_dataset(
                features, image_type=image_type, threshold=threshold)

    def _load_features_from_array(self, features):
        """ Load feature data from a 2D ndarray on disk. """
        self.feature_images = np.load(features)
        self.feature_names = range(self.feature_images.shape[1])

    def _load_features_from_dataset(self, features=None, image_type=None,
                                    threshold=0.001):
        """ Load feature image data from the current Dataset instance. See
        load_features() for documentation.
        """
        self.feature_names = self.dataset.feature_table.feature_names
        if features is not None:
            self.feature_names = [f for f in features
                                  if f in self.feature_names]
        from neurosynth.analysis import meta
        self.feature_images = meta.analyze_features(
            self.dataset, self.feature_names, image_type=image_type,
            threshold=threshold)
        # Apply a mask if one was originally passed
        if self.masker.layers:
            in_mask = self.masker.get_mask(in_global_mask=True)
            self.feature_images = self.feature_images[in_mask, :]

    def _load_features_from_images(self, images, names=None):
        """ Load feature image data from image files.

        Args:
          images: A list of image filenames.
          names: An optional list of strings to use as the feature names. Must
            be in the same order as the images.
        """
        if names is not None and len(names) != len(images):
            raise Exception(
                "Lists of feature names and images must be of same length!")
        self.feature_names = names if names is not None else images
        self.feature_images = imageutils.load_imgs(images, self.masker)

    def train_classifiers(self, features=None):
        ''' Train a set of classifiers '''
        # for f in features:
        #     clf = Classifier(None)
        #     self.classifiers.append(clf)
        pass

    def _pearson_correlation(self, imgs_to_decode):
        """ Decode images using Pearson's r.

        Computes the correlation between each input image and each feature
        image across voxels.

        Args:
            imgs_to_decode: An ndarray of images to decode, with voxels in rows
                and images in columns.

        Returns:
            An n_features x n_images 2D array, with each cell representing the
            pearson correlation between the i'th feature and the j'th image
            across all voxels.
        """
        x, y = imgs_to_decode.astype(float), self.feature_images.astype(float)
        return self._xy_corr(x, y)

    def _dot_product(self, imgs_to_decode):
        """ Decoding using the dot product.
        """
        return np.dot(imgs_to_decode.T, self.feature_images).T

    def _roi_association(self, imgs_to_decode, value='z', binarize=None):
        """ Computes the strength of association between activation in a mask
        and presence/absence of a semantic feature. This is essentially a
        generalization of the voxel-wise reverse inference z-score to the
        multivoxel case.
        """
        imgs_to_decode = imgs_to_decode.squeeze()
        x = average_within_regions(self.dataset, imgs_to_decode).astype(float)
        y = self.dataset.feature_table.data[self.feature_names].values
        if binarize is not None:
            y[y > binarize] = 1.
            y[y < 1.] = 0.
        r = self._xy_corr(x.T, y)
        if value == 'r':
            return r
        elif value == 'z':
            f_r = np.arctanh(r)
            return f_r * np.sqrt(y.shape[0] - 3)

    def _xy_corr(self, x, y):
        x, y = x - x.mean(0), y - y.mean(0)
        x, y = x / np.sqrt((x ** 2).sum(0)), y / np.sqrt((y ** 2).sum(0))
        return x.T.dot(y).T

    def plot_polar(self, data, n_top=3, overplot=False, labels=None,
                   palette='husl'):

        n_panels = data.shape[1]

        if labels is None:
            labels = []
            for i in range(n_panels):
                labels.extend(data.iloc[:, i].order(ascending=False)
                              .index[:n_top])
            labels = np.unique(labels)

        data = data.loc[labels, :]

        # Use hierarchical clustering to order
        from scipy.spatial.distance import pdist
        from scipy.cluster.hierarchy import linkage, leaves_list
        dists = pdist(data, metric='correlation')
        pairs = linkage(dists)
        order = leaves_list(pairs)
        data = data.iloc[order, :]
        labels = [labels[i] for i in order]

        theta = np.linspace(0.0, 2 * np.pi, len(labels), endpoint=False)
        if overplot:
            fig, ax = plt.subplots(1, 1, subplot_kw=dict(polar=True))
            fig.set_size_inches(10, 10)
        else:
            fig, axes = plt.subplots(1, n_panels, sharex=False, sharey=False,
                                     subplot_kw=dict(polar=True))
            fig.set_size_inches((6 * n_panels, 6))
        # A bit silly to import seaborn just for this...
        # should extract just the color_palette functionality.
        import seaborn as sns
        colors = sns.color_palette(palette, n_panels)
        for i in range(n_panels):
            if overplot:
                alpha = 0.2
            else:
                ax = axes[i]
                alpha = 0.8
            ax.set_ylim(data.values.min(), data.values.max())
            d = data.iloc[:, i].values
            ax.fill(theta, d, color=colors[i], alpha=alpha, ec='k',
                    linewidth=0)
            ax.fill(theta, d, alpha=1.0, ec=colors[i],
                    linewidth=2, fill=False)
            ax.set_xticks(theta)
            ax.set_xticklabels(labels, fontsize=18)
            [lab.set_fontsize(18) for lab in ax.get_yticklabels()]
            ax.set_title('Cluster %d' % i, fontsize=22, y=1.12)
        plt.tight_layout()
        return plt


""" Classification and decoding related tools """

import numpy as np
from functools import reduce
from sklearn.feature_selection.univariate_selection import SelectKBest
import re
from six import string_types

def feature_selection(feat_select, X, y):
    """" Implements various kinds of feature selection """
    # K-best
    if re.match('.*-best', feat_select) is not None:
        n = int(feat_select.split('-')[0])

        selector = SelectKBest(k=n)

        import warnings
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', category=UserWarning)
            features_selected = np.where(
                selector.fit(X, y).get_support() is True)[0]

    elif re.match('.*-randombest', feat_select) is not None:
        n = int(feat_select.split('-')[0])

        from random import shuffle
        features = range(0, X.shape[1])
        shuffle(features)

        features_selected = features[:n]

    return features_selected

def get_score(X, y, clf, scoring='accuracy'):
    prediction = clf.predict(X)

    if scoring == 'accuracy':
        from sklearn.metrics import accuracy_score
        score = accuracy_score(y, prediction)
    elif scoring == 'f1':
        from sklearn.metrics import f1_score
        score = f1_score(y, prediction)
    else:
        score = scoring(y, prediction.squeeze())

    return prediction, score

def classify_by_features(dataset, features, studies=None, method='SVM',
                         scikit_classifier=None):
    pass

def regularize(X, method='scale'):
    if method == 'scale':
        from sklearn import preprocessing
        return preprocessing.scale(X, with_mean=False)
    else:
        raise Exception('Unrecognized regularization method')

def get_studies_by_regions(dataset, masks, threshold=0.08, remove_overlap=True,
                           studies=None, features=None,
                           regularization="scale"):
    """ Set up data for a classification task given a set of masks

        Given a set of masks, this function retrieves studies associated with
        each mask at the specified threshold, optionally removes overlap and
        filters by studies and features, and returns studies by feature matrix
        (X) and class labels (y)

        Args:
            dataset: a Neurosynth dataset
            maks: a list of paths to Nifti masks
            threshold: percentage of voxels active within the mask for study
                to be included
            remove_overlap: A boolean indicating if studies studies that
                appear in more than one mask should be excluded
            studies: An optional list of study names used to constrain the set
                used in classification. If None, will use all features in the
                dataset.
            features: An optional list of feature names used to constrain the
                set used in classification. If None, will use all features in
                the dataset.
            regularize: Optional boolean indicating if X should be regularized

        Returns:
            A tuple (X, y) of np arrays.
            X is a feature by studies matrix and y is a vector of class labels
    """

    import nibabel as nib
    import os

    # Load masks using NiBabel

    try:
        loaded_masks = [nib.load(os.path.relpath(m)) for m in masks]
    except OSError:
        print('Error loading masks. Check the path')

    # Get a list of studies that activate for each mask file--i.e.,  a list of
    # lists

    grouped_ids = [dataset.get_studies(mask=m, activation_threshold=threshold)
                   for m in loaded_masks]

    # Flattened ids

    flat_ids = reduce(lambda a, b: a + b, grouped_ids)

    # Remove duplicates

    if remove_overlap:
        import collections
        flat_ids = [id for (id, count) in
                    collections.Counter(flat_ids).items() if count == 1]
        grouped_ids = [[x for x in m if x in flat_ids] for m in
                       grouped_ids]  # Remove

    # Create class label(y)
    y = [[idx] * len(ids) for (idx, ids) in enumerate(grouped_ids)]
    y = reduce(lambda a, b: a + b, y)  # Flatten
    y = np.array(y)

    # Extract feature set for each class separately
    X = [dataset.get_feature_data(ids=group_ids, features=features)
         for group_ids in grouped_ids]

    X = np.vstack(tuple(X))

    if regularization:
        X = regularize(X, method=regularization)

    return (X, y)

def get_feature_order(dataset, features):
    """ Returns a list with the order that features requested appear in
    dataset """
    all_features = dataset.get_feature_names()

    i = [all_features.index(f) for f in features]

    return i

def classify_regions(dataset, masks, method='ERF', threshold=0.08,
                     remove_overlap=True, regularization='scale',
                     output='summary', studies=None, features=None,
                     class_weight='auto', classifier=None,
                     cross_val='4-Fold', param_grid=None, scoring='accuracy'):
    """ Perform classification on specified regions

        Given a set of masks, this function retrieves studies associated with
        each mask at the specified threshold, optionally removes overlap and
        filters by studies and features. Then it trains an algorithm to
        classify studies based on features and tests performance.

        Args:
            dataset: a Neurosynth dataset
            maks: a list of paths to Nifti masks
            method: a string indicating which method to used.
                'SVM': Support Vector Classifier with rbf kernel
                'ERF': Extremely Randomized Forest classifier
                'Dummy': A dummy classifier using stratified classes as
                    predictor
            threshold: percentage of voxels active within the mask for study
                to be included
            remove_overlap: A boolean indicating if studies studies that
                appear in more than one mask should be excluded
            regularization: A string indicating type of regularization to use.
                If None, performs no regularization.
                'scale': Unit scale without demeaning
            output: A string indicating output type
                'summary': Dictionary with summary statistics including score
                    and n
                'summary_clf': Same as above but also includes classifier
                'clf': Only returns classifier
                Warning: using cv without grid will return an untrained
                classifier
            studies: An optional list of study names used to constrain the set
                used in classification. If None, will use all features in the
                dataset.
            features: An optional list of feature names used to constrain the
                set used in classification. If None, will use all features in
                the dataset.
            class_weight: Parameter to pass to classifier determining how to
                weight classes
            classifier: An optional sci-kit learn classifier to use instead of
                pre-set up classifiers set up using 'method'
            cross_val: A string indicating type of cross validation to use.
                Can also pass a scikit_classifier
            param_grid: A dictionary indicating which parameters to optimize
                using GridSearchCV. If None, no GridSearch will be used

        Returns:
            A tuple (X, y) of np arrays.
            X is a feature by studies matrix and y is a vector of class labels
    """

    (X, y) = get_studies_by_regions(dataset, masks, threshold, remove_overlap,
                                    studies, features,
                                    regularization=regularization)

    return classify(X, y, method, classifier, output, cross_val,
                    class_weight, scoring=scoring, param_grid=param_grid)

def classify(X, y, clf_method='ERF', classifier=None, output='summary_clf',
             cross_val=None, class_weight=None, regularization=None,
             param_grid=None, scoring='accuracy', refit_all=True,
             feat_select=None):
    """ Wrapper for scikit-learn classification functions
    Imlements various types of classification and cross validation """

    # Build classifier
    clf = Classifier(clf_method, classifier, param_grid)

    # Fit & test model with or without cross-validation
    if cross_val is not None:
        score = clf.cross_val_fit(X, y, cross_val, scoring=scoring,
                                  feat_select=feat_select,
                                  class_weight=class_weight)
    else:
        # Does not support scoring function
        score = clf.fit(X, y, class_weight=class_weight).score(X, y)

    # Return some stuff...
    from collections import Counter

    if output == 'clf':
        return clf
    else:
        if output == 'summary':
            output = {'score': score, 'n': dict(Counter(y))}
        elif output == 'summary_clf':
            output = {
                'score': score,
                'n': dict(Counter(y)),
                'clf': clf,
                'features_selected': clf.features_selected,
                'predictions': clf.predictions
            }

        return output

class Classifier:

    def __init__(self, clf_method='ERF', classifier=None, param_grid=None):
        """ Initialize a new classifier instance """

        # Set classifier
        self.features_selected = None
        self.predictions = None

        if classifier is not None:
            self.clf = classifier

            from sklearn.svm import LinearSVC
            import random
            if isinstance(self.clf, LinearSVC):
                self.clf.set_params().random_state = random.randint(0, 200)
        else:
            if clf_method == 'SVM':
                from sklearn import svm
                self.clf = svm.SVC()
            elif clf_method == 'ERF':
                from sklearn.ensemble import ExtraTreesClassifier
                self.clf = ExtraTreesClassifier(
                    n_estimators=100, max_depth=None, min_samples_split=1,
                    random_state=0)
            elif clf_method == 'GBC':
                from sklearn.ensemble import GradientBoostingClassifier
                self.clf = GradientBoostingClassifier(n_estimators=100,
                                                      max_depth=1)
            elif clf_method == 'Dummy':
                from sklearn.dummy import DummyClassifier
                self.clf = DummyClassifier(strategy='stratified')
            else:
                raise Exception('Unrecognized classification method')

        if isinstance(param_grid, dict):
            from sklearn.grid_search import GridSearchCV
            self.clf = GridSearchCV(estimator=self.clf,
                                    param_grid=param_grid)

    def fit(self, X, y, cv=None, class_weight='auto'):
        """ Fits X to outcomes y, using clf """

        # Incorporate error checking such as :
        # if isinstance(self.classifier, ScikitClassifier):
        #     do one thingNone
        # otherwiseNone.

        self.X = X
        self.y = y

        self.set_class_weight(class_weight=class_weight, y=y)

        self.clf = self.clf.fit(X, y)

        return self.clf

    def set_class_weight(self, class_weight='auto', y=None):
        """ Sets the class_weight of the classifier to match y """

        if class_weight is None:
            cw = None

            try:
                self.clf.set_params(class_weight=cw)
            except ValueError:
                pass

        elif class_weight == 'auto':
            c = np.bincount(y)
            ii = np.nonzero(c)[0]
            c = c / float(c.sum())
            cw = dict(zip(ii[::-1], c[ii]))

            try:
                self.clf.set_params(class_weight=cw)
            except ValueError:
                import warnings
                warnings.warn(
                    "Tried to set class_weight, but failed. The classifier "
                    "probably doesn't support it")

    def cross_val_fit(self, X, y, cross_val='4-Fold', scoring='accuracy',
                      feat_select=None, class_weight='auto'):
        """ Fits X to outcomes y, using clf and cv_method """

        from sklearn import cross_validation

        self.X = X
        self.y = y

        self.set_class_weight(class_weight=class_weight, y=y)

        # Set cross validator
        if isinstance(cross_val, string_types):
            if re.match('.*-Fold', cross_val) is not None:
                n = int(cross_val.split('-')[0])
                self.cver = cross_validation.StratifiedKFold(self.y, n)
            else:
                raise Exception('Unrecognized cross validation method')
        else:
            self.cver = cross_val

        if feat_select is not None:
            self.features_selected = []

        # Perform cross-validated classification
        from sklearn.grid_search import GridSearchCV
        if isinstance(self.clf, GridSearchCV):
            import warnings

            if feat_select is not None:
                warnings.warn(
                    "Cross-validated feature selection not supported with "
                    "GridSearchCV")
            self.clf.set_params(cv=self.cver, scoring=scoring)

            with warnings.catch_warnings():
                warnings.simplefilter('ignore', category=UserWarning)
                self.clf = self.clf.fit(X, y)

            self.cvs = self.clf.best_score_
        else:
            self.cvs = self.feat_select_cvs(
                feat_select=feat_select, scoring=scoring)

        if feat_select is not None:
            fs = feature_selection(
                feat_select, X, y)
            self.features_selected.append(fs)

            X = X[:, fs]

        self.clf.fit(X, y)

        return self.cvs.mean()

    def feat_select_cvs(self, scoring=None, feat_select=None):
        """ Returns cross validated scores (just like cross_val_score),
        but includes feature selection as part of the cross validation loop """

        scores = []
        self.predictions = []

        for train, test in self.cver:
            X_train, X_test, y_train, y_test = self.X[
                train], self.X[test], self.y[train], self.y[test]

            if feat_select is not None:
                # Get which features are kept
                fs = feature_selection(
                    feat_select, X_train, y_train)

                self.features_selected.append(fs)

                # Filter X to only keep selected features
                X_train, X_test = X_train[
                    :, fs], X_test[:, fs]

            # Set scoring (not implement as accuracy is default)

            # Train classifier
            self.clf.fit(X_train, y_train)

            # Test classifier
            predicition, s = get_score(
                X_test, y_test, self.clf, scoring=scoring)

            scores.append(s)
            self.predictions.append((y_test, predicition))

        return np.array(scores)

    def fit_dataset(self, dataset, y, features=None,
                    feature_type='features'):
        """ Given a dataset, fits either features or voxels to y """

        # Get data from dataset

        if feature_type == 'features':
            X = np.rot90(dataset.feature_table.data.toarray())
        elif feature_type == 'voxels':
            X = np.rot90(dataset.image_table.data.toarray())

        self.sk_classifier.fit(X, y)


""" Parsing expression grammar for feature-based study selection. """

from ply import lex
from ply import yacc
import pandas as pd
import logging

logger = logging.getLogger('neurosynth.lexparser')

class Lexer(object):

    tokens = (
        'WORD', 'FLOAT', 'ANDNOT', 'OR', 'AND', 'LPAR', 'RPAR', 'LT', 'RT'
    )

    t_ANDNOT = r'\&\~'
    t_AND = r'\&'
    t_OR = r'\|'
    t_LPAR = r'\('
    t_RPAR = r'\)'
    t_LT = r'\<'
    t_RT = r'\>'

    t_ignore = ' \t'

    def __init__(self, dataset=None):
        self.dataset = dataset

    def t_WORD(self, t):
        r'[a-zA-Z\_\-\*]+'
        return t

    def t_FLOAT(self, t):
        r'[0-9\.]+'
        t.value = float(t.value)
        return t

    def build(self, **kwargs):
        self.lexer = lex.lex(module=self, optimize=1, **kwargs)

    def t_error(self, t):
        logger.error("Illegal character %s!" % t.value[0])
        t.lexer.skip(1)

    def test(self, data):

        self.lexer.input(data)
        while True:
            tok = self.lexer.token()
            if not tok:
                break      # No more input
            print(tok)

class Parser(object):

    def __init__(self, lexer, dataset, threshold=0.001, func=sum):

        self.lexer = lexer
        self.dataset = dataset
        self.threshold = threshold
        self.func = func
        self.tokens = lexer.tokens

    def p_list_andnot(self, p):
        'list : list ANDNOT list'
        p[0] = p[1].loc[set(p[1].index) - set(p[3].index)]

    def p_list_and(self, p):
        'list : list AND list'
        p[0] = pd.concat(
            [p[1], p[3]], axis=1).dropna().apply(self.func, axis=1)

    def p_list_or(self, p):
        'list : list OR list'
        p[0] = pd.concat(
            [p[1], p[3]], axis=1).fillna(0.0).apply(self.func, axis=1)

    def p_list_lt(self, p):
        'list : list LT freq'
        p[0] = p[1][p[1] < p[3]]

    def p_list_rt(self, p):
        'list : list RT freq'
        p[0] = p[1][p[1] >= p[3]]

    def p_feature_words(self, p):
        '''feature : WORD WORD
                | feature WORD'''
        p[0] = ' '.join([p[1], p[2]])

    def p_list_feature(self, p):
        '''list : feature
            | WORD '''
        p[0] = self.dataset.get_studies(
            features=p[1], frequency_threshold=self.threshold, func=self.func,
            return_type='weights')

    def p_list_expr(self, p):
        'list : LPAR list RPAR'
        p[0] = p[2]

    def p_freq_float(self, p):
        'freq : FLOAT'
        p[0] = p[1]

    def build(self, **kwargs):
        self.parser = yacc.yacc(module=self, **kwargs)

    def p_error(self, p):
        print(p)

    def parse(self, input):
        return self.parser.parse(input)


""" Dimensionality reduction methods"""

import os
import numpy as np
from neurosynth.base.dataset import Dataset
from neurosynth.base import imageutils
from neurosynth.tests.utils import get_resource_path
import logging
import subprocess
import pandas as pd
import shutil
from os.path import dirname, join

logger = logging.getLogger('neurosynth.cluster')

def average_within_regions(dataset, regions, masker=None, threshold=None,
                           remove_zero=True):
    """ Aggregates over all voxels within each ROI in the input image.

    Takes a Dataset and a Nifti image that defines distinct regions, and
    returns a numpy matrix of  ROIs x mappables, where the value at each
    ROI is the proportion of active voxels in that ROI. Each distinct ROI
    must have a unique value in the image; non-contiguous voxels with the
    same value will be assigned to the same ROI.

    Args:
        dataset: Either a Dataset instance from which image data are
            extracted, or a Numpy array containing image data to use. If
            the latter, the array contains voxels in rows and
            features/studies in columns. The number of voxels must be equal
            to the length of the vectorized image mask in the regions
            image.
        regions: An image defining the boundaries of the regions to use.
            Can be one of:
            1) A string name of the NIFTI or Analyze-format image
            2) A NiBabel SpatialImage
            3) A list of NiBabel images
            4) A 1D numpy array of the same length as the mask vector in
                the Dataset's current Masker.
        masker: Optional masker used to load image if regions is not a
            numpy array. Must be passed if dataset is a numpy array.
        threshold: An optional float in the range of 0 - 1 or integer. If
            passed, the array will be binarized, with ROI values above the
            threshold assigned to True and values below the threshold
            assigned to False. (E.g., if threshold = 0.05, only ROIs in
            which more than 5% of voxels are active will be considered
            active.) If threshold is integer, studies will only be
            considered active if they activate more than that number of
            voxels in the ROI.
        remove_zero: An optional boolean; when True, assume that voxels
            with value of 0 should not be considered as a separate ROI, and
            will be ignored.

    Returns:
        A 2D numpy array with ROIs in rows and mappables in columns.
    """

    if masker is not None:
        masker = masker
    else:
        if isinstance(dataset, Dataset):
            masker = dataset.masker
        else:
            if not type(regions).__module__.startswith('numpy'):
                raise ValueError(
                    "If dataset is a numpy array and regions is not a numpy "
                    "array, a masker must be provided.")

    if not type(regions).__module__.startswith('numpy'):
        regions = masker.mask(regions)

    if isinstance(dataset, Dataset):
        dataset = dataset.get_image_data(dense=False)

    # If multiple images are passed, give each one a unique value
    if regions.ndim == 2:
        m = regions
        for i in range(regions.shape[1]):
            _nz = np.nonzero(m[:, i])[0]
            if isinstance(threshold, int):
                m[_nz, i] = 1.0
            else:
                m[_nz, i] = 1.0 / np.count_nonzero(m[:, i])

    # Otherwise create an ROI-coding matrix
    else:
        labels = np.unique(regions)

        if remove_zero:
            labels = labels[np.nonzero(labels)]

        n_regions = labels.size

        m = np.zeros((regions.size, n_regions))
        for i in range(n_regions):
            if isinstance(threshold, int):
                m[regions == labels[i], i] = 1.0
            else:
                m[regions == labels[i], i] = 1.0 / \
                    np.sum(regions == labels[i])

    # Call dot() on the array itself as this will use sparse matrix
    # multiplication if possible.
    result = dataset.T.dot(m).T

    if threshold is not None:
        result[result < threshold] = 0.0
        result = result.astype(bool)

    return result

def apply_grid(dataset, masker=None, scale=5, threshold=None):
    """ Imposes a 3D grid on the brain volume and averages across all voxels
    that fall within each cell.
    Args:
        dataset: Data to apply grid to. Either a Dataset instance, or a numpy
            array with voxels in rows and features in columns.
        masker: Optional Masker instance used to map between the created grid
            and the dataset. This is only needed if dataset is a numpy array;
            if dataset is a Dataset instance, the Masker in the dataset will
            be used.
        scale: int; scaling factor (in mm) to pass onto create_grid().
        threshold: Optional float to pass to reduce.average_within_regions().
    Returns:
        A tuple of length 2, where the first element is a numpy array of
        dimensions n_cubes x n_studies, and the second element is a numpy
        array, with the same dimensions as the Masker instance in the current
        Dataset, that maps voxel identities onto cell IDs in the grid.
    """
    if masker is None:
        if isinstance(dataset, Dataset):
            masker = dataset.masker
        else:
            raise ValueError(
                "If dataset is a numpy array, a masker must be provided.")

    grid = imageutils.create_grid(masker.volume, scale)
    cm = masker.mask(grid, in_global_mask=True)
    data = average_within_regions(dataset, cm, threshold)
    return (data, grid)

def get_random_voxels(dataset, n_voxels):
    """ Returns mappable data for a random subset of voxels.

    May be useful as a baseline in predictive analyses--e.g., to compare
    performance of a more principled feature selection method with simple
    random selection.

    Args:
        dataset: A Dataset instance
        n_voxels: An integer specifying the number of random voxels to select.

    Returns:
        A 2D numpy array with (randomly-selected) voxels in rows and mappables
        in columns.
    """
    voxels = np.arange(dataset.masker.n_vox_in_vol)
    np.random.shuffle(voxels)
    selected = voxels[0:n_voxels]
    return dataset.get_image_data(voxels=selected)

def _get_top_words(model, feature_names, n_top_words=40):
    """ Return top forty words from each topic in trained topic model.
    """
    topic_words = []
    for topic in model.components_:
        top_words = [feature_names[i] for i in topic.argsort()[:-n_top_words-1:-1]]
        topic_words += [top_words]
    return topic_words

def run_lda(abstracts, n_topics=50, n_words=31, n_iters=1000, alpha=None,
            beta=0.001):
    """ Perform topic modeling using Latent Dirichlet Allocation with the
    Java toolbox MALLET.

    Args:
        abstracts:  A pandas DataFrame with two columns ('pmid' and 'abstract')
                    containing article abstracts.
        n_topics:   Number of topics to generate. Default=50.
        n_words:    Number of top words to return for each topic. Default=31,
                    based on Poldrack et al. (2012).
        n_iters:    Number of iterations to run in training topic model.
                    Default=1000.
        alpha:      The Dirichlet prior on the per-document topic
                    distributions.
                    Default: 50 / n_topics, based on Poldrack et al. (2012).
        beta:       The Dirichlet prior on the per-topic word distribution.
                    Default: 0.001, based on Poldrack et al. (2012).

    Returns:
        weights_df: A pandas DataFrame derived from the MALLET
                    output-doc-topics output file. Contains the weight assigned
                    to each article for each topic, which can be used to select
                    articles for topic-based meta-analyses (accepted threshold
                    from Poldrack article is 0.001). [n_topics]+1 columns:
                    'pmid' is the first column and the following columns are
                    the topic names. The names of the topics match the names
                    in df (e.g., topic_000).
        keys_df:    A pandas DataFrame derived from the MALLET
                    output-topic-keys output file. Contains the top [n_words]
                    words for each topic, which can act as a summary of the
                    topic's content. Two columns: 'topic' and 'terms'. The
                    names of the topics match the names in weights (e.g.,
                    topic_000).
    """
    if abstracts.index.name != 'pmid':
        abstracts.index = abstracts['pmid']

    resdir = os.path.abspath(get_resource_path())
    tempdir = os.path.join(resdir, 'topic_models')
    absdir = os.path.join(tempdir, 'abstracts')
    if not os.path.isdir(tempdir):
        os.mkdir(tempdir)

    if alpha is None:
        alpha = 50. / n_topics

    # Check for presence of abstract files and convert if necessary
    if not os.path.isdir(absdir):
        print('Abstracts folder not found. Creating abstract files...')
        os.mkdir(absdir)
        for pmid in abstracts.index.values:
            abstract = abstracts.loc[pmid]['abstract']
            with open(os.path.join(absdir, str(pmid) + '.txt'), 'w') as fo:
                fo.write(abstract)

    # Run MALLET topic modeling
    print('Generating topics...')
    mallet_bin = join(dirname(dirname(__file__)),
                      'resources/mallet/bin/mallet')
    import_str = ('{mallet} import-dir '
                  '--input {absdir} '
                  '--output {outdir}/topic-input.mallet '
                  '--keep-sequence '
                  '--remove-stopwords').format(mallet=mallet_bin,
                                               absdir=absdir,
                                               outdir=tempdir)

    train_str = ('{mallet} train-topics '
                 '--input {out}/topic-input.mallet '
                 '--num-topics {n_topics} '
                 '--num-top-words {n_words} '
                 '--output-topic-keys {out}/topic_keys.txt '
                 '--output-doc-topics {out}/doc_topics.txt '
                 '--num-iterations {n_iters} '
                 '--output-model {out}/saved_model.mallet '
                 '--random-seed 1 '
                 '--alpha {alpha} '
                 '--beta {beta}').format(mallet=mallet_bin, out=tempdir,
                                         n_topics=n_topics, n_words=n_words,
                                         n_iters=n_iters,
                                         alpha=alpha, beta=beta)

    subprocess.call(import_str, shell=True)
    subprocess.call(train_str, shell=True)

    # Read in and convert doc_topics and topic_keys.
    def clean_str(string):
        return os.path.basename(os.path.splitext(string)[0])

    def get_sort(lst):
        return [i[0] for i in sorted(enumerate(lst), key=lambda x:x[1])]

    topic_names = ['topic_{0:03d}'.format(i) for i in range(n_topics)]

    # doc_topics: Topic weights for each paper.
    # The conversion here is pretty ugly at the moment.
    # First row should be dropped. First column is row number and can be used
    # as the index.
    # Second column is 'file: /full/path/to/pmid.txt' <-- Parse to get pmid.
    # After that, odd columns are topic numbers and even columns are the
    # weights for the topics in the preceding column. These columns are sorted
    # on an individual pmid basis by the weights.
    n_cols = (2 * n_topics) + 1
    dt_df = pd.read_csv(os.path.join(tempdir, 'doc_topics.txt'),
                        delimiter='\t', skiprows=1, header=None, index_col=0)
    dt_df = dt_df[dt_df.columns[:n_cols]]

    # Get pmids from filenames
    dt_df[1] = dt_df[1].apply(clean_str)

    # Put weights (even cols) and topics (odd cols) into separate dfs.
    weights_df = dt_df[dt_df.columns[2::2]]
    weights_df.index = dt_df[1]
    weights_df.columns = range(n_topics)

    topics_df = dt_df[dt_df.columns[1::2]]
    topics_df.index = dt_df[1]
    topics_df.columns = range(n_topics)

    # Sort columns in weights_df separately for each row using topics_df.
    sorters_df = topics_df.apply(get_sort, axis=1)
    weights = weights_df.as_matrix()
    sorters = sorters_df.as_matrix()
    # there has to be a better way to do this.
    for i in range(sorters.shape[0]):
        weights[i, :] = weights[i, sorters[i, :]]

    # Define topic names (e.g., topic_000)
    index = dt_df[1]
    weights_df = pd.DataFrame(columns=topic_names, data=weights, index=index)
    weights_df.index.name = 'pmid'

    # topic_keys: Top [n_words] words for each topic.
    keys_df = pd.read_csv(os.path.join(tempdir, 'topic_keys.txt'),
                          delimiter='\t', header=None, index_col=0)

    # Second column is a list of the terms.
    keys_df = keys_df[[2]]
    keys_df.rename(columns={2: 'terms'}, inplace=True)
    keys_df.index = topic_names
    keys_df.index.name = 'topic'

    # Remove all temporary files (abstract files, model, and outputs).
    shutil.rmtree(tempdir)

    # Return article topic weights and topic keys.
    return weights_df, keys_df


"""Various statistical helper functions"""

import warnings
import numpy as np
from scipy import special

def pearson(x, y):
    """ Correlates row vector x with each row vector in 2D array y. """
    data = np.vstack((x, y))
    ms = data.mean(axis=1)[(slice(None, None, None), None)]
    datam = data - ms
    datass = np.sqrt(np.sum(datam**2, axis=1))
    temp = np.dot(datam[1:], datam[0].T)
    rs = temp / (datass[1:] * datass[0])
    return rs

def two_way(cells):
    """ Two-way chi-square test of independence.
    Takes a 3D array as input: N(voxels) x 2 x 2, where the last two dimensions
    are the contingency table for each of N voxels. Returns an array of
    p-values.
    """
    # Mute divide-by-zero warning for bad voxels since we account for that
    # later
    warnings.simplefilter("ignore", RuntimeWarning)

    cells = cells.astype('float64')  # Make sure we don't overflow
    total = np.apply_over_axes(np.sum, cells, [1, 2]).ravel()
    chi_sq = np.zeros(cells.shape, dtype='float64')
    for i in range(2):
        for j in range(2):
            exp = np.sum(cells[:, i, :], 1).ravel() * \
                np.sum(cells[:, :, j], 1).ravel() / total
            bad_vox = np.where(exp == 0)[0]
            chi_sq[:, i, j] = (cells[:, i, j] - exp) ** 2 / exp
            chi_sq[bad_vox, i, j] = 1.0  # Set p-value for invalid voxels to 1
    chi_sq = np.apply_over_axes(np.sum, chi_sq, [1, 2]).ravel()
    return special.chdtrc(1, chi_sq)

def one_way(data, n):
    """ One-way chi-square test of independence.
    Takes a 1D array as input and compares activation at each voxel to
    proportion expected under a uniform distribution throughout the array. Note
    that if you're testing activation with this, make sure that only valid
    voxels (e.g., in-mask gray matter voxels) are included in the array, or
    results won't make any sense!
    """
    term = data.astype('float64')
    no_term = n - term
    t_exp = np.mean(term, 0)
    t_exp = np.array([t_exp, ] * data.shape[0])
    nt_exp = n - t_exp
    t_mss = (term - t_exp) ** 2 / t_exp
    nt_mss = (no_term - nt_exp) ** 2 / nt_exp
    chi2 = t_mss + nt_mss
    return special.chdtrc(1, chi2)

def fdr(p, q=.05):
    """ Determine FDR threshold given a p value array and desired false
    discovery rate q. """
    s = np.sort(p)
    nvox = p.shape[0]
    null = np.array(range(1, nvox + 1), dtype='float') * q / nvox
    below = np.where(s <= null)[0]
    return s[max(below)] if len(below) else -1


""" Base classes for representing and manipulating data. """

import logging
import re
import os
import sys
import numpy as np
import pandas as pd
from scipy import sparse
from neurosynth.base import mask, imageutils, transformations
from neurosynth.base import lexparser as lp
from neurosynth.utils import deprecated

# For Python 2/3 compatibility
from six import string_types
from functools import reduce
try:
    from urllib2 import urlopen
except ImportError:
    from urllib.request import urlopen
try:
    import cPickle as pickle
except:
    import pickle

logger = logging.getLogger('neurosynth.dataset')

def download(path='.', url=None, unpack=False):
    """ Download the latest data files.
    Args:
        path (str): Location to save the retrieved data files. Defaults to
            current directory.
        unpack (bool): If True, unzips the data file post-download.
    """

    if url is None:
        url = 'https://github.com/neurosynth/neurosynth-data/blob/master/current_data.tar.gz?raw=true'
    if os.path.exists(path) and os.path.isdir(path):
        basename = os.path.basename(url).split('?')[0]
        filename = os.path.join(path, basename)
    else:
        filename = path

    f = open(filename, 'wb')

    u = urlopen(url)
    file_size = int(u.headers["Content-Length"][0])
    print("Downloading the latest Neurosynth files: {0} bytes: {1}".format(
        url, file_size))

    bytes_dl = 0
    block_size = 8192
    while True:
        buffer = u.read(block_size)
        if not buffer:
            break
        bytes_dl += len(buffer)
        f.write(buffer)
        p = float(bytes_dl) / file_size
        status = r"{0}  [{1:.2%}]".format(bytes_dl, p)
        status = status + chr(8) * (len(status) + 1)
        sys.stdout.write(status)

    f.close()

    if unpack:
        import tarfile
        tarfile.open(filename, 'r:gz').extractall(os.path.dirname(filename))

def download_abstracts(dataset, path='.', email=None, out_file=None):
    """ Download the abstracts for a dataset/list of pmids
    """
    try:
        from Bio import Entrez, Medline
    except:
        raise Exception(
            'Module biopython is required for downloading abstracts from PubMed.')

    if email is None:
        raise Exception('No email address provided.')
    Entrez.email = email

    if isinstance(dataset, Dataset):
        pmids = dataset.image_table.ids.astype(str).tolist()
    elif isinstance(dataset, list):
        pmids = [str(pmid) for pmid in dataset]
    else:
        raise Exception(
            'Dataset type not recognized: {0}'.format(type(dataset)))

    records = []
    # PubMed only allows you to search ~1000 at a time. I chose 900 to be safe.
    chunks = [pmids[x: x + 900] for x in range(0, len(pmids), 900)]
    for chunk in chunks:
        h = Entrez.efetch(db='pubmed', id=chunk, rettype='medline',
                          retmode='text')
        records += list(Medline.parse(h))

    # Pull data for studies with abstracts
    data = [[study['PMID'], study['AB']]
            for study in records if study.get('AB', None)]
    df = pd.DataFrame(columns=['pmid', 'abstract'], data=data)
    if out_file is not None:
        df.to_csv(os.path.join(os.path.abspath(path), out_file), index=False)
    return df

class Dataset(object):

    """ Base Dataset class.

    The core data-representing object in Neurosynth. Internally stores
    information about both reported activations and tagged features. Provides a
    variety of methods for manipulating and retrieving various kinds of data.

    The Dataset is typically initialized by passing in a database file as the
    first argument. At minimum, the input file must contain tab-delimited
    columns named x, y, z, id, and space (case-insensitive). The x/y/z columns
    indicate the coordinates of the activation center or peak, the id column is
    used to group multiple activations from a single study. Typically the id
    should be a uniquely identifying field accessible to others, e.g., a PubMed
    ID in the case of entire articles. The space column indicates the nominal
    atlas used to produce each activation. Currently all values except 'TAL'
    (Talairach) will be ignored. If space == TAL and the transform argument is
    True, all activations reported in Talairach space will be converted to MNI
    space using the Lancaster et al. transform.

    Args:
        filename (str): The name of a database file containing a list of
            activations.
        feature_filename (str): An optional filename to construct a
            FeatureTable from.
        masker (str): An optional Nifti/Analyze image name defining the space
            to use for all operations. If no image is passed, defaults to the
            MNI152 2 mm template packaged with FSL.
        r (int): An optional integer specifying the radius of the smoothing
            kernel, in mm. Defaults to 6 mm.
        transform (bool, dict): Optional argument specifying how to handle
            transformation between coordinates reported in different
            stereotactic spaces. When True (default), activations in Talairach
            (T88) space will be converted to MNI space using the Lancaster et
            al (2007) transform; no other transformations will be applied. When
            False, no transformation will be applied. Alternatively, the user
            can pass their own dictionary of named transformations to apply, in
            which case each activation will be checked against the dictionary
            as it is read in and the specified transformation will be applied
            if found (for further details, see transformations.Transformer).
        target (str): The name of the target space within which activation
            coordinates are represented. By default, MNI.
        kwargs (dict): Additional optional arguments passed to add_features().
    """

    def __init__(
            self, filename, feature_filename=None, masker=None, r=6,
            transform=True, target='MNI', **kwargs):

        # Instance properties
        self.r = r

        # Set up transformations between different image spaces
        if transform:
            if not isinstance(transform, dict):
                transform = {'T88': transformations.t88_to_mni(),
                             'TAL': transformations.t88_to_mni()
                             }
            self.transformer = transformations.Transformer(transform, target)
        else:
            self.transformer = None

        # Load and process activation data
        self.activations = self._load_activations(filename)

        # Load the volume into a new Masker
        if masker is None:
            resource_dir = os.path.join(os.path.dirname(__file__),
                                        os.path.pardir,
                                        'resources')
            masker = os.path.join(
                resource_dir, 'MNI152_T1_2mm_brain.nii.gz')
        self.masker = mask.Masker(masker)

        # Create supporting tables for images and features
        self.create_image_table()
        if feature_filename is not None:
            self.add_features(feature_filename, **kwargs)

    def _load_activations(self, filename):
        """ Load activation data from a text file.

        Args:
            filename (str): a string pointing to the location of the txt file
                to read from.
        """
        logger.info("Loading activation data from %s..." % filename)

        activations = pd.read_csv(filename, sep='\t')
        activations.columns = [col.lower()
                               for col in list(activations.columns)]

        # Make sure all mandatory columns exist
        mc = ['x', 'y', 'z', 'id', 'space']
        if (set(mc) - set(list(activations.columns))):
            logger.error(
                "At least one of mandatory columns (x, y, z, id, and space) "
                "is missing from input file.")
            return

        # Transform to target space where needed
        spaces = activations['space'].unique()
        xyz = activations[['x', 'y', 'z']].values
        for s in spaces:
            if s != self.transformer.target:
                inds = activations['space'] == s
                xyz[inds] = self.transformer.apply(s, xyz[inds])
        activations[['x', 'y', 'z']] = xyz

        # xyz --> ijk
        ijk = pd.DataFrame(
            transformations.xyz_to_mat(xyz), columns=['i', 'j', 'k'])
        activations = pd.concat([activations, ijk], axis=1)
        return activations

    def create_image_table(self, r=None):
        """ Create and store a new ImageTable instance based on the current
        Dataset. Will generally be called privately, but may be useful as a
        convenience method in cases where the user wants to re-generate the
        table with a new smoothing kernel of different radius.

        Args:
            r (int): An optional integer indicating the radius of the smoothing
                kernel. By default, this is None, which will keep whatever
                value is currently set in the Dataset instance.
        """
        logger.info("Creating image table...")
        if r is not None:
            self.r = r
        self.image_table = ImageTable(self)

    def get_studies(self, features=None, expression=None, mask=None,
                    peaks=None, frequency_threshold=0.001,
                    activation_threshold=0.0, func=np.sum, return_type='ids',
                    r=6
                    ):
        """ Get IDs or data for studies that meet specific criteria.

        If multiple criteria are passed, the set intersection is returned. For
        example, passing expression='emotion' and mask='my_mask.nii.gz' would
        return only those studies that are associated with emotion AND report
        activation within the voxels indicated in the passed image.

        Args:
            ids (list): A list of IDs of studies to retrieve.
            features (list or str): The name of a feature, or a list of
                features, to use for selecting studies.
            expression (str): A string expression to pass to the PEG for study
                retrieval.
            mask: the mask image (see Masker documentation for valid data
                types).
            peaks (ndarray or list): Either an n x 3 numpy array, or a list of
                lists or tuples (e.g., [(-10, 22, 14)]) specifying the world
                (x/y/z) coordinates of the target location(s).
            frequency_threshold (float): For feature-based or expression-based
                selection, the threshold for selecting studies--i.e., the
                cut-off for a study to be included. Must be a float in range
                [0, 1].
            activation_threshold (int or float): For mask-based selection,
                threshold for a study to be included based on amount of
                activation displayed. If an integer, represents the absolute
                number of voxels that must be active within the mask in order
                for a study to be selected. If a float, it represents the
                proportion of voxels that must be active.
            func (Callable): The function to use when aggregating over the list
                of features. See documentation in FeatureTable.get_ids() for a
                full explanation. Only used for feature- or expression-based
                selection.
            return_type (str): A string specifying what data to return. Valid
                options are:
                'ids': returns a list of IDs of selected studies.
                'images': returns a voxel x study matrix of data for all
                selected studies.
                'weights': returns a dict where the keys are study IDs and the
                values are the computed weights. Only valid when performing
                feature-based selection.
            r (int): For peak-based selection, the distance cut-off (in mm)
                for inclusion (i.e., only studies with one or more activations
                within r mm of one of the passed foci will be returned).

        Returns:
            When return_type is 'ids' (default), returns a list of IDs of the
            selected studies. When return_type is 'data', returns a 2D numpy
            array, with voxels in rows and studies in columns. When return_type
            is 'weights' (valid only for expression-based selection), returns
            a dict, where the keys are study IDs, and the values are the
            computed weights.

        Examples
        --------
        Select all studies tagged with the feature 'emotion':

            >>> ids = dataset.get_studies(features='emotion')

        Select all studies that activate at least 20% of voxels in an amygdala
        mask, and retrieve activation data rather than IDs:

            >>> data = dataset.get_studies(mask='amygdala_mask.nii.gz',
                threshold=0.2, return_type='images')

        Select studies that report at least one activation within 12 mm of at
        least one of three specific foci:

            >>> ids = dataset.get_studies(peaks=[[12, -20, 30], [-26, 22, 22],
                                                [0, 36, -20]], r=12)

        """
        results = []

        # Feature-based selection
        if features is not None:
            # Need to handle weights as a special case, because we can't
            # retrieve the weights later using just the IDs.
            if return_type == 'weights':
                if expression is not None or mask is not None or \
                        peaks is not None:
                    raise ValueError(
                        "return_type cannot be 'weights' when feature-based "
                        "search is used in conjunction with other search "
                        "modes.")
                return self.feature_table.get_ids(
                    features, frequency_threshold, func, get_weights=True)
            else:
                results.append(self.feature_table.get_ids(
                    features, frequency_threshold, func))

        # Logical expression-based selection
        if expression is not None:
            _ids = self.feature_table.get_ids_by_expression(
                expression, frequency_threshold, func)
            results.append(list(_ids))

        # Mask-based selection
        if mask is not None:
            mask = self.masker.mask(mask, in_global_mask=True).astype(bool)
            num_vox = np.sum(mask)
            prop_mask_active = self.image_table.data.T.dot(mask).astype(float)
            if isinstance(activation_threshold, float):
                prop_mask_active /= num_vox
            indices = np.where(prop_mask_active > activation_threshold)[0]
            results.append([self.image_table.ids[ind] for ind in indices])

        # Peak-based selection
        if peaks is not None:
            r = float(r)
            found = set()
            for p in peaks:
                xyz = np.array(p, dtype=float)
                x = self.activations['x']
                y = self.activations['y']
                z = self.activations['z']
                dists = np.sqrt(np.square(x - xyz[0]) + np.square(y - xyz[1]) +
                                np.square(z - xyz[2]))
                inds = np.where((dists > 5.5) & (dists < 6.5))[0]
                tmp = dists[inds]
                found |= set(self.activations[dists <= r]['id'].unique())
            results.append(found)

        # Get intersection of all sets
        ids = list(reduce(lambda x, y: set(x) & set(y), results))

        if return_type == 'ids':
            return ids
        elif return_type == 'data':
            return self.get_image_data(ids)

    def add_features(self, features, append=True, merge='outer',
                     duplicates='ignore', min_studies=0.0, threshold=0.001):
        """ Construct a new FeatureTable from file.

        Args:
            features: Feature data to add. Can be:
                (a) A text file containing the feature data, where each row is
                a study in the database, with features in columns. The first
                column must contain the IDs of the studies to match up with the
                image data.
                (b) A pandas DataFrame, where studies are in rows, features are
                in columns, and the index provides the study IDs.
            append (bool): If True, adds new features to existing ones
                incrementally. If False, replaces old features.
            merge, duplicates, min_studies, threshold: Additional arguments
                passed to FeatureTable.add_features().
         """
        if (not append) or not hasattr(self, 'feature_table'):
            self.feature_table = FeatureTable(self)

        self.feature_table.add_features(features, merge=merge,
                                        duplicates=duplicates,
                                        min_studies=min_studies,
                                        threshold=threshold)

    def get_image_data(self, ids=None, voxels=None, dense=True):
        """ A convenience wrapper for ImageTable.get_image_data().

        Args:
            ids (list, array): A list or 1D numpy array of study ids to
                return. If None, returns data for all studies.
            voxels (list, array): A list or 1D numpy array of voxel indices
                (i.e., rows) to return. If None, returns data for all voxels.
        """
        return self.image_table.get_image_data(ids, voxels=voxels, dense=dense)

    def get_feature_data(self, ids=None, **kwargs):
        """ A convenience wrapper for FeatureTable.get_image_data(). """
        return self.feature_table.get_feature_data(ids, **kwargs)

    def get_feature_names(self, features=None):
        """ Returns names of features. If features is None, returns all
        features. Otherwise assumes the user is trying to find the order of the
        features.  """
        if features:
            return self.feature_table.get_ordered_names(features)
        else:
            return self.feature_table.feature_names

    def get_feature_counts(self, threshold=0.001):
        """ Returns a dictionary, where the keys are the feature names
        and the values are the number of studies tagged with the feature. """
        counts = np.sum(self.get_feature_data() >= threshold, 0)
        return dict(zip(self.get_feature_names(), list(counts)))

    @classmethod
    def load(cls, filename):
        """ Load a pickled Dataset instance from file. """
        try:
            dataset = pickle.load(open(filename, 'rb'))
        except UnicodeDecodeError:
            # Need to try this for python3
            dataset = pickle.load(open(filename, 'rb'), encoding='latin')

        if hasattr(dataset, 'feature_table'):
            dataset.feature_table._csr_to_sdf()
        return dataset

    def save(self, filename):
        """ Pickle the Dataset instance to the provided file.
        """
        if hasattr(self, 'feature_table'):
            self.feature_table._sdf_to_csr()

        pickle.dump(self, open(filename, 'wb'), -1)

        if hasattr(self, 'feature_table'):
            self.feature_table._csr_to_sdf()

class ImageTable(object):

    """ Represents image data from multiple studies in an accessible form.

    Args:
        dataset (Dataset): Dataset instance to pull inputs from.
        r (int): The radius of the sphere used for smoothing (default = 6 mm).
        use_sparse (bool): Flag indicating whether or not to represent the data
            as a sparse array (generally this should be left to True, as these
            data are quite sparse and computation is often considerably slower
            in dense form.)
    """

    def __init__(self, dataset, r=6, use_sparse=True):
        activations, masker, r = dataset.activations, dataset.masker, dataset.r
        for var in [activations, masker, r]:
            assert var is not None
        self.ids = activations['id'].unique()
        self.masker = masker
        self.r = r

        n_studies = len(self.ids)
        data_shape = (self.masker.n_vox_in_vol, n_studies)
        if use_sparse:
            # Fancy indexing assignment is not supported for sparse matrices,
            # so let's keep lists of values and their indices (rows, cols) to
            # later construct the csr_matrix.
            vals, rows, cols = [], [], []
        else:
            self.data = np.zeros(data_shape, dtype=int)

        logger.info("Mapping %d studies into image space..." % (n_studies,))
        for i, (name, data) in enumerate(activations.groupby('id')):
            logger.debug("%s/%s..." % (str(i + 1), str(n_studies)))
            img = imageutils.map_peaks_to_image(
                data[['i', 'j', 'k']].values, r=r,
                header=self.masker.get_header())
            img_masked = self.masker.mask(img)
            if use_sparse:
                nz = np.nonzero(img_masked)
                assert(len(nz) == 1)
                vals += list(img_masked[nz])
                rows += list(nz[0])
                cols += [i] * len(nz[0])
            else:
                self.data[:, i] = img_masked

        if use_sparse:
            self.data = sparse.csr_matrix((
                vals, (rows, cols)), shape=data_shape)

    def get_image_data(self, ids=None, voxels=None, dense=True):
        """ Slices and returns a subset of image data.

        Args:
            ids (list, array): A list or 1D numpy array of study ids to
                return. If None, returns data for all studies.
            voxels (list, array): A list or 1D numpy array of voxel indices
                (i.e., rows) to return. If None, returns data for all voxels.
            dense (bool): Optional boolean. When True (default), convert the
                result to a dense array before returning. When False, keep as
                sparse matrix.

        Returns:
          A 2D numpy array with voxels in rows and studies in columns.
        """
        if dense and ids is None and voxels is None:
            logger.warning(
                "Warning: get_image_data() is being called without specifying "
                "a subset of studies or voxels to retrieve. This may result in"
                " a very large amount of data (several GB) being read into "
                "memory. If you experience any problems, consider returning a "
                "sparse matrix by passing dense=False, or pass in a list of "
                "ids of voxels to retrieve only a portion of the data.")

        result = self.data
        if ids is not None:
            idxs = np.where(np.in1d(np.array(self.ids), np.array(ids)))[0]
            result = result[:, idxs]
        if voxels is not None:
            result = result[voxels, :]
        return result.toarray() if dense else result

    def trim(self, ids):
        """ Trim ImageTable to keep only the passed studies. This is a
        convenience method, and should generally be avoided in favor of
        non-destructive alternatives that don't require slicing (e.g.,
            matrix multiplication). """
        self.data = self.get_image_data(ids, dense=False)  # .tocoo()
        idxs = np.where(np.in1d(np.array(self.ids), np.array(ids)))[0]
        self.ids = [self.ids[i] for i in idxs]

    def save_images_to_file(self, ids, outroot='./'):
        """ Reconstructs vectorized images corresponding to the specified
        study ids and saves them to file, prepending with the outroot
        (default: current directory). """
        pass

    def save(self, filename):
        pickle.dump(self, open(filename, 'wb'), -1)

class FeatureTable(object):

    """ A FeatureTable instance stores a matrix of studies x features,
    along with associated manipulation methods. """

    def __init__(self, dataset):
        """ Initialize a new FeatureTable. Takes as input a parent DataSet
        instance and feature data (if provided). """
        self.dataset = dataset
        self.data = pd.DataFrame()

    def add_features(self, features, merge='outer', duplicates='ignore',
                     min_studies=0, threshold=0.0001):
        """ Add new features to FeatureTable.
        Args:
            features (str, DataFrame): A filename to load data from, or a
                pandas DataFrame. In either case, studies are in rows and
                features are in columns. Values in cells reflect the weight of
                the intersecting feature for the intersecting study. Feature
                names and study IDs should be included as the first column
                and first row, respectively.
            merge (str): The merge strategy to use when merging new features
                with old. This is passed to pandas.merge, so can be 'left',
                'right', 'outer', or 'inner'. Defaults to outer (i.e., all data
                in both new and old will be kept, and missing values will be
                assigned zeros.)
            duplicates (str): string indicating how to handle features whose
                name matches an existing feature. Valid options:
                    'ignore' (default): ignores the new feature, keeps old data
                    'replace': replace the old feature's data with the new data
                    'merge': keeps both features, renaming them so they're
                        different
            min_studies (int): minimum number of studies that pass threshold in
                order to add feature.
            threshold (float): minimum frequency threshold each study must
                exceed in order to count towards min_studies.
        """
        if isinstance(features, string_types):
            if not os.path.exists(features):
                raise ValueError("%s cannot be found." % features)
            try:
                features = pd.read_csv(features, sep='\t', index_col=0)
            except Exception as e:
                logger.error("%s cannot be parsed: %s" % (features, e))

        if min_studies:
            valid = np.where(
                (features.values >= threshold).sum(0) >= min_studies)[0]
            features = features.iloc[:, valid]

        # Warn user if no/few IDs match between the FeatureTable and the
        # Dataset. This most commonly happens because older database.txt files
        # used doi's as IDs whereas we now use PMIDs throughout.
        n_studies = len(features)
        n_common_ids = len(
            set(features.index) & set(self.dataset.image_table.ids))
        if float(n_common_ids) / n_studies < 0.01:  # Minimum 1% overlap
            msg = "Only %d" % n_common_ids if n_common_ids else "None of the"
            logger.warning(
                msg + " studies in the feature file matched studies currently "
                "the Dataset. The most likely cause for this is that you're "
                "pairing a newer feature set with an older, incompatible "
                "database file. You may want to try regenerating the Dataset "
                "instance from a newer database file that uses PMIDs rather "
                "than doi's as the study identifiers in the first column.")

        old_data = self.data.to_dense()
        # Handle features with duplicate names
        common_features = list(set(old_data.columns) & set(features.columns))
        if duplicates == 'ignore':
            features = features.drop(common_features, axis=1)
        elif duplicates == 'replace':
            old_data = old_data.drop(common_features, axis=1)

        data = old_data.merge(
            features, how=merge, left_index=True, right_index=True)
        self.data = data.fillna(0.0).to_sparse()

    @property
    def feature_names(self):
        return list(self.data.columns)

    def get_feature_data(self, ids=None, features=None, dense=True):
        """ Slices and returns a subset of feature data.

        Args:
            ids (list, array): A list or 1D numpy array of study ids to
                return rows for. If None, returns data for all studies
                (i.e., all rows in array).
            features (list, array): A list or 1D numpy array of named features
                to return. If None, returns data for all features (i.e., all
                columns in array).
            dense (bool): Optional boolean. When True (default), convert the
                result to a dense array before returning. When False, keep as
                sparse matrix. Note that if ids is not None, the returned array
                will always be dense.
        Returns:
          A pandas DataFrame with study IDs in rows and features incolumns.
        """
        result = self.data

        if ids is not None:
            result = result.ix[ids]

        if features is not None:
            result = result.ix[:, features]

        return result.to_dense() if dense else result

    def get_ordered_names(self, features):
        """ Given a list of features, returns features in order that they
        appear in database.

        Args:
            features (list): A list or 1D numpy array of named features to
            return.

        Returns:
            A list of features in order they appear in database.
        """

        idxs = np.where(
            np.in1d(self.data.columns.values, np.array(features)))[0]
        return list(self.data.columns[idxs].values)

    def get_ids(self, features, threshold=0.0, func=np.sum, get_weights=False):
        """ Returns a list of all studies in the table that meet the desired
        feature-based criteria.

        Will most commonly be used to retrieve studies that use one or more
        features with some minimum frequency; e.g.,:

            get_ids(['fear', 'anxiety'], threshold=0.001)

        Args:
            features (lists): a list of feature names to search on.
            threshold (float): optional float indicating threshold features
                must pass to be included.
            func (Callable): any numpy function to use for thresholding
                (default: sum). The function will be applied to the list of
                features and the result compared to the threshold. This can be
                used to change the meaning of the query in powerful ways. E.g,:
                    max: any of the features have to pass threshold
                        (i.e., max > thresh)
                    min: all features must each individually pass threshold
                        (i.e., min > thresh)
                    sum: the summed weight of all features must pass threshold
                        (i.e., sum > thresh)
            get_weights (bool): if True, returns a dict with ids => weights.

        Returns:
            When get_weights is false (default), returns a list of study
                names. When true, returns a dict, with study names as keys
                and feature weights as values.
        """
        if isinstance(features, str):
            features = [features]
        features = self.search_features(features)  # Expand wild cards
        feature_weights = self.data.ix[:, features]
        weights = feature_weights.apply(func, 1)
        above_thresh = weights[weights >= threshold]
        # ids_to_keep = self.ids[above_thresh]
        return above_thresh if get_weights else list(above_thresh.index)

    def search_features(self, search):
        ''' Returns all features that match any of the elements in the input
        list.

        Args:
            search (str, list): A string or list of strings defining the query.

        Returns:
            A list of matching feature names.
        '''
        if isinstance(search, string_types):
            search = [search]
        search = [s.replace('*', '.*') for s in search]
        cols = list(self.data.columns)
        results = []
        for s in search:
            results.extend([f for f in cols if re.match(s + '$', f)])
        return list(set(results))

    def get_ids_by_expression(self, expression, threshold=0.001, func=np.sum):
        """ Use a PEG to parse expression and return study IDs."""
        lexer = lp.Lexer()
        lexer.build()
        parser = lp.Parser(
            lexer, self.dataset, threshold=threshold, func=func)
        parser.build()
        return parser.parse(expression).keys().values

    def get_features_by_ids(self, ids=None, threshold=0.0001, func=np.mean,
                            get_weights=False):
        ''' Returns features for which the mean loading across all specified
        studies (in ids) is >= threshold. '''
        weights = self.data.ix[ids].apply(func, 0)
        above_thresh = weights[weights >= threshold]
        return above_thresh if get_weights else list(above_thresh.index)

    def _sdf_to_csr(self):
        """ Convert FeatureTable to SciPy CSR matrix. """
        data = self.data.to_dense()
        self.data = {
            'columns': list(data.columns),
            'index': list(data.index),
            'values': sparse.csr_matrix(data.values)
        }

    def _csr_to_sdf(self):
        """ Inverse of _sdf_to_csr(). """
        self.data = pd.DataFrame(self.data['values'].todense(),
                                 index=self.data['index'],
                                 columns=self.data['columns']).to_sparse()


import warnings

def deprecated(*args):
    """ Deprecation warning decorator. Takes optional deprecation message,
    otherwise will use a generic warning. """
    def wrap(func):
        def wrapped_func(*args, **kwargs):
            warnings.warn(msg, category=DeprecationWarning)
            return func(*args, **kwargs)
        return wrapped_func

    if len(args) == 1 and callable(args[0]):
        msg = "Function '%s' will be deprecated in future versions of " \
            "Neurosynth." % args[0].__name__
        return wrap(args[0])
    else:
        msg = args[0]
        return wrap


from copy import deepcopy
import numpy as np
from six import string_types
from sklearn import decomposition as sk_decomp
from sklearn import cluster as sk_cluster
from sklearn.metrics import pairwise_distances
from os.path import exists, join
from os import makedirs
from nibabel import nifti1
from neurosynth.analysis import meta

class Clusterable(object):

    '''
    Args:
        dataset: The Dataset instance to extract data from.
        mask: A mask defining the voxels to cluster. Can be a filename,
            nibabel image, or numpy array (see Masker.mask() for details).
        features (str or list): Optional string or list of strings specifying
            any feature names to use for study selection. E.g., passing
            ['emotion', 'reward'] would retain for analysis only those studies
            associated with the features emotion or reward at a frequency
            greater than feature_threshold.
        feature_threshold (float): The threshold to use when selecting studies
            on the basis of features.
        min_voxels (int): Minimum number of active voxels a study
            must report in order to be retained in the dataset. By default,
            all studies are used.
        min_studies (int): Minimum number of studies a voxel must be
            active in in order to be retained in analysis. By default, all
            voxels are used.
    '''

    def __init__(self, dataset, mask=None, features=None,
                 feature_threshold=None, min_voxels=None, min_studies=None):

        self.dataset = dataset
        self.masker = deepcopy(dataset.masker)

        # Condition study inclusion on specific features
        if features is not None:
            ids = dataset.get_studies(features=features,
                                      frequency_threshold=feature_threshold)
            data = dataset.get_image_data(ids, dense=False)
        else:
            data = dataset.image_table.data

        # Trim data based on minimum number of voxels or studies
        if min_studies is not None:
            av = self.masker.unmask(
                data.sum(1) >= min_studies, output='vector')
            self.masker.add({'voxels': av})

        if min_voxels is not None:
            data = data[:, np.array(data.sum(0) >= min_voxels).squeeze()]

        if mask is not None:
            self.masker.add({'roi': mask})

        self.data = data[self.masker.get_mask(['voxels', 'roi']), :].toarray()

    def transform(self, transformer, transpose=False):
        ''' Apply a transformation to the Clusterable instance. Accepts any
        scikit-learn-style class that implements a fit_transform() method. '''
        data = self.data.T if transpose else self.data
        data = transformer.fit_transform(data)
        self.data = data.T if transpose else data
        return self

def magic(dataset, method='coactivation', roi_mask=None,
          coactivation_mask=None, features=None, feature_threshold=0.05,
          min_voxels_per_study=None, min_studies_per_voxel=None,
          reduce_reference='pca', n_components=100,
          distance_metric='correlation', clustering_algorithm='kmeans',
          n_clusters=5, clustering_kwargs={}, output_dir=None, filename=None,
          coactivation_images=False, coactivation_threshold=0.1):
    ''' Execute a full clustering analysis pipeline.
    Args:
        dataset: a Dataset instance to extract all data from.
        method (str): the overall clustering approach to use. Valid options:
            'coactivation' (default): Clusters voxel within the ROI mask based
                on shared pattern of coactivation with the rest of the brain.
            'studies': Treat each study as a feature in an n-dimensional space.
                I.e., voxels will be assigned to the same cluster if they tend
                to be co-reported in similar studies.
            'features': Voxels will be assigned to the same cluster if they
                tend to have similar feature vectors (i.e., the studies that
                activate those voxels tend to use similar terms).
        roi_mask: A string, nibabel image, or numpy array providing an
            inclusion mask of voxels to cluster. If None, the default mask
            in the Dataset instance is used (typically, all in-brain voxels).
        coactivation_mask: If method='coactivation', this mask defines the
            voxels to use when generating the pairwise distance matrix. For
            example, if a PFC mask is passed, all voxels in the roi_mask will
            be clustered based on how similar their patterns of coactivation
            with PFC voxels are. Can be a str, nibabel image, or numpy array.
        features (str or list): Optional string or list of strings specifying
            any feature names to use for study selection. E.g., passing
            ['emotion', 'reward'] would retain for analysis only those studies
            associated with the features emotion or reward at a frequency
            greater than feature_threshold.
        feature_threshold (float): The threshold to use when selecting studies
            on the basis of features.
        min_voxels_per_study (int): Minimum number of active voxels a study
            must report in order to be retained in the dataset. By default,
            all studies are used.
        min_studies_per_voxel (int): Minimum number of studies a voxel must be
            active in in order to be retained in analysis. By default, all
            voxels are used.
        reduce_reference (str, scikit-learn object or None): The dimensionality
            reduction algorithm to apply to the feature space prior to the
            computation of pairwise distances. If a string is passed (either
            'pca' or 'ica'), n_components must be specified. If None, no
            dimensionality reduction will be applied. Otherwise, must be a
            scikit-learn-style object that exposes a transform() method.
        n_components (int): Number of components to extract during the
            dimensionality reduction step. Only used if reduce_reference is
            a string.
        distance_metric (str): The distance metric to use when computing
            pairwise distances on the to-be-clustered voxels. Can be any of the
            metrics supported by sklearn.metrics.pairwise_distances.
        clustering_algorithm (str or scikit-learn object): the clustering
            algorithm to use. If a string, must be one of 'kmeans' or 'minik'.
            Otherwise, any sklearn class that exposes a fit_predict() method.
        n_clusters (int): If clustering_algorithm is a string, the number of
            clusters to extract.
        clustering_kwargs (dict): Additional keywords to pass to the clustering
            object.
        output_dir (str): The directory to write results to. If None (default),
            returns the cluster label image rather than saving to disk.
        filename (str): Name of cluster label image file. Defaults to
            cluster_labels_k{k}.nii.gz, where k is the number of clusters.
        coactivation_images (bool): If True, saves a meta-analytic coactivation
            map for every ROI in the resulting cluster map.
        coactivation_threshold (float or int): If coactivation_images is True,
            this is the threshold used to define whether or not a study is
            considered to activation within a cluster ROI. Integer values are
            interpreted as minimum number of voxels within the ROI; floats
            are interpreted as the proportion of voxels. Defaults to 0.1 (i.e.,
            10% of all voxels within ROI must be active).
    '''

    roi = Clusterable(dataset, roi_mask, min_voxels=min_voxels_per_study,
                      min_studies=min_studies_per_voxel, features=features,
                      feature_threshold=feature_threshold)

    if method == 'coactivation':
        reference = Clusterable(dataset, coactivation_mask,
                                min_voxels=min_voxels_per_study,
                                min_studies=min_studies_per_voxel,
                                features=features,
                                feature_threshold=feature_threshold)
    elif method == 'features':
        reference = deepcopy(roi)
        feature_data = dataset.feature_table.data
        n_studies = len(feature_data)
        reference.data = reference.data.dot(feature_data.values) / n_studies
    elif method == 'studies':
        reference = roi

    if reduce_reference is not None:

        if isinstance(reduce_reference, string_types):

            # Number of components can't exceed feature count or cluster count
            n_feat = reference.data.shape[1]
            n_components = min(n_components, n_feat)

            reduce_reference = {
                'pca': sk_decomp.PCA,
                'ica': sk_decomp.FastICA
            }[reduce_reference](n_components)

        # For non-coactivation-based approaches, transpose the data matrix
        transpose = (method == 'coactivation')
        reference = reference.transform(reduce_reference, transpose=transpose)

    if method == 'coactivation':
        distances = pairwise_distances(roi.data, reference.data,
                                       metric=distance_metric)
    else:
        distances = reference.data

    # TODO: add additional clustering methods
    if isinstance(clustering_algorithm, string_types):
        clustering_algorithm = {
            'kmeans': sk_cluster.KMeans,
            'minik': sk_cluster.MiniBatchKMeans
        }[clustering_algorithm](n_clusters, **clustering_kwargs)

    labels = clustering_algorithm.fit_predict(distances) + 1.
    header = roi.masker.get_header()
    header['cal_max'] = labels.max()
    header['cal_min'] = labels.min()
    voxel_labels = roi.masker.unmask(labels)
    img = nifti1.Nifti1Image(voxel_labels, None, header)

    if output_dir is not None:
        if not exists(output_dir):
            makedirs(output_dir)
        if filename is None:
            filename = 'cluster_labels_k%d.nii.gz' % n_clusters
        outfile = join(output_dir, filename)
        img.to_filename(outfile)

        # Write coactivation images
        if coactivation_images:
            for l in np.unique(voxel_labels):
                roi_mask = np.copy(voxel_labels)
                roi_mask[roi_mask != l] = 0
                ids = dataset.get_studies(
                    mask=roi_mask, activation_threshold=coactivation_threshold)
                ma = meta.MetaAnalysis(dataset, ids)
                ma.save_results(output_dir=join(output_dir, 'coactivation'),
                                prefix='cluster_%d_coactivation' % l)
    else:
        return img