# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""
The rapidart module provides routines for artifact detection and region of
interest analysis.

These functions include:

  * ArtifactDetect: performs artifact detection on functional images

  * StimulusCorrelation: determines correlation between stimuli
    schedule and movement/intensity parameters

   Change directory to provide relative paths for doctests
   >>> import os
   >>> filepath = os.path.dirname( os.path.realpath( __file__ ) )
   >>> datadir = os.path.realpath(os.path.join(filepath, '../testing/data'))
   >>> os.chdir(datadir)
"""

import os
from copy import deepcopy
from warnings import warn

from nibabel import load, funcs, Nifti1Image
import numpy as np
from scipy import signal
import scipy.io as sio
from nipype.external import six

from ..interfaces.base import (BaseInterface, traits, InputMultiPath,
                                    OutputMultiPath, TraitedSpec, File,
                                    BaseInterfaceInputSpec, isdefined)
from ..utils.filemanip import filename_to_list, save_json, split_filename
from ..utils.misc import find_indices

from .. import logging, config
iflogger = logging.getLogger('interface')

def _get_affine_matrix(params, source):
    """Return affine matrix given a set of translation and rotation parameters

    params : np.array (upto 12 long) in native package format
    source : the package that generated the parameters
             supports SPM, AFNI, FSFAST, FSL, NIPY
    """
    if source == 'FSL':
        params = params[[3, 4, 5, 0, 1, 2]]
    elif source in ('AFNI', 'FSFAST'):
        params = params[np.asarray([4, 5, 3, 1, 2, 0]) + (len(params) > 6)]
        params[3:] = params[3:] * np.pi / 180.
    if source == 'NIPY':
        # nipy does not store typical euler angles, use nipy to convert
        from nipy.algorithms.registration import to_matrix44
        return to_matrix44(params)
    #process for FSL, SPM, AFNI and FSFAST
    rotfunc = lambda x: np.array([[np.cos(x), np.sin(x)],
                                  [-np.sin(x), np.cos(x)]])
    q = np.array([0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0])
    if len(params) < 12:
        params = np.hstack((params, q[len(params):]))
    params.shape = (len(params),)
    # Translation
    T = np.eye(4)
    T[0:3, -1] = params[0:3]
    # Rotation
    Rx = np.eye(4)
    Rx[1:3, 1:3] = rotfunc(params[3])
    Ry = np.eye(4)
    Ry[(0, 0, 2, 2), (0, 2, 0, 2)] = rotfunc(params[4]).ravel()
    Rz = np.eye(4)
    Rz[0:2, 0:2] = rotfunc(params[5])
    # Scaling
    S = np.eye(4)
    S[0:3, 0:3] = np.diag(params[6:9])
    # Shear
    Sh = np.eye(4)
    Sh[(0, 0, 1), (1, 2, 2)] = params[9:12]
    if source in ('AFNI', 'FSFAST'):
        return np.dot(T, np.dot(Ry, np.dot(Rx, np.dot(Rz, np.dot(S, Sh)))))
    return np.dot(T, np.dot(Rx, np.dot(Ry, np.dot(Rz, np.dot(S, Sh)))))

def _calc_norm(mc, use_differences, source, brain_pts=None):
    """Calculates the maximum overall displacement of the midpoints
    of the faces of a cube due to translation and rotation.

    Parameters
    ----------
    mc : motion parameter estimates
        [3 translation, 3 rotation (radians)]
    use_differences : boolean
    brain_pts : [4 x n_points] of coordinates

    Returns
    -------

    norm : at each time point
    displacement : euclidean distance (mm) of displacement at each coordinate

    """

    if brain_pts is None:
        respos = np.diag([70, 70, 75])
        resneg = np.diag([-70, -110, -45])
        all_pts = np.vstack((np.hstack((respos, resneg)), np.ones((1, 6))))
        displacement = None
    else:
        all_pts = brain_pts
    n_pts = all_pts.size - all_pts.shape[1]
    newpos = np.zeros((mc.shape[0], n_pts))
    if brain_pts is not None:
        displacement = np.zeros((mc.shape[0], n_pts / 3))
    for i in range(mc.shape[0]):
        affine = _get_affine_matrix(mc[i, :], source)
        newpos[i, :] = np.dot(affine,
                              all_pts)[0:3, :].ravel()
        if brain_pts is not None:
            displacement[i, :] = \
                  np.sqrt(np.sum(np.power(np.reshape(newpos[i, :],
                                                     (3, all_pts.shape[1])) -
                                          all_pts[0:3, :],
                                          2),
                                 axis=0))
    # np.savez('displacement.npz', newpos=newpos, pts=all_pts)
    normdata = np.zeros(mc.shape[0])
    if use_differences:
        newpos = np.concatenate((np.zeros((1, n_pts)),
                                 np.diff(newpos, n=1, axis=0)), axis=0)
        for i in range(newpos.shape[0]):
            normdata[i] = \
             np.max(np.sqrt(np.sum(np.reshape(np.power(np.abs(newpos[i, :]), 2),
                                              (3, all_pts.shape[1])), axis=0)))
    else:
        newpos = np.abs(signal.detrend(newpos, axis=0, type='constant'))
        normdata = np.sqrt(np.mean(np.power(newpos, 2), axis=1))
    return normdata, displacement

def _nanmean(a, axis=None):
    """Return the mean excluding items that are nan

    >>> a = [1, 2, np.nan]
    >>> _nanmean(a)
    1.5

    """
    if axis:
        return np.nansum(a, axis) / np.sum(1 - np.isnan(a), axis)
    else:
        return np.nansum(a) / np.sum(1 - np.isnan(a))

class ArtifactDetectInputSpec(BaseInterfaceInputSpec):
    realigned_files = InputMultiPath(File(exists=True),
                                desc="Names of realigned functional data files",
                                     mandatory=True)
    realignment_parameters = InputMultiPath(File(exists=True), mandatory=True,
                            desc=("Names of realignment parameters"
                                  "corresponding to the functional data files"))
    parameter_source = traits.Enum("SPM", "FSL", "AFNI", "NiPy", "FSFAST",
                                   desc="Source of movement parameters",
                                   mandatory=True)
    use_differences = traits.ListBool([True, False], minlen=2, maxlen=2,
                                      usedefault=True,
            desc=("Use differences between successive motion (first element)"
                  "and intensity paramter (second element) estimates in order"
                  "to determine outliers.  (default is [True, False])"))
    use_norm = traits.Bool(True, requires=['norm_threshold'],
                           desc=("Uses a composite of the motion parameters in "
                                 "order to determine outliers."),
                           usedefault=True)
    norm_threshold = traits.Float(desc=("Threshold to use to detect motion-rela"
                                        "ted outliers when composite motion is "
                                        "being used"), mandatory=True,
                                  xor=['rotation_threshold',
                                       'translation_threshold'])
    rotation_threshold = traits.Float(mandatory=True, xor=['norm_threshold'],
            desc=("Threshold (in radians) to use to detect rotation-related "
                  "outliers"))
    translation_threshold = traits.Float(mandatory=True, xor=['norm_threshold'],
            desc=("Threshold (in mm) to use to detect translation-related "
                 "outliers"))
    zintensity_threshold = traits.Float(mandatory=True,
            desc=("Intensity Z-threshold use to detection images that deviate "
                  "from the mean"))
    mask_type = traits.Enum('spm_global', 'file', 'thresh',
            desc=("Type of mask that should be used to mask the functional "
                  "data. *spm_global* uses an spm_global like calculation to "
                  "determine the brain mask. *file* specifies a brain mask "
                  "file (should be an image file consisting of 0s and 1s). "
                  "*thresh* specifies a threshold to use. By default all voxels"
                  "are used, unless one of these mask types are defined."),
                            mandatory=True)
    mask_file = File(exists=True,
                     desc="Mask file to be used if mask_type is 'file'.")
    mask_threshold = traits.Float(desc=("Mask threshold to be used if mask_type"
                                        " is 'thresh'."))
    intersect_mask = traits.Bool(True,
                                 desc=("Intersect the masks when computed from "
                                       "spm_global."))
    save_plot = traits.Bool(True, desc="save plots containing outliers",
                            usedefault=True)
    plot_type = traits.Enum('png', 'svg', 'eps', 'pdf',
                            desc="file type of the outlier plot",
                            usedefault=True)
    bound_by_brainmask = traits.Bool(False, desc=("use the brain mask to "
                                                 "determine bounding box"
                                                 "for composite norm (works"
                                                 "for SPM and Nipy - currently"
                                                 "inaccurate for FSL, AFNI"),
                                     usedefault=True)
    global_threshold = traits.Float(8.0, desc=("use this threshold when mask "
                                               "type equal's spm_global"),
                                    usedefault=True)

class ArtifactDetectOutputSpec(TraitedSpec):
    outlier_files = OutputMultiPath(File(exists=True),
            desc=("One file for each functional run containing a list of "
                  "0-based indices corresponding to outlier volumes"))
    intensity_files = OutputMultiPath(File(exists=True),
            desc=("One file for each functional run containing the global "
                  "intensity values determined from the brainmask"))
    norm_files = OutputMultiPath(File,
            desc=("One file for each functional run containing the composite "
                  "norm"))
    statistic_files = OutputMultiPath(File(exists=True),
            desc=("One file for each functional run containing information "
                  "about the different types of artifacts and if design info is"
                  " provided then details of stimulus correlated motion and a "
                  "listing or artifacts by event type."))
    plot_files = OutputMultiPath(File,
            desc=("One image file for each functional run containing the "
                  "detected outliers"))
    mask_files = OutputMultiPath(File,
            desc=("One image file for each functional run containing the mask"
                  "used for global signal calculation"))
    displacement_files = OutputMultiPath(File,
            desc=("One image file for each functional run containing the voxel"
                  "displacement timeseries"))

class ArtifactDetect(BaseInterface):
    """Detects outliers in a functional imaging series

    Uses intensity and motion parameters to infer outliers. If `use_norm` is
    True, it computes the movement of the center of each face a cuboid centered
    around the head and returns the maximal movement across the centers.

    Examples
    --------

    >>> ad = ArtifactDetect()
    >>> ad.inputs.realigned_files = 'functional.nii'
    >>> ad.inputs.realignment_parameters = 'functional.par'
    >>> ad.inputs.parameter_source = 'FSL'
    >>> ad.inputs.norm_threshold = 1
    >>> ad.inputs.use_differences = [True, False]
    >>> ad.inputs.zintensity_threshold = 3
    >>> ad.run() # doctest: +SKIP
    """

    input_spec = ArtifactDetectInputSpec
    output_spec = ArtifactDetectOutputSpec

    def __init__(self, **inputs):
        super(ArtifactDetect, self).__init__(**inputs)

    def _get_output_filenames(self, motionfile, output_dir):
        """Generate output files based on motion filenames

        Parameters
        ----------

        motionfile: file/string
            Filename for motion parameter file
        output_dir: string
            output directory in which the files will be generated
        """
        if isinstance(motionfile, six.string_types):
            infile = motionfile
        elif isinstance(motionfile, list):
            infile = motionfile[0]
        else:
            raise Exception("Unknown type of file")
        _, filename, ext = split_filename(infile)
        artifactfile = os.path.join(output_dir, ''.join(('art.', filename,
                                                         '_outliers.txt')))
        intensityfile = os.path.join(output_dir, ''.join(('global_intensity.',
                                                          filename, '.txt')))
        statsfile = os.path.join(output_dir, ''.join(('stats.', filename,
                                                      '.txt')))
        normfile = os.path.join(output_dir, ''.join(('norm.', filename,
                                                     '.txt')))
        plotfile = os.path.join(output_dir, ''.join(('plot.', filename, '.',
                                                     self.inputs.plot_type)))
        displacementfile = os.path.join(output_dir, ''.join(('disp.',
                                                             filename, ext)))
        maskfile = os.path.join(output_dir, ''.join(('mask.', filename, ext)))
        return (artifactfile, intensityfile, statsfile, normfile, plotfile,
                displacementfile, maskfile)

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['outlier_files'] = []
        outputs['intensity_files'] = []
        outputs['statistic_files'] = []
        outputs['mask_files'] = []
        if isdefined(self.inputs.use_norm) and self.inputs.use_norm:
            outputs['norm_files'] = []
            if self.inputs.bound_by_brainmask:
                outputs['displacement_files'] = []
        if isdefined(self.inputs.save_plot) and self.inputs.save_plot:
            outputs['plot_files'] = []
        for i, f in enumerate(filename_to_list(self.inputs.realigned_files)):
            (outlierfile, intensityfile, statsfile, normfile, plotfile,
             displacementfile, maskfile) = \
                                      self._get_output_filenames(f, os.getcwd())
            outputs['outlier_files'].insert(i, outlierfile)
            outputs['intensity_files'].insert(i, intensityfile)
            outputs['statistic_files'].insert(i, statsfile)
            outputs['mask_files'].insert(i, maskfile)
            if isdefined(self.inputs.use_norm) and self.inputs.use_norm:
                outputs['norm_files'].insert(i, normfile)
                if self.inputs.bound_by_brainmask:
                    outputs['displacement_files'].insert(i, displacementfile)
            if isdefined(self.inputs.save_plot) and self.inputs.save_plot:
                outputs['plot_files'].insert(i, plotfile)
        return outputs

    def _plot_outliers_with_wave(self, wave, outliers, name):
        import matplotlib.pyplot as plt
        plt.plot(wave)
        plt.ylim([wave.min(), wave.max()])
        plt.xlim([0, len(wave) - 1])
        if len(outliers):
            plt.plot(np.tile(outliers[:, None], (1, 2)).T,
                     np.tile([wave.min(), wave.max()], (len(outliers), 1)).T,
                     'r')
        plt.xlabel('Scans - 0-based')
        plt.ylabel(name)

    def _detect_outliers_core(self, imgfile, motionfile, runidx, cwd=None):
        """
        Core routine for detecting outliers
        """
        if not cwd:
            cwd = os.getcwd()

        # read in functional image
        if isinstance(imgfile, six.string_types):
            nim = load(imgfile)
        elif isinstance(imgfile, list):
            if len(imgfile) == 1:
                nim = load(imgfile[0])
            else:
                images = [load(f) for f in imgfile]
                nim = funcs.concat_images(images)

        # compute global intensity signal
        (x, y, z, timepoints) = nim.get_shape()

        data = nim.get_data()
        affine = nim.get_affine()
        g = np.zeros((timepoints, 1))
        masktype = self.inputs.mask_type
        if masktype == 'spm_global':  # spm_global like calculation
            iflogger.debug('art: using spm global')
            intersect_mask = self.inputs.intersect_mask
            if intersect_mask:
                mask = np.ones((x, y, z), dtype=bool)
                for t0 in range(timepoints):
                    vol = data[:, :, :, t0]
                    # Use an SPM like approach
                    mask_tmp = vol > \
                               (_nanmean(vol) / self.inputs.global_threshold)
                    mask = mask * mask_tmp
                for t0 in range(timepoints):
                    vol = data[:, :, :, t0]
                    g[t0] = _nanmean(vol[mask])
                if len(find_indices(mask)) < (np.prod((x, y, z)) / 10):
                    intersect_mask = False
                    g = np.zeros((timepoints, 1))
            if not intersect_mask:
                iflogger.info('not intersect_mask is True')
                mask = np.zeros((x, y, z, timepoints))
                for t0 in range(timepoints):
                    vol = data[:, :, :, t0]
                    mask_tmp = vol > \
                                 (_nanmean(vol) / self.inputs.global_threshold)
                    mask[:, :, :, t0] = mask_tmp
                    g[t0] = np.nansum(vol * mask_tmp)/np.nansum(mask_tmp)
        elif masktype == 'file':  # uses a mask image to determine intensity
            maskimg = load(self.inputs.mask_file)
            mask = maskimg.get_data()
            affine = maskimg.get_affine()
            mask = mask > 0.5
            for t0 in range(timepoints):
                vol = data[:, :, :, t0]
                g[t0] = _nanmean(vol[mask])
        elif masktype == 'thresh':  # uses a fixed signal threshold
            for t0 in range(timepoints):
                vol = data[:, :, :, t0]
                mask = vol > self.inputs.mask_threshold
                g[t0] = _nanmean(vol[mask])
        else:
            mask = np.ones((x, y, z))
            g = _nanmean(data[mask > 0, :], 1)

        # compute normalized intensity values
        gz = signal.detrend(g, axis=0)  # detrend the signal
        if self.inputs.use_differences[1]:
            gz = np.concatenate((np.zeros((1, 1)), np.diff(gz, n=1, axis=0)),
                                axis=0)
        gz = (gz - np.mean(gz)) / np.std(gz)  # normalize the detrended signal
        iidx = find_indices(abs(gz) > self.inputs.zintensity_threshold)

        # read in motion parameters
        mc_in = np.loadtxt(motionfile)
        mc = deepcopy(mc_in)

        (artifactfile, intensityfile, statsfile, normfile, plotfile,
         displacementfile, maskfile) = self._get_output_filenames(imgfile, cwd)
        mask_img = Nifti1Image(mask.astype(np.uint8), affine)
        mask_img.to_filename(maskfile)

        if self.inputs.use_norm:
            brain_pts = None
            if self.inputs.bound_by_brainmask:
                voxel_coords = np.nonzero(mask)
                coords = np.vstack((voxel_coords[0],
                                    np.vstack((voxel_coords[1],
                                               voxel_coords[2])))).T
                brain_pts = np.dot(affine,
                                   np.hstack((coords,
                                              np.ones((coords.shape[0], 1)))).T)
            # calculate the norm of the motion parameters
            normval, displacement = _calc_norm(mc,
                                               self.inputs.use_differences[0],
                                               self.inputs.parameter_source,
                                               brain_pts=brain_pts)
            tidx = find_indices(normval > self.inputs.norm_threshold)
            ridx = find_indices(normval < 0)
            if displacement is not None:
                dmap = np.zeros((x, y, z, timepoints), dtype=np.float)
                for i in range(timepoints):
                    dmap[voxel_coords[0],
                         voxel_coords[1],
                         voxel_coords[2], i] = displacement[i, :]
                dimg = Nifti1Image(dmap, affine)
                dimg.to_filename(displacementfile)
        else:
            if self.inputs.use_differences[0]:
                mc = np.concatenate((np.zeros((1, 6)),
                                     np.diff(mc_in, n=1, axis=0)),
                                    axis=0)
            traval = mc[:, 0:3]  # translation parameters (mm)
            rotval = mc[:, 3:6]  # rotation parameters (rad)
            tidx = find_indices(np.sum(abs(traval) >
                                       self.inputs.translation_threshold, 1)
                                > 0)
            ridx = find_indices(np.sum(abs(rotval) >
                                       self.inputs.rotation_threshold, 1) > 0)

        outliers = np.unique(np.union1d(iidx, np.union1d(tidx, ridx)))

        # write output to outputfile
        np.savetxt(artifactfile, outliers, fmt='%d', delimiter=' ')
        np.savetxt(intensityfile, g, fmt='%.2f', delimiter=' ')
        if self.inputs.use_norm:
            np.savetxt(normfile, normval, fmt='%.4f', delimiter=' ')

        if isdefined(self.inputs.save_plot) and self.inputs.save_plot:
            import matplotlib
            matplotlib.use(config.get("execution", "matplotlib_backend"))
            import matplotlib.pyplot as plt
            fig = plt.figure()
            if isdefined(self.inputs.use_norm) and self.inputs.use_norm:
                plt.subplot(211)
            else:
                plt.subplot(311)
            self._plot_outliers_with_wave(gz, iidx, 'Intensity')
            if isdefined(self.inputs.use_norm) and self.inputs.use_norm:
                plt.subplot(212)
                self._plot_outliers_with_wave(normval, np.union1d(tidx, ridx),
                                              'Norm (mm)')
            else:
                diff = ''
                if self.inputs.use_differences[0]:
                    diff = 'diff'
                plt.subplot(312)
                self._plot_outliers_with_wave(traval, tidx,
                                              'Translation (mm)' + diff)
                plt.subplot(313)
                self._plot_outliers_with_wave(rotval, ridx,
                                              'Rotation (rad)' + diff)
            plt.savefig(plotfile)
            plt.close(fig)

        motion_outliers = np.union1d(tidx, ridx)
        stats = [{'motion_file': motionfile,
                  'functional_file': imgfile},
                 {'common_outliers': len(np.intersect1d(iidx, motion_outliers)),
                  'intensity_outliers': len(np.setdiff1d(iidx,
                                                         motion_outliers)),
                  'motion_outliers': len(np.setdiff1d(motion_outliers, iidx)),
                  },
         {'motion': [{'using differences': self.inputs.use_differences[0]},
                      {'mean': np.mean(mc_in, axis=0).tolist(),
                       'min': np.min(mc_in, axis=0).tolist(),
                       'max': np.max(mc_in, axis=0).tolist(),
                       'std': np.std(mc_in, axis=0).tolist()},
                      ]},
         {'intensity': [{'using differences': self.inputs.use_differences[1]},
                         {'mean': np.mean(gz, axis=0).tolist(),
                          'min': np.min(gz, axis=0).tolist(),
                          'max': np.max(gz, axis=0).tolist(),
                          'std': np.std(gz, axis=0).tolist()},
                         ]},
                 ]
        if self.inputs.use_norm:
            stats.insert(3, {'motion_norm':
                                 {'mean': np.mean(normval, axis=0).tolist(),
                                  'min': np.min(normval, axis=0).tolist(),
                                  'max': np.max(normval, axis=0).tolist(),
                                  'std': np.std(normval, axis=0).tolist(),
                                  }})
        save_json(statsfile, stats)

    def _run_interface(self, runtime):
        """Execute this module.
        """
        funcfilelist = filename_to_list(self.inputs.realigned_files)
        motparamlist = filename_to_list(self.inputs.realignment_parameters)
        for i, imgf in enumerate(funcfilelist):
            self._detect_outliers_core(imgf, motparamlist[i], i,
                                       cwd=os.getcwd())
        return runtime

class StimCorrInputSpec(BaseInterfaceInputSpec):
    realignment_parameters = InputMultiPath(File(exists=True), mandatory=True,
        desc=('Names of realignment parameters corresponding to the functional '
              'data files'))
    intensity_values = InputMultiPath(File(exists=True), mandatory=True,
              desc='Name of file containing intensity values')
    spm_mat_file = File(exists=True, mandatory=True,
                        desc='SPM mat file (use pre-estimate SPM.mat file)')
    concatenated_design = traits.Bool(mandatory=True,
              desc='state if the design matrix contains concatenated sessions')

class StimCorrOutputSpec(TraitedSpec):
    stimcorr_files = OutputMultiPath(File(exists=True),
                     desc='List of files containing correlation values')

class StimulusCorrelation(BaseInterface):
    """Determines if stimuli are correlated with motion or intensity
    parameters.

    Currently this class supports an SPM generated design matrix and requires
    intensity parameters. This implies that one must run
    :ref:`ArtifactDetect <nipype.algorithms.rapidart.ArtifactDetect>`
    and :ref:`Level1Design <nipype.interfaces.spm.model.Level1Design>` prior to running this or
    provide an SPM.mat file and intensity parameters through some other means.

    Examples
    --------

    >>> sc = StimulusCorrelation()
    >>> sc.inputs.realignment_parameters = 'functional.par'
    >>> sc.inputs.intensity_values = 'functional.rms'
    >>> sc.inputs.spm_mat_file = 'SPM.mat'
    >>> sc.inputs.concatenated_design = False
    >>> sc.run() # doctest: +SKIP

    """

    input_spec = StimCorrInputSpec
    output_spec = StimCorrOutputSpec

    def _get_output_filenames(self, motionfile, output_dir):
        """Generate output files based on motion filenames

        Parameters
        ----------
        motionfile: file/string
            Filename for motion parameter file
        output_dir: string
            output directory in which the files will be generated
        """
        (_, filename) = os.path.split(motionfile)
        (filename, _) = os.path.splitext(filename)
        corrfile = os.path.join(output_dir, ''.join(('qa.', filename,
                                                     '_stimcorr.txt')))
        return corrfile

    def _stimcorr_core(self, motionfile, intensityfile, designmatrix, cwd=None):
        """
        Core routine for determining stimulus correlation

        """
        if not cwd:
            cwd = os.getcwd()
        # read in motion parameters
        mc_in = np.loadtxt(motionfile)
        g_in = np.loadtxt(intensityfile)
        g_in.shape = g_in.shape[0], 1
        dcol = designmatrix.shape[1]
        mccol = mc_in.shape[1]
        concat_matrix = np.hstack((np.hstack((designmatrix, mc_in)), g_in))
        cm = np.corrcoef(concat_matrix, rowvar=0)
        corrfile = self._get_output_filenames(motionfile, cwd)
        # write output to outputfile
        file = open(corrfile, 'w')
        file.write("Stats for:\n")
        file.write("Stimulus correlated motion:\n%s\n" % motionfile)
        for i in range(dcol):
            file.write("SCM.%d:" % i)
            for v in cm[i, dcol + np.arange(mccol)]:
                file.write(" %.2f" % v)
            file.write('\n')
        file.write("Stimulus correlated intensity:\n%s\n" % intensityfile)
        for i in range(dcol):
            file.write("SCI.%d: %.2f\n" % (i, cm[i, -1]))
        file.close()

    def _get_spm_submatrix(self, spmmat, sessidx, rows=None):
        """
        Parameters
        ----------
        spmmat: scipy matlab object
            full SPM.mat file loaded into a scipy object
        sessidx: int
            index to session that needs to be extracted.
        """
        designmatrix = spmmat['SPM'][0][0].xX[0][0].X
        U = spmmat['SPM'][0][0].Sess[0][sessidx].U[0]
        if rows is None:
            rows = spmmat['SPM'][0][0].Sess[0][sessidx].row[0] - 1
        cols = spmmat['SPM'][0][0].Sess[0][sessidx].col[0][range(len(U))] - 1
        outmatrix = designmatrix.take(rows.tolist(), axis=0).take(cols.tolist(),
                                                                  axis=1)
        return outmatrix

    def _run_interface(self, runtime):
        """Execute this module.
        """
        motparamlist = self.inputs.realignment_parameters
        intensityfiles = self.inputs.intensity_values
        spmmat = sio.loadmat(self.inputs.spm_mat_file, struct_as_record=False)
        nrows = []
        for i in range(len(motparamlist)):
            sessidx = i
            rows = None
            if self.inputs.concatenated_design:
                sessidx = 0
                mc_in = np.loadtxt(motparamlist[i])
                rows = np.sum(nrows) + np.arange(mc_in.shape[0])
                nrows.append(mc_in.shape[0])
            matrix = self._get_spm_submatrix(spmmat, sessidx, rows)
            self._stimcorr_core(motparamlist[i], intensityfiles[i],
                                matrix, os.getcwd())
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        files = []
        for i, f in enumerate(self.inputs.realignment_parameters):
            files.insert(i, self._get_output_filenames(f, os.getcwd()))
        if files:
            outputs['stimcorr_files'] = files
        return outputs


# fifteenwordsnippets.py

# A script that searches a HathiTrust corpus of 6,942 volumes (1700-1923), plus Hoyt & Richard's
# corpus of 808 vols (1923-1950), for words related to money. It takes seven words on either
# side of those words to create a snippet.

# In cases where the possibly-monetary word is ambiguous, e.g. "pounds," it runs the central
# seven words of the snippet through a regularized logistic model (created by model_contexts)
# in order to make a prediction about the likelihood that this word refers to money. The
# model I used is based on 700 manually-tagged snippets; it's about about 87% accurate,
# five-fold crossvalidated.

import modelingcounter
import os, sys
import SonicScrewdriver as utils
import csv
import pickle
from bagofwords import WordVector, StandardizingVector
from sklearn.linear_model import LogisticRegression

# We start with a number of functions borrowed from other scripts; these were used to
# generate the logistic model, so we use them also here to clean and normalize snippets.

punctuple = ('.', ',', '?', '!', ';', '"', '“', '”', ':', '--', '—', ')', '(', "'", "`", "[", "]", "{", "}")

def all_nonalphanumeric(astring):
    nonalphanum = True
    for character in astring:
        if character.isalpha() or character.isdigit():
            nonalphanum = False
            break
    return nonalphanum

def strip_punctuation(astring):
    global punctuple
    keepclipping = True
    suffix = ""
    while keepclipping == True and len(astring) > 1:
        keepclipping = False
        if astring.endswith(punctuple):
            suffix = astring[-1:] + suffix
            astring = astring[:-1]
            keepclipping = True
    keepclipping = True
    prefix = ""
    while keepclipping == True and len(astring) > 1:
        keepclipping = False
        if astring.startswith(punctuple):
            prefix = prefix + astring[:1]
            astring = astring[1:]
            keepclipping = True
    return(prefix, astring, suffix)

def as_wordlist(line):
    ''' Converts a line into a list of words, splitting
    tokens brutally and unreflectively at punctuation.
    One of the effects will be to split possessives into noun
    and s. But this might not be a bad thing for current
    purposes.
    '''

    line = line.replace('”', ' ')
    line = line.replace(':', ' ')
    line = line.replace(';', ' ')
    line = line.replace('—', ' ')
    line = line.replace('--', ' ')
    line = line.replace('.', ' ')
    line = line.replace(',', ' ')
    line = line.replace('-', ' ')
    line = line.replace('—', ' ')
    line = line.replace("'", ' ')
    line = line.replace('"', ' ')

    # That's not the most efficient way to do this computationally,
    # but it prevents me from having to look up the .translate
    # method.

    words = line.split(' ')

    wordlist = list()

    for word in words:
        word = word.lower()
        prefix, word, suffix = strip_punctuation(word)
        # In case we missed anything.

        if len(word) > 0 and not all_nonalphanumeric(word):
            wordlist.append(word)

    return wordlist

def is_money(wordlist, WINDOWRADIUS, model, features, standardizer):
    # We're getting a wordlist generated by WINDOWRADIUS, but
    # we only want the central seven words for our model.

    startindex = WINDOWRADIUS - 3
    endindex = WINDOWRADIUS + 4

    modelsegment = ' '.join(wordlist[startindex : endindex])

    # You'd think we could just take a list of words, but in
    # generating the model we ran strings through a particular
    # tokenizing process, and we should replicate that here.

    normalizedlist = as_wordlist(modelsegment)
    vector = WordVector(normalizedlist)
    vector.selectfeatures(features)
    # This turns a sparse dictionary into an array with zeroes
    # for missing features.

    vector.normalizefrequencies()
    # raw counts are divided by total counts.

    vector.standardizefrequencies(standardizer)
    # features are now centered on the means, and divided by
    # standard deviations, calculated on the training set

    classlabel = model.predict(vector.features)[0]

    if classlabel == 1:
        return True
    elif classlabel == 0:
        return False
    else:
        print("ANOMALY!")
        print(classlabel)
        return False

        # Cause that's how I do error handling.

# Main script.

# Let's load the model.

modelfolder = "/Volumes/TARDIS/work/moneycontext/"
modelpath = modelfolder + "logisticmodel.p"
with open(modelpath, mode = 'rb') as f:
    logisticmodel = pickle.load(f)

standardizerpath = modelfolder + 'standardizer.p'
with open(standardizerpath, mode = 'rb') as f:
    standardizer = pickle.load(f)

featurepath = modelfolder + 'featurelist.p'
with open(featurepath, mode = 'rb') as f:
    features = pickle.load(f)

# Now load HathiTrust metadata.

rows, columns, table = utils.readtsv('/Volumes/TARDIS/work/metadata/MergedMonographs.tsv')

ambiguouswords = {'crown', 'crowns', 'guinea', 'guineas', 'nickel', 'sovereign', 'sovereigns', 'pound', 'pounds', 'quid'}

moneywords = {'dollar', 'dollars', 'dime', 'dimes', 'nickel', 'nickels', 'pound', 'pounds', 'shilling', 'shillings', 'sovereign', 'sovereigns','cent', 'cents', 'centime', 'centimes', 'crown', 'crowns', 'halfcrown', 'half-crown','penny', 'pennies', 'pence', 'farthing', 'farthings', 'franc', 'francs', 'guilder', 'guilders', 'florin', 'florins', 'guinea', 'guineas', "ha'penny", 'tuppence', 'twopence', 'sixpence', '|arabicprice|', '|price|', 'quid'}

# Words I explicitly decided not to include: 'quarter', 'quarters', 'mark', 'marks.' Monetary uses
# seemed rare enough relative to others that they'd be more likely to introduce noise than to help.
# |arabicprice| is a code the tokenizer in modelingcounter produces whenever it encounters
# a number connected to £, $, ¢, s, or d. In the output we convert that to |price|, for no very
# good reason.

wealthwords = {'fortune', 'fortunes', 'wealth', 'rich', 'riches', 'money', 'moneys', 'fund', 'funds', 'sum', 'sums', 'price', 'prices', 'priced'}

# This is by no means an exhaustive list. Owe, loan, borrowed, etc.
# If we really want to get at the full range of words potentially
# associated with money, topic modeling would be an appropriate lever.
# We can perhaps enumerate currency terms intuitively, but not these.

alltargetwords = moneywords

sourcedir = "/Volumes/TARDIS/work/moneytexts/"
filelist = os.listdir(sourcedir)
filelist = [x for x in filelist if x.endswith(".txt")]
contexts = []

WINDOWRADIUS = 7

ctr = 0

for filename in filelist:

    htid = utils.pairtreelabel(filename.replace('.fic.txt', ''))

    if htid not in rows:
        print(htid)
        continue
    else:
        date = utils.simple_date(htid, table)

    filepath = os.path.join(sourcedir, filename)
    with open(filepath, encoding = 'utf-8') as f:
        filelines = f.readlines()
    pagelist = [filelines]

    # The wordcounter module expects a list of pages, each of which is a list of lines.
    # Ebooks have no pages -- at least as I currently receive them -- so we treat it
    # all as one giant page.

    tokenstream = modelingcounter.makestream(pagelist)

    newcontexts = modelingcounter.extract_snippets(tokenstream,  WINDOWRADIUS, alltargetwords)

    approvedcontexts = []

    for snippet, snippettomodel in newcontexts:

        keyword = snippettomodel[WINDOWRADIUS]
        keyword = keyword.lower()
        prefix, keyword, suffix = strip_punctuation(keyword)

        if keyword in wealthwords:
            category = 'wealth'
        elif keyword in ambiguouswords:
            currency = is_money(snippettomodel, WINDOWRADIUS, logisticmodel, features, standardizer)
            if currency:
                category = 'money'
            else:
                category = "notmoney"
        elif keyword in moneywords:
            category = 'money'
        else:
            print('ANOMALY: ' + keyword)
            # Cause that's how I do error handling.
            category = 'null'

        if category == 'money':
            approvedcontexts.append((htid, date, snippet, keyword, category))

    print(ctr)
    ctr += 1

    outfile = "/Volumes/TARDIS/work/moneycontext/twentyfivesnippets.tsv"
    with open(outfile, mode='a', encoding='utf-8') as f:
        for context in approvedcontexts:
            htid, date, alist, keyword, category = context
            snippet = " ".join(alist)
            snippet = snippet.replace('\t', '')
            # Because we don't want stray tabs in our tab-separated values.

            f.write(htid + '\t' + str(date) + '\t' + keyword + '\t' + category + '\t' + snippet + '\n')

sourcedir = "/Volumes/TARDIS/work/US_NOVELS_1923-1950/"

filelist = os.listdir(sourcedir)
fileset = set([x for x in filelist if x.endswith(".txt")])
filelist = list(fileset)

metafile = os.path.join(sourcedir, "US_NOVELS_1923-1950_META.txt")

datedict = dict()
dateset = set()

with open(metafile, newline='', encoding = 'utf-8') as f:
    reader = csv.reader(f)
    for fields in reader:
        idcode = fields[0]
        date = int(fields[8])
        datedict[idcode] = date
        dateset.add(date)

for filename in filelist:

    htid = utils.pairtreelabel(filename.replace('.txt', ''))

    if htid not in datedict:
        print(htid)
        continue
    else:
        date = datedict[htid]

    filepath = os.path.join(sourcedir, filename)
    with open(filepath, encoding = 'utf-8') as f:
        filelines = f.readlines()
    pagelist = [filelines]

    # The wordcounter module expects a list of pages, each of which is a list of lines.
    # Ebooks have no pages -- at least as I currently receive them -- so we treat it
    # all as one giant page.

    tokenstream = modelingcounter.makestream(pagelist)

    newcontexts = modelingcounter.extract_snippets(tokenstream, WINDOWRADIUS, alltargetwords)
    approvedcontexts = []

    for snippet, snippettomodel in newcontexts:

        keyword = snippettomodel[WINDOWRADIUS]
        keyword = keyword.lower()
        prefix, keyword, suffix = strip_punctuation(keyword)

        if keyword in wealthwords:
            category = 'wealth'
        elif keyword in ambiguouswords:
            currency = is_money(snippettomodel, WINDOWRADIUS, logisticmodel, features, standardizer)
            if currency:
                category = 'money'
            else:
                category = "notmoney"
        elif keyword in moneywords:
            category = 'money'
        else:
            print('ANOMALY: ' + keyword)
            # Cause that's how I do error handling.
            category = 'null'

        if category == 'money':
            approvedcontexts.append((htid, date, snippet, keyword, category))

    outfile = "/Volumes/TARDIS/work/moneycontext/twentyfivesnippets.tsv"
    with open(outfile, mode='a', encoding='utf-8') as f:
        for context in approvedcontexts:
            htid, date, alist, keyword, category = context
            snippet = " ".join(alist)
            snippet = snippet.replace('\t', '')
            # Because we don't want stray tabs in our tab-separated values.

            f.write(htid + '\t' + str(date) + '\t' + keyword + '\t' + category + '\t' + snippet + '\n')


import numpy as np
import as_classification.ann_models
import as_sound.features.extractFeatures as ef
import as_classification.utilities
import os
import matplotlib.pyplot as plt

# -------------------------------
# CREATE MODEL
# -------------------------------

model = as_classification.ann_models.ANN_5FCL()
model.initialize(15,2)

# -------------------------------
# READ AUDIO FILES
# -------------------------------

speech_data = ef.computeSupervectorForFile(os.path.dirname(os.path.realpath(__file__)) + '/data/Speech.wav', 8000, 2048, 2049)
noise_data = ef.computeSupervectorForFile(os.path.dirname(os.path.realpath(__file__)) + '/data/Noise.wav', 8000, 2048, 2049)

whole_data = np.hstack((speech_data, noise_data))
whole_data = np.swapaxes(whole_data,0,1)
whole_labels = np.zeros((whole_data.shape[0], 2))
whole_labels[:speech_data.shape[1],0] = 1
whole_labels[speech_data.shape[1]:,1] = 1

training_data = {"labels": whole_labels,
                 "data": whole_data}

training_data, test_data = as_classification.utilities.divideTrainingData(training_data, 0.6)

# -------------------------------
# TRAIN
# -------------------------------

model.train(training_data, test_data, 10, 20, 100)

model.saveCheckpoint(os.path.dirname(os.path.realpath(__file__)) + '/data/vadModel_ANN_5FCL.chkp')

#xp = np.arange(0,prediction.shape[0])

#plt.plot(xp, test_data[:,14], '-b', label='RMS')
#plt.plot(xp, prediction[:,0], '-r', label='ANN Output')
#plt.legend(loc='upper left')

#plt.show()

#feat = as_sound.features.extractFeatures.computeSupervector(normalized_data)


import pandas as pd
import numpy as np

import pytest

import dask
import dask.dataframe as dd

from dask.dataframe._compat import tm, PANDAS_GT_100
from dask.dataframe.indexing import _coerce_loc_index
from dask.dataframe.utils import assert_eq, make_meta, PANDAS_VERSION

dsk = {
    ("x", 0): pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[0, 1, 3]),
    ("x", 1): pd.DataFrame({"a": [4, 5, 6], "b": [3, 2, 1]}, index=[5, 6, 8]),
    ("x", 2): pd.DataFrame({"a": [7, 8, 9], "b": [0, 0, 0]}, index=[9, 9, 9]),
}
meta = make_meta({"a": "i8", "b": "i8"}, index=pd.Index([], "i8"))
d = dd.DataFrame(dsk, "x", meta, [0, 5, 9, 9])
full = d.compute()
CHECK_FREQ = {}
if dd._compat.PANDAS_GT_110:
    CHECK_FREQ["check_freq"] = False

def test_loc():
    assert d.loc[3:8].divisions[0] == 3
    assert d.loc[3:8].divisions[-1] == 8

    assert d.loc[5].divisions == (5, 5)

    assert_eq(d.loc[5], full.loc[5:5])
    assert_eq(d.loc[3:8], full.loc[3:8])
    assert_eq(d.loc[:8], full.loc[:8])
    assert_eq(d.loc[3:], full.loc[3:])
    assert_eq(d.loc[[5]], full.loc[[5]])

    expected_warning = FutureWarning

    if not PANDAS_GT_100:
        # removed in pandas 1.0
        with pytest.warns(expected_warning):
            assert_eq(d.loc[[3, 4, 1, 8]], full.loc[[3, 4, 1, 8]])
        with pytest.warns(expected_warning):
            assert_eq(d.loc[[3, 4, 1, 9]], full.loc[[3, 4, 1, 9]])
        with pytest.warns(expected_warning):
            assert_eq(d.loc[np.array([3, 4, 1, 9])], full.loc[np.array([3, 4, 1, 9])])

    assert_eq(d.a.loc[5], full.a.loc[5:5])
    assert_eq(d.a.loc[3:8], full.a.loc[3:8])
    assert_eq(d.a.loc[:8], full.a.loc[:8])
    assert_eq(d.a.loc[3:], full.a.loc[3:])
    assert_eq(d.a.loc[[5]], full.a.loc[[5]])
    if not PANDAS_GT_100:
        # removed in pandas 1.0
        with pytest.warns(expected_warning):
            assert_eq(d.a.loc[[3, 4, 1, 8]], full.a.loc[[3, 4, 1, 8]])
        with pytest.warns(expected_warning):
            assert_eq(d.a.loc[[3, 4, 1, 9]], full.a.loc[[3, 4, 1, 9]])
        with pytest.warns(expected_warning):
            assert_eq(
                d.a.loc[np.array([3, 4, 1, 9])], full.a.loc[np.array([3, 4, 1, 9])]
            )
    assert_eq(d.a.loc[[]], full.a.loc[[]])
    assert_eq(d.a.loc[np.array([])], full.a.loc[np.array([])])

    pytest.raises(KeyError, lambda: d.loc[1000])
    assert_eq(d.loc[1000:], full.loc[1000:])
    assert_eq(d.loc[-2000:-1000], full.loc[-2000:-1000])

    assert sorted(d.loc[5].dask) == sorted(d.loc[5].dask)
    assert sorted(d.loc[5].dask) != sorted(d.loc[6].dask)

def test_loc_non_informative_index():
    df = pd.DataFrame({"x": [1, 2, 3, 4]}, index=[10, 20, 30, 40])
    ddf = dd.from_pandas(df, npartitions=2, sort=True)
    ddf.divisions = (None,) * 3
    assert not ddf.known_divisions

    ddf.loc[20:30].compute(scheduler="sync")

    assert_eq(ddf.loc[20:30], df.loc[20:30])

    df = pd.DataFrame({"x": [1, 2, 3, 4]}, index=[10, 20, 20, 40])
    ddf = dd.from_pandas(df, npartitions=2, sort=True)
    assert_eq(ddf.loc[20], df.loc[20:20])

def test_loc_with_text_dates():
    A = dd._compat.makeTimeSeries().iloc[:5]
    B = dd._compat.makeTimeSeries().iloc[5:]
    s = dd.Series(
        {("df", 0): A, ("df", 1): B},
        "df",
        A,
        [A.index.min(), B.index.min(), B.index.max()],
    )

    assert s.loc["2000":"2010"].divisions == s.divisions
    assert_eq(s.loc["2000":"2010"], s)
    assert len(s.loc["2000-01-03":"2000-01-05"].compute()) == 3

def test_loc_with_series():
    assert_eq(d.loc[d.a % 2 == 0], full.loc[full.a % 2 == 0])

    assert sorted(d.loc[d.a % 2].dask) == sorted(d.loc[d.a % 2].dask)
    assert sorted(d.loc[d.a % 2].dask) != sorted(d.loc[d.a % 3].dask)

def test_loc_with_array():
    assert_eq(d.loc[(d.a % 2 == 0).values], full.loc[(full.a % 2 == 0).values])

    assert sorted(d.loc[(d.a % 2).values].dask) == sorted(d.loc[(d.a % 2).values].dask)
    assert sorted(d.loc[(d.a % 2).values].dask) != sorted(d.loc[(d.a % 3).values].dask)

def test_loc_with_function():
    assert_eq(d.loc[lambda df: df["a"] > 3, :], full.loc[lambda df: df["a"] > 3, :])

    def _col_loc_fun(_df):
        return _df.columns.str.contains("b")

    assert_eq(d.loc[:, _col_loc_fun], full.loc[:, _col_loc_fun])

def test_loc_with_array_different_partition():
    df = pd.DataFrame(
        np.random.randn(20, 5),
        index=list("abcdefghijklmnopqrst"),
        columns=list("ABCDE"),
    )
    ddf = dd.from_pandas(df, 3)

    assert_eq(ddf.loc[(ddf.A > 0).values], df.loc[(df.A > 0).values])
    with pytest.raises(ValueError):
        ddf.loc[(ddf.A > 0).repartition(["a", "g", "k", "o", "t"]).values]

def test_loc_with_series_different_partition():
    df = pd.DataFrame(
        np.random.randn(20, 5),
        index=list("abcdefghijklmnopqrst"),
        columns=list("ABCDE"),
    )
    ddf = dd.from_pandas(df, 3)

    assert_eq(ddf.loc[ddf.A > 0], df.loc[df.A > 0])
    assert_eq(
        ddf.loc[(ddf.A > 0).repartition(["a", "g", "k", "o", "t"])], df.loc[df.A > 0]
    )

def test_loc2d():
    # index indexer is always regarded as slice for duplicated values
    assert_eq(d.loc[5, "a"], full.loc[5:5, "a"])
    # assert_eq(d.loc[[5], 'a'], full.loc[[5], 'a'])
    assert_eq(d.loc[5, ["a"]], full.loc[5:5, ["a"]])
    # assert_eq(d.loc[[5], ['a']], full.loc[[5], ['a']])

    assert_eq(d.loc[3:8, "a"], full.loc[3:8, "a"])
    assert_eq(d.loc[:8, "a"], full.loc[:8, "a"])
    assert_eq(d.loc[3:, "a"], full.loc[3:, "a"])
    assert_eq(d.loc[[8], "a"], full.loc[[8], "a"])

    assert_eq(d.loc[3:8, ["a"]], full.loc[3:8, ["a"]])
    assert_eq(d.loc[:8, ["a"]], full.loc[:8, ["a"]])
    assert_eq(d.loc[3:, ["a"]], full.loc[3:, ["a"]])

    # 3d
    with pytest.raises(pd.core.indexing.IndexingError):
        d.loc[3, 3, 3]

    # Series should raise
    with pytest.raises(pd.core.indexing.IndexingError):
        d.a.loc[3, 3]

    with pytest.raises(pd.core.indexing.IndexingError):
        d.a.loc[3:, 3]

    with pytest.raises(pd.core.indexing.IndexingError):
        d.a.loc[d.a % 2 == 0, 3]

@pytest.mark.skip(PANDAS_GT_100, reason="Removed in pandas 1.0")
def test_loc2d_some_missing():
    with pytest.warns(FutureWarning):
        assert_eq(d.loc[[3, 4, 3], ["a"]], full.loc[[3, 4, 3], ["a"]])

def test_loc2d_with_known_divisions():
    df = pd.DataFrame(
        np.random.randn(20, 5),
        index=list("abcdefghijklmnopqrst"),
        columns=list("ABCDE"),
    )
    ddf = dd.from_pandas(df, 3)

    assert_eq(ddf.loc["a", "A"], df.loc[["a"], "A"])
    assert_eq(ddf.loc["a", ["A"]], df.loc[["a"], ["A"]])
    assert_eq(ddf.loc["a":"o", "A"], df.loc["a":"o", "A"])
    assert_eq(ddf.loc["a":"o", ["A"]], df.loc["a":"o", ["A"]])
    assert_eq(ddf.loc[["n"], ["A"]], df.loc[["n"], ["A"]])
    assert_eq(ddf.loc[["a", "c", "n"], ["A"]], df.loc[["a", "c", "n"], ["A"]])
    assert_eq(ddf.loc[["t", "b"], ["A"]], df.loc[["t", "b"], ["A"]])
    assert_eq(
        ddf.loc[["r", "r", "c", "g", "h"], ["A"]],
        df.loc[["r", "r", "c", "g", "h"], ["A"]],
    )

def test_loc2d_with_unknown_divisions():
    df = pd.DataFrame(
        np.random.randn(20, 5),
        index=list("abcdefghijklmnopqrst"),
        columns=list("ABCDE"),
    )
    ddf = dd.from_pandas(df, 3)

    ddf.divisions = (None,) * len(ddf.divisions)
    assert ddf.known_divisions is False

    assert_eq(ddf.loc["a", "A"], df.loc[["a"], "A"])
    assert_eq(ddf.loc["a", ["A"]], df.loc[["a"], ["A"]])
    assert_eq(ddf.loc["a":"o", "A"], df.loc["a":"o", "A"])
    assert_eq(ddf.loc["a":"o", ["A"]], df.loc["a":"o", ["A"]])

def test_loc2d_duplicated_columns():
    df = pd.DataFrame(
        np.random.randn(20, 5),
        index=list("abcdefghijklmnopqrst"),
        columns=list("AABCD"),
    )
    ddf = dd.from_pandas(df, 3)

    assert_eq(ddf.loc["a", "A"], df.loc[["a"], "A"])
    assert_eq(ddf.loc["a", ["A"]], df.loc[["a"], ["A"]])
    assert_eq(ddf.loc["j", "B"], df.loc[["j"], "B"])
    assert_eq(ddf.loc["j", ["B"]], df.loc[["j"], ["B"]])

    assert_eq(ddf.loc["a":"o", "A"], df.loc["a":"o", "A"])
    assert_eq(ddf.loc["a":"o", ["A"]], df.loc["a":"o", ["A"]])
    assert_eq(ddf.loc["j":"q", "B"], df.loc["j":"q", "B"])
    assert_eq(ddf.loc["j":"q", ["B"]], df.loc["j":"q", ["B"]])

    assert_eq(ddf.loc["a":"o", "B":"D"], df.loc["a":"o", "B":"D"])
    assert_eq(ddf.loc["a":"o", "B":"D"], df.loc["a":"o", "B":"D"])
    assert_eq(ddf.loc["j":"q", "B":"A"], df.loc["j":"q", "B":"A"])
    assert_eq(ddf.loc["j":"q", "B":"A"], df.loc["j":"q", "B":"A"])

    assert_eq(ddf.loc[ddf.B > 0, "B"], df.loc[df.B > 0, "B"])
    assert_eq(ddf.loc[ddf.B > 0, ["A", "C"]], df.loc[df.B > 0, ["A", "C"]])

def test_getitem():
    df = pd.DataFrame(
        {
            "A": [1, 2, 3, 4, 5, 6, 7, 8, 9],
            "B": [9, 8, 7, 6, 5, 4, 3, 2, 1],
            "C": [True, False, True] * 3,
        },
        columns=list("ABC"),
    )
    ddf = dd.from_pandas(df, 2)
    assert_eq(ddf["A"], df["A"])
    # check cache consistency
    tm.assert_series_equal(ddf["A"]._meta, ddf._meta["A"])

    assert_eq(ddf[["A", "B"]], df[["A", "B"]])
    tm.assert_frame_equal(ddf[["A", "B"]]._meta, ddf._meta[["A", "B"]])

    assert_eq(ddf[ddf.C], df[df.C])
    tm.assert_series_equal(ddf.C._meta, ddf._meta.C)

    assert_eq(ddf[ddf.C.repartition([0, 2, 5, 8])], df[df.C])

    pytest.raises(KeyError, lambda: df["X"])
    pytest.raises(KeyError, lambda: df[["A", "X"]])
    pytest.raises(AttributeError, lambda: df.X)

    # not str/unicode
    df = pd.DataFrame(np.random.randn(10, 5))
    ddf = dd.from_pandas(df, 2)
    assert_eq(ddf[0], df[0])
    assert_eq(ddf[[1, 2]], df[[1, 2]])

    pytest.raises(KeyError, lambda: df[8])
    pytest.raises(KeyError, lambda: df[[1, 8]])

def test_getitem_slice():
    df = pd.DataFrame(
        {
            "A": [1, 2, 3, 4, 5, 6, 7, 8, 9],
            "B": [9, 8, 7, 6, 5, 4, 3, 2, 1],
            "C": [True, False, True] * 3,
        },
        index=list("abcdefghi"),
    )
    ddf = dd.from_pandas(df, 3)
    assert_eq(ddf["a":"e"], df["a":"e"])
    assert_eq(ddf["a":"b"], df["a":"b"])
    assert_eq(ddf["f":], df["f":])

def test_getitem_integer_slice():
    df = pd.DataFrame({"A": range(6)})
    ddf = dd.from_pandas(df, 2)
    # integer slicing is iloc based
    with pytest.raises(NotImplementedError):
        ddf[1:3]

    df = pd.DataFrame({"A": range(6)}, index=[1.0, 2.0, 3.0, 5.0, 10.0, 11.0])
    ddf = dd.from_pandas(df, 2)
    # except for float dtype indexes
    assert_eq(ddf[2:8], df[2:8])
    assert_eq(ddf[2:], df[2:])
    assert_eq(ddf[:8], df[:8])

def test_loc_on_numpy_datetimes():
    df = pd.DataFrame(
        {"x": [1, 2, 3]}, index=list(map(np.datetime64, ["2014", "2015", "2016"]))
    )
    a = dd.from_pandas(df, 2)
    a.divisions = list(map(np.datetime64, a.divisions))

    assert_eq(a.loc["2014":"2015"], a.loc["2014":"2015"])

def test_loc_on_pandas_datetimes():
    df = pd.DataFrame(
        {"x": [1, 2, 3]}, index=list(map(pd.Timestamp, ["2014", "2015", "2016"]))
    )
    a = dd.from_pandas(df, 2)
    a.divisions = list(map(pd.Timestamp, a.divisions))

    assert_eq(a.loc["2014":"2015"], a.loc["2014":"2015"])

def test_loc_datetime_no_freq():
    # https://github.com/dask/dask/issues/2389

    datetime_index = pd.date_range("2016-01-01", "2016-01-31", freq="12h")
    datetime_index.freq = None  # FORGET FREQUENCY
    df = pd.DataFrame({"num": range(len(datetime_index))}, index=datetime_index)

    ddf = dd.from_pandas(df, npartitions=1)
    slice_ = slice("2016-01-03", "2016-01-05")
    result = ddf.loc[slice_, :]
    expected = df.loc[slice_, :]
    assert_eq(result, expected)

def test_coerce_loc_index():
    for t in [pd.Timestamp, np.datetime64]:
        assert isinstance(_coerce_loc_index([t("2014")], "2014"), t)

def test_loc_timestamp_str():

    df = pd.DataFrame(
        {"A": np.random.randn(100), "B": np.random.randn(100)},
        index=pd.date_range("2011-01-01", freq="H", periods=100),
    )
    ddf = dd.from_pandas(df, 10)

    # partial string slice
    assert_eq(df.loc["2011-01-02"], ddf.loc["2011-01-02"])
    assert_eq(df.loc["2011-01-02":"2011-01-10"], ddf.loc["2011-01-02":"2011-01-10"])
    # same reso, dask result is always DataFrame
    assert_eq(
        df.loc["2011-01-02 10:00"].to_frame().T,
        ddf.loc["2011-01-02 10:00"],
        **CHECK_FREQ
    )

    # series
    assert_eq(df.A.loc["2011-01-02"], ddf.A.loc["2011-01-02"], **CHECK_FREQ)
    assert_eq(
        df.A.loc["2011-01-02":"2011-01-10"],
        ddf.A.loc["2011-01-02":"2011-01-10"],
        **CHECK_FREQ
    )

    # slice with timestamp (dask result must be DataFrame)
    assert_eq(
        df.loc[pd.Timestamp("2011-01-02")].to_frame().T,
        ddf.loc[pd.Timestamp("2011-01-02")],
        **CHECK_FREQ
    )
    assert_eq(
        df.loc[pd.Timestamp("2011-01-02") : pd.Timestamp("2011-01-10")],
        ddf.loc[pd.Timestamp("2011-01-02") : pd.Timestamp("2011-01-10")],
        **CHECK_FREQ
    )
    assert_eq(
        df.loc[pd.Timestamp("2011-01-02 10:00")].to_frame().T,
        ddf.loc[pd.Timestamp("2011-01-02 10:00")],
        **CHECK_FREQ
    )

    df = pd.DataFrame(
        {"A": np.random.randn(100), "B": np.random.randn(100)},
        index=pd.date_range("2011-01-01", freq="M", periods=100),
    )
    ddf = dd.from_pandas(df, 50)
    assert_eq(df.loc["2011-01"], ddf.loc["2011-01"])
    assert_eq(df.loc["2011"], ddf.loc["2011"])

    assert_eq(df.loc["2011-01":"2012-05"], ddf.loc["2011-01":"2012-05"])
    assert_eq(df.loc["2011":"2015"], ddf.loc["2011":"2015"])

    # series
    assert_eq(df.B.loc["2011-01"], ddf.B.loc["2011-01"])
    assert_eq(df.B.loc["2011"], ddf.B.loc["2011"])

    assert_eq(df.B.loc["2011-01":"2012-05"], ddf.B.loc["2011-01":"2012-05"])
    assert_eq(df.B.loc["2011":"2015"], ddf.B.loc["2011":"2015"])

def test_getitem_timestamp_str():

    df = pd.DataFrame(
        {"A": np.random.randn(100), "B": np.random.randn(100)},
        index=pd.date_range("2011-01-01", freq="H", periods=100),
    )
    ddf = dd.from_pandas(df, 10)

    # partial string slice
    assert_eq(df["2011-01-02"], ddf["2011-01-02"])
    assert_eq(df["2011-01-02":"2011-01-10"], df["2011-01-02":"2011-01-10"])

    df = pd.DataFrame(
        {"A": np.random.randn(100), "B": np.random.randn(100)},
        index=pd.date_range("2011-01-01", freq="D", periods=100),
    )
    ddf = dd.from_pandas(df, 50)
    assert_eq(df["2011-01"], ddf["2011-01"])
    assert_eq(df["2011"], ddf["2011"])

    assert_eq(df["2011-01":"2012-05"], ddf["2011-01":"2012-05"])
    assert_eq(df["2011":"2015"], ddf["2011":"2015"])

def test_loc_period_str():
    # .loc with PeriodIndex doesn't support partial string indexing
    # https://github.com/pydata/pandas/issues/13429
    pass

def test_getitem_period_str():

    df = pd.DataFrame(
        {"A": np.random.randn(100), "B": np.random.randn(100)},
        index=pd.period_range("2011-01-01", freq="H", periods=100),
    )
    ddf = dd.from_pandas(df, 10)

    # partial string slice
    assert_eq(df["2011-01-02"], ddf["2011-01-02"])
    assert_eq(df["2011-01-02":"2011-01-10"], df["2011-01-02":"2011-01-10"])
    # same reso, dask result is always DataFrame

    df = pd.DataFrame(
        {"A": np.random.randn(100), "B": np.random.randn(100)},
        index=pd.period_range("2011-01-01", freq="D", periods=100),
    )
    ddf = dd.from_pandas(df, 50)
    assert_eq(df["2011-01"], ddf["2011-01"])
    assert_eq(df["2011"], ddf["2011"])

    assert_eq(df["2011-01":"2012-05"], ddf["2011-01":"2012-05"])
    assert_eq(df["2011":"2015"], ddf["2011":"2015"])

def test_to_series():

    # Test for time index
    df = pd.DataFrame(
        {"A": np.random.randn(100)},
        index=pd.date_range("2011-01-01", freq="H", periods=100),
    )
    ddf = dd.from_pandas(df, 10)

    assert_eq(df.index.to_series(), ddf.index.to_series())

    # Test for numerical index
    df = pd.DataFrame({"A": np.random.randn(100)}, index=range(100))
    ddf = dd.from_pandas(df, 10)

    assert_eq(df.index.to_series(), ddf.index.to_series())

def test_to_frame():

    # Test for time index
    df = pd.DataFrame(
        {"A": np.random.randn(100)},
        index=pd.date_range("2011-01-01", freq="H", periods=100),
    )
    ddf = dd.from_pandas(df, 10)

    assert_eq(df.index.to_frame(), ddf.index.to_frame())

    # Test for numerical index
    df = pd.DataFrame({"A": np.random.randn(100)}, index=range(100))
    ddf = dd.from_pandas(df, 10)

    assert_eq(df.index.to_frame(), ddf.index.to_frame())

@pytest.mark.skipif(PANDAS_VERSION < "0.24.0", reason="No renaming for index")
def test_to_frame_name():
    # Test for time index
    df = pd.DataFrame(
        {"A": np.random.randn(100)},
        index=pd.date_range("2011-01-01", freq="H", periods=100),
    )
    ddf = dd.from_pandas(df, 10)

    assert_eq(df.index.to_frame(name="foo"), ddf.index.to_frame(name="foo"))

    # Test for numerical index
    df = pd.DataFrame({"A": np.random.randn(100)}, index=range(100))
    ddf = dd.from_pandas(df, 10)

    assert_eq(df.index.to_frame(name="bar"), ddf.index.to_frame(name="bar"))

@pytest.mark.parametrize("indexer", [0, [0], [0, 1], [1, 0], [False, True, True]])
def test_iloc(indexer):
    df = pd.DataFrame({"A": [1, 2], "B": [3, 4], "C": [5, 6]})
    ddf = dd.from_pandas(df, 2)

    result = ddf.iloc[:, indexer]
    expected = df.iloc[:, indexer]

    assert_eq(result, expected)

def test_iloc_series():
    s = pd.Series([1, 2, 3])
    ds = dd.from_pandas(s, 2)
    with pytest.raises(AttributeError):
        ds.iloc[:]

def test_iloc_raises():
    df = pd.DataFrame({"A": [1, 2], "B": [3, 4], "C": [5, 6]})
    ddf = dd.from_pandas(df, 2)

    with pytest.raises(NotImplementedError):
        ddf.iloc[[0, 1], :]

    with pytest.raises(NotImplementedError):
        ddf.iloc[[0, 1], [0, 1]]

    with pytest.raises(ValueError):
        ddf.iloc[[0, 1], [0, 1], [1, 2]]

    with pytest.raises(IndexError):
        ddf.iloc[:, [5, 6]]

def test_iloc_duplicate_columns():
    df = pd.DataFrame({"A": [1, 2], "B": [3, 4], "C": [5, 6]})
    ddf = dd.from_pandas(df, 2)
    df.columns = ["A", "A", "C"]
    ddf.columns = ["A", "A", "C"]

    selection = ddf.iloc[:, 2]
    # Check that `iloc` is called instead of getitem
    assert any([key.startswith("iloc") for key in selection.dask.layers.keys()])

    select_first = ddf.iloc[:, 1]
    assert_eq(select_first, df.iloc[:, 1])

    select_zeroth = ddf.iloc[:, 0]
    assert_eq(select_zeroth, df.iloc[:, 0])

    select_list_cols = ddf.iloc[:, [0, 2]]
    assert_eq(select_list_cols, df.iloc[:, [0, 2]])

    select_negative = ddf.iloc[:, -1:-3:-1]
    assert_eq(select_negative, df.iloc[:, -1:-3:-1])

def test_iloc_dispatch_to_getitem():
    df = pd.DataFrame({"A": [1, 2], "B": [3, 4], "C": [5, 6]})
    ddf = dd.from_pandas(df, 2)

    selection = ddf.iloc[:, 2]

    assert all([not key.startswith("iloc") for key in selection.dask.layers.keys()])
    assert any([key.startswith("getitem") for key in selection.dask.layers.keys()])

    select_first = ddf.iloc[:, 1]
    assert_eq(select_first, df.iloc[:, 1])

    select_zeroth = ddf.iloc[:, 0]
    assert_eq(select_zeroth, df.iloc[:, 0])

    select_list_cols = ddf.iloc[:, [0, 2]]
    assert_eq(select_list_cols, df.iloc[:, [0, 2]])

    select_negative = ddf.iloc[:, -1:-3:-1]
    assert_eq(select_negative, df.iloc[:, -1:-3:-1])

def test_iloc_out_of_order_selection():
    df = pd.DataFrame({"A": [1] * 100, "B": [2] * 100, "C": [3] * 100, "D": [4] * 100})
    ddf = dd.from_pandas(df, 2)
    ddf = ddf[["C", "A", "B"]]
    a = ddf.iloc[:, 0]
    b = ddf.iloc[:, 1]
    c = ddf.iloc[:, 2]

    assert a.name == "C"
    assert b.name == "A"
    assert c.name == "B"

    a1, b1, c1 = dask.compute(a, b, c)

    assert a1.name == "C"
    assert b1.name == "A"
    assert c1.name == "B"


# coding: utf-8
"""Interface to sourmash"""

import os
import numpy as np
import pandas as pd
import shutil

import anvio
import anvio.utils as utils
import anvio.terminal as terminal
import anvio.filesnpaths as filesnpaths

from scipy.stats import entropy, skew, kurtosis

from anvio.errors import ConfigError

__author__ = "Developers of anvi'o (see AUTHORS.txt)"
__copyright__ = "Copyleft 2015-2019, the Meren Lab (http://merenlab.org/)"
__credits__ = []
__license__ = "GPL 3.0"
__version__ = anvio.__version__
__maintainer__ = "Mahmoud Yousef"
__email__ = "mahmoudyousef@uchicago.edu"

class Sourmash:
    """This calculates a single kmer signature, and computes similarities.

    Feel free to buff this to suit your needs
    """

    def __init__(self, args={}, run=terminal.Run(), progress=terminal.Progress(), program_name='sourmash'):
        self.run = run
        self.progress = progress
        self.program_name = program_name
        self.check_program()

        self.results = {}

        A = lambda x: args.__dict__[x] if x in args.__dict__ else None
        self.log_file_path = os.path.abspath(A('log_file') or filesnpaths.get_temp_file_path())
        self.num_threads = A('num_threads') or 1
        self.kmer_size = A('kmer_size') or 51
        self.scale = A('scale') or 1000

        self.run.warning("Anvi'o will use 'sourmash' by Brown et al. (DOI: 10.21105/joss.00027) to compute kmer sequences and determine mash distances. "
                        "If you publish your findings, please do not forget to properly credit their work",
                         lc='green', header="CITATION")

        if self.num_threads != 1:
            self.num_threads = 1
            self.run.warning("Anvi'o speaking: sourmash currently doesn't support multithreading. "
                            "Anvi'o will have to reduce your number of threads to one :(")

        self.run.info('[sourmash] Log file path', self.log_file_path, nl_after=1)

    def check_program(self):
        utils.is_program_exists(self.program_name)

    def process(self, input_path, fasta_files):
        self.run.info('[sourmash] Kmer size', self.kmer_size, nl_before=1)
        self.run.info('[sourmash] Compression ratio', self.scale)

        report_name = 'kmer_%d_mash_similarity' % self.kmer_size

        # backup the old working directory before changing the directory
        old_wd = os.getcwd()
        os.chdir(input_path)
        if not os.path.exists('output'):
            os.mkdir('output')
        else:
            pass

        self.progress.new('Sourmash')
        self.progress.update('Computing fasta signatures for kmer=%d, scale=%d' % (self.kmer_size, self.scale))

        scale = '--scaled=%i' % self.scale
        compute_command = [self.program_name, 'compute',
                           '-k', self.kmer_size,
                           '-f', scale]
        compute_command.extend(fasta_files)

        exit_code = utils.run_command(compute_command, self.log_file_path, remove_log_file_if_exists=False)
        if int(exit_code):
            self.progress.end()
            raise ConfigError("sourmash returned with non-zero exit code, there may be some errors. "
                             "Please check the log file `%s` for details. Offending command: "
                             "`%s` ..." % (self.log_file_path, ' '.join([str(x) for x in compute_command[:7]])))

        self.progress.update('Computing similarity matrix for kmer=%d, scale=%d' % (self.kmer_size, self.scale))
        compare_command = [self.program_name, 'compare',
                           '-k', self.kmer_size,
                           '--csv', os.path.join('output', report_name + '.txt')]
        for f in fasta_files:
            compare_command.append(f + ".sig")

        exit_code = utils.run_command(compare_command, self.log_file_path, remove_log_file_if_exists=False)
        if int(exit_code):
            self.progress.end()
            raise ConfigError("sourmash returned with non-zero exit code, there may be some errors. "
                             "Please check the log file `%s` for details. Offending command: "
                             "`%s` ..." % (self.log_file_path, ' '.join([str(x) for x in compute_command[:7]])))

        self.results[report_name] = utils.get_TAB_delimited_file_as_dictionary(os.path.join('output', report_name + '.txt'),
                                                                               indexing_field=-1,
                                                                               separator=',')

        self.progress.end()

        # restore old working directory
        os.chdir(old_wd)

        return self.results


"""Build a sentiment analysis / polarity model

Sentiment analysis can be casted as a binary text classification problem,
that is fitting a linear classifier on features extracted from the text
of the user messages so as to guess wether the opinion of the author is
positive or negative.

In this examples we will use a movie review dataset.

"""
# Author: Olivier Grisel <olivier.grisel@ensta.org>
# License: Simplified BSD

import sys
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.svm import LinearSVC
from sklearn.pipeline import Pipeline
from sklearn.grid_search import GridSearchCV
from sklearn.datasets import load_files
from sklearn.cross_validation import train_test_split
from sklearn import metrics

if __name__ == "__main__":
    # NOTE: we put the following in a 'if __name__ == "__main__"' protected
    # block to be able to use a multi-core grid search that also works under
    # Windows, see: http://docs.python.org/library/multiprocessing.html#windows
    # The multiprocessing module is used as the backend of joblib.Parallel
    # that is used when n_jobs != 1 in GridSearchCV

    # the training data folder must be passed as first argument
    movie_reviews_data_folder = sys.argv[1]
    dataset = load_files(movie_reviews_data_folder, shuffle=False)
    print("n_samples: %d" % len(dataset.data))

    # split the dataset in training and test set:
    docs_train, docs_test, y_train, y_test = train_test_split(
        dataset.data, dataset.target, test_size=0.25, random_state=None)

    # TASK: Build a vectorizer / classifier pipeline that filters out tokens
    # that are too rare or too frequent
    pipeline = Pipeline([
        ('vect', TfidfVectorizer(min_df=3, max_df=0.95)),
        ('clf', LinearSVC(C=1000)),
    ])

    # TASK: Build a grid search to find out whether unigrams or bigrams are
    # more useful.
    # Fit the pipeline on the training set using grid search for the parameters
    parameters = {
        'vect__ngram_range': [(1, 1), (1, 2)],
    }
    grid_search = GridSearchCV(pipeline, parameters, n_jobs=-1)
    grid_search.fit(docs_train, y_train)

    # TASK: print the cross-validated scores for the each parameters set
    # explored by the grid search
    print(grid_search.grid_scores_)

    # TASK: Predict the outcome on the testing set and store it in a variable
    # named y_predicted
    y_predicted = grid_search.predict(docs_test)

    # Print the classification report
    print(metrics.classification_report(y_test, y_predicted,
                                        target_names=dataset.target_names))

    # Print and plot the confusion matrix
    cm = metrics.confusion_matrix(y_test, y_predicted)
    print(cm)

    # import matplotlib.pyplot as plt
    # plt.matshow(cm)
    # plt.show()


import nose

from datetime import datetime
from numpy.random import randn
import numpy as np

from pandas.core.api import Series, DataFrame, date_range
import pandas.util.testing as tm
import pandas.stats.math as pmath
from pandas import ols

N, K = 100, 10

_have_statsmodels = True
try:
    import statsmodels.api as sm
except ImportError:
    try:
        import scikits.statsmodels.api as sm  # noqa
    except ImportError:
        _have_statsmodels = False

class TestMath(tm.TestCase):

    _nan_locs = np.arange(20, 40)
    _inf_locs = np.array([])

    def setUp(self):
        arr = randn(N)
        arr[self._nan_locs] = np.NaN

        self.arr = arr
        self.rng = date_range(datetime(2009, 1, 1), periods=N)

        self.series = Series(arr.copy(), index=self.rng)

        self.frame = DataFrame(randn(N, K), index=self.rng,
                               columns=np.arange(K))

    def test_rank_1d(self):
        self.assertEqual(1, pmath.rank(self.series))
        self.assertEqual(0, pmath.rank(Series(0, self.series.index)))

    def test_solve_rect(self):
        if not _have_statsmodels:
            raise nose.SkipTest("no statsmodels")

        b = Series(np.random.randn(N), self.frame.index)
        result = pmath.solve(self.frame, b)
        with tm.assert_produces_warning(FutureWarning, check_stacklevel=False):
            expected = ols(y=b, x=self.frame, intercept=False).beta
        self.assertTrue(np.allclose(result, expected))

    def test_inv_illformed(self):
        singular = DataFrame(np.array([[1, 1], [2, 2]]))
        rs = pmath.inv(singular)
        expected = np.array([[0.1, 0.2], [0.1, 0.2]])
        self.assertTrue(np.allclose(rs, expected))

if __name__ == '__main__':
    nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],
                   exit=False)