"""
=====================
Lasso and Elastic Net
=====================

Lasso and elastic net (L1 and L2 penalisation) implemented using a
coordinate descent.

The coefficients can be forced to be positive.
"""
print(__doc__)

# Author: Alexandre Gramfort <alexandre.gramfort@inria.fr>
# License: BSD 3 clause

from itertools import cycle
import numpy as np
import matplotlib.pyplot as plt

from sklearn.linear_model import lasso_path, enet_path
from sklearn import datasets

diabetes = datasets.load_diabetes()
X = diabetes.data
y = diabetes.target

X /= X.std(axis=0)  # Standardize data (easier to set the l1_ratio parameter)

# Compute paths

eps = 5e-3  # the smaller it is the longer is the path

print("Computing regularization path using the lasso...")
alphas_lasso, coefs_lasso, _ = lasso_path(X, y, eps, fit_intercept=False)

print("Computing regularization path using the positive lasso...")
alphas_positive_lasso, coefs_positive_lasso, _ = lasso_path(
    X, y, eps, positive=True, fit_intercept=False)
print("Computing regularization path using the elastic net...")
alphas_enet, coefs_enet, _ = enet_path(
    X, y, eps=eps, l1_ratio=0.8, fit_intercept=False)

print("Computing regularization path using the positve elastic net...")
alphas_positive_enet, coefs_positive_enet, _ = enet_path(
    X, y, eps=eps, l1_ratio=0.8, positive=True, fit_intercept=False)

# Display results

plt.figure(1)
ax = plt.gca()

colors = cycle(['b', 'r', 'g', 'c', 'k'])
neg_log_alphas_lasso = -np.log10(alphas_lasso)
neg_log_alphas_enet = -np.log10(alphas_enet)
for coef_l, coef_e, c in zip(coefs_lasso, coefs_enet, colors):
    l1 = plt.plot(neg_log_alphas_lasso, coef_l, c=c)
    l2 = plt.plot(neg_log_alphas_enet, coef_e, linestyle='--', c=c)

plt.xlabel('-Log(alpha)')
plt.ylabel('coefficients')
plt.title('Lasso and Elastic-Net Paths')
plt.legend((l1[-1], l2[-1]), ('Lasso', 'Elastic-Net'), loc='lower left')
plt.axis('tight')

plt.figure(2)
ax = plt.gca()
neg_log_alphas_positive_lasso = -np.log10(alphas_positive_lasso)
for coef_l, coef_pl, c in zip(coefs_lasso, coefs_positive_lasso, colors):
    l1 = plt.plot(neg_log_alphas_lasso, coef_l, c=c)
    l2 = plt.plot(neg_log_alphas_positive_lasso, coef_pl, linestyle='--', c=c)

plt.xlabel('-Log(alpha)')
plt.ylabel('coefficients')
plt.title('Lasso and positive Lasso')
plt.legend((l1[-1], l2[-1]), ('Lasso', 'positive Lasso'), loc='lower left')
plt.axis('tight')

plt.figure(3)
ax = plt.gca()
neg_log_alphas_positive_enet = -np.log10(alphas_positive_enet)
for (coef_e, coef_pe, c) in zip(coefs_enet, coefs_positive_enet, colors):
    l1 = plt.plot(neg_log_alphas_enet, coef_e, c=c)
    l2 = plt.plot(neg_log_alphas_positive_enet, coef_pe, linestyle='--', c=c)

plt.xlabel('-Log(alpha)')
plt.ylabel('coefficients')
plt.title('Elastic-Net and positive Elastic-Net')
plt.legend((l1[-1], l2[-1]), ('Elastic-Net', 'positive Elastic-Net'),
           loc='lower left')
plt.axis('tight')
plt.show()


"""Fixture module to skip the datasets loading when offline

Mock urllib2 access to mldata.org and create a temporary data folder.
"""

from os import makedirs
from os.path import join
import numpy as np
import tempfile
import shutil

from sklearn import datasets
from sklearn.utils.testing import install_mldata_mock
from sklearn.utils.testing import uninstall_mldata_mock

def globs(globs):
    # Create a temporary folder for the data fetcher
    global custom_data_home
    custom_data_home = tempfile.mkdtemp()
    makedirs(join(custom_data_home, 'mldata'))
    globs['custom_data_home'] = custom_data_home
    return globs

def setup_module():
    # setup mock urllib2 module to avoid downloading from mldata.org
    install_mldata_mock({
        'mnist-original': {
            'data': np.empty((70000, 784)),
            'label': np.repeat(np.arange(10, dtype='d'), 7000),
        },
        'iris': {
            'data': np.empty((150, 4)),
        },
        'datasets-uci-iris': {
            'double0': np.empty((150, 4)),
            'class': np.empty((150,)),
        },
    })

def teardown_module():
    uninstall_mldata_mock()
    shutil.rmtree(custom_data_home)


"""
worker classes

@author: Jev Kuznetsov
Licence: GPL v2
"""

__docformat__ = 'restructuredtext'

import os
import logger as logger
import yahooFinance as yahoo
from functions import returns, rank
from datetime import date
from pandas import DataFrame, Series
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

class Symbol(object):
    '''
    Symbol class, the foundation of Trading With Python library,
    This class acts as an interface to Yahoo data, Interactive Brokers etc
    '''
    def __init__(self,name):
        self.name = name
        self.log = logger.getLogger(self.name)
        self.log.debug('class created.')

        self.dataDir = os.getenv("USERPROFILE")+'\\twpData\\symbols\\'+self.name
        self.log.debug('Data dir:'+self.dataDir)
        self.ohlc = None # historic OHLC data

    def downloadHistData(self, startDate=(2010,1,1),endDate=date.today().timetuple()[:3],\
                    source = 'yahoo'):
        '''
        get historical OHLC data from a data source (yahoo is default)
        startDate and endDate are tuples in form (d,m,y)
        '''
        self.log.debug('Getting OHLC data')
        self.ohlc = yahoo.getHistoricData(self.name,startDate,endDate)

    def histData(self,column='adj_close'):
        '''
        Return a column of historic data.

        Returns
        -------------
        df : DataFrame
        '''
        s = self.ohlc[column]
        return DataFrame(s.values,s.index,[self.name])

    @property
    def dayReturns(self):
        ''' close-close returns '''
        return (self.ohlc['adj_close']/self.ohlc['adj_close'].shift(1)-1)
        #return DataFrame(s.values,s.index,[self.name])

class Portfolio(object):
    def __init__(self,histPrice,name=''):
        """
        Constructor

        Parameters
        ----------
        histPrice : historic price

        """
        self.histPrice = histPrice
        self.params = DataFrame(index=self.symbols)
        self.params['capital'] = 100*np.ones(self.histPrice.shape[1],dtype=np.float)
        self.params['last'] = self.histPrice.tail(1).T.ix[:,0]
        self.params['shares'] = self.params['capital']/self.params['last']
        self.name= name

    def setHistPrice(self,histPrice):
        self.histPrice = histPrice

    def setShares(self,shares):
        """ set number of shares, adjust capital
        shares: list, np array or Series
        """

        if len(shares) != self.histPrice.shape[1]:
            raise AttributeError('Wrong size of shares vector.')
        self.params['shares'] = shares
        self.params['capital'] = self.params['shares']*self.params['last']

    def setCapital(self,capital):
        """ Set target captial, adjust number of shares """
        if len(capital) != self.histPrice.shape[1]:
            raise AttributeError('Wrong size of shares vector.')
        self.params['capital'] = capital
        self.params['shares'] = self.params['capital']/self.params['last']

    def calculateStatistics(self,other=None):
        ''' calculate spread statistics, save internally '''
        res = {}
        res['micro'] = rank(self.returns[-1],self.returns)
        res['macro'] = rank(self.value[-1], self.value)

        res['last'] = self.value[-1]

        if other is not None:
            res['corr'] = self.returns.corr(returns(other))

        return   Series(res,name=self.name)

    @property
    def symbols(self):
        return self.histPrice.columns.tolist()

    @property
    def returns(self):
        return (returns(self.histPrice)*self.params['capital']).sum(axis=1)

    @property
    def value(self):
        return (self.histPrice*self.params['shares']).sum(axis=1)

    def __repr__(self):
        return ("Portfolio %s \n" % self.name ) + str(self.params)
        #return ('Spread %s :' % self.name ) + str.join(',',
        #        ['%s*%.2f' % t  for t in zip(self.symbols,self.capital)])

class Spread(object):
    '''
    Spread class, used to build a spread out of two symbols.
    '''

    def __init__(self,stock,hedge,beta=None):
        ''' init with symbols or price series '''

        if isinstance(stock,str) and isinstance(hedge,str):
            self.symbols = [stock,hedge]
            self._getYahooData()
        elif isinstance(stock,pd.Series) and isinstance(hedge,pd.Series):
            self.symbols = [stock.name,hedge.name]
            self.price = pd.DataFrame(dict(zip(self.symbols,[stock,hedge]))).dropna()
        else:
            raise ValueError('Both stock and hedge should be of the same type, symbol string or Series')

        # calculate returns
        self.returns = self.price.pct_change()

        if beta is not None:
            self.beta = beta
        else:
            self.estimateBeta()

        # set data
        self.data = pd.DataFrame(index = self.symbols)
        self.data['beta'] = pd.Series({self.symbols[0]:1., self.symbols[1]:-self.beta})

    def calculateShares(self,bet):
        ''' set number of shares based on last quote '''
        if 'price' not in self.data.columns:
            print 'Getting quote...'
            self.getQuote()
        self.data['shares'] = bet*self.data['beta']/self.data['price']

    def estimateBeta(self,plotOn=False):
        """ linear estimation of beta """
        x = self.returns[self.symbols[1]] # hedge
        y = self.returns[self.symbols[0]] # stock

        #avoid extremes
        low = np.percentile(x,20)
        high = np.percentile(x,80)
        iValid = (x>low) & (x<high)

        x = x[iValid]
        y = y[iValid]

        if plotOn:
            plt.plot(x,y,'o')
            plt.grid(True)

        iteration = 1
        nrOutliers = 1
        while iteration < 3 and nrOutliers > 0 :

            (a,b) = np.polyfit(x,y,1)
            yf = np.polyval([a,b],x)

            err = yf-y
            idxOutlier = abs(err) > 3*np.std(err)
            nrOutliers =sum(idxOutlier)
            beta = a
            #print 'Iteration: %i beta: %.2f outliers: %i' % (iteration,beta, nrOutliers)
            x = x[~idxOutlier]
            y = y[~idxOutlier]
            iteration += 1

        if plotOn:
            yf = x*beta
            plt.plot(x,yf,'-',color='red')
            plt.xlabel(self.symbols[1])
            plt.ylabel(self.symbols[0])

        self.beta = beta
        return beta

    @property
    def spread(self):
        ''' return daily returns of the pair '''
        return (self.returns*self.data['beta']).sum(1)

    def getQuote(self):
        ''' get current quote from yahoo '''
        q = yahoo.getQuote(self.symbols)
        self.data['price'] = q['last']

    def _getYahooData(self, startDate=(2007,1,1)):
        """ fetch historic data """
        data = {}
        for symbol in self.symbols:
            print 'Downloading %s' % symbol
            data[symbol]=(yahoo.getHistoricData(symbol,sDate=startDate)['adj_close'] )

        self.price = pd.DataFrame(data).dropna()

    def __repr__(self):
        return 'Spread 1*%s & %.2f*%s ' % (self.symbols[0],-self.beta,self.symbols[1])

    @property
    def name(self):
        return str.join('_',self.symbols)

if __name__=='__main__':

    s = Spread(['SPY','IWM'])


"""Partial dependence plots for tree ensembles. """

# Authors: Peter Prettenhofer
# License: BSD 3 clause

from itertools import count
import numbers

import numpy as np
from scipy.stats.mstats import mquantiles

from ..utils.extmath import cartesian
from ..externals.joblib import Parallel, delayed
from ..externals import six
from ..externals.six.moves import map, range, zip
from ..utils import check_array
from ..tree._tree import DTYPE

from ._gradient_boosting import _partial_dependence_tree
from .gradient_boosting import BaseGradientBoosting

def _grid_from_X(X, percentiles=(0.05, 0.95), grid_resolution=100):
    """Generate a grid of points based on the ``percentiles of ``X``.

    The grid is generated by placing ``grid_resolution`` equally
    spaced points between the ``percentiles`` of each column
    of ``X``.

    Parameters
    ----------
    X : ndarray
        The data
    percentiles : tuple of floats
        The percentiles which are used to construct the extreme
        values of the grid axes.
    grid_resolution : int
        The number of equally spaced points that are placed
        on the grid.

    Returns
    -------
    grid : ndarray
        All data points on the grid; ``grid.shape[1] == X.shape[1]``
        and ``grid.shape[0] == grid_resolution * X.shape[1]``.
    axes : seq of ndarray
        The axes with which the grid has been created.
    """
    if len(percentiles) != 2:
        raise ValueError('percentile must be tuple of len 2')
    if not all(0. <= x <= 1. for x in percentiles):
        raise ValueError('percentile values must be in [0, 1]')

    axes = []
    for col in range(X.shape[1]):
        uniques = np.unique(X[:, col])
        if uniques.shape[0] < grid_resolution:
            # feature has low resolution use unique vals
            axis = uniques
        else:
            emp_percentiles = mquantiles(X, prob=percentiles, axis=0)
            # create axis based on percentiles and grid resolution
            axis = np.linspace(emp_percentiles[0, col],
                               emp_percentiles[1, col],
                               num=grid_resolution, endpoint=True)
        axes.append(axis)

    return cartesian(axes), axes

def partial_dependence(gbrt, target_variables, grid=None, X=None,
                       percentiles=(0.05, 0.95), grid_resolution=100):
    """Partial dependence of ``target_variables``.

    Partial dependence plots show the dependence between the joint values
    of the ``target_variables`` and the function represented
    by the ``gbrt``.

    Read more in the :ref:`User Guide <partial_dependence>`.

    Parameters
    ----------
    gbrt : BaseGradientBoosting
        A fitted gradient boosting model.
    target_variables : array-like, dtype=int
        The target features for which the partial dependecy should be
        computed (size should be smaller than 3 for visual renderings).
    grid : array-like, shape=(n_points, len(target_variables))
        The grid of ``target_variables`` values for which the
        partial dependecy should be evaluated (either ``grid`` or ``X``
        must be specified).
    X : array-like, shape=(n_samples, n_features)
        The data on which ``gbrt`` was trained. It is used to generate
        a ``grid`` for the ``target_variables``. The ``grid`` comprises
        ``grid_resolution`` equally spaced points between the two
        ``percentiles``.
    percentiles : (low, high), default=(0.05, 0.95)
        The lower and upper percentile used create the extreme values
        for the ``grid``. Only if ``X`` is not None.
    grid_resolution : int, default=100
        The number of equally spaced points on the ``grid``.

    Returns
    -------
    pdp : array, shape=(n_classes, n_points)
        The partial dependence function evaluated on the ``grid``.
        For regression and binary classification ``n_classes==1``.
    axes : seq of ndarray or None
        The axes with which the grid has been created or None if
        the grid has been given.

    Examples
    --------
    >>> samples = [[0, 0, 2], [1, 0, 0]]
    >>> labels = [0, 1]
    >>> from sklearn.ensemble import GradientBoostingClassifier
    >>> gb = GradientBoostingClassifier(random_state=0).fit(samples, labels)
    >>> kwargs = dict(X=samples, percentiles=(0, 1), grid_resolution=2)
    >>> partial_dependence(gb, [0], **kwargs) # doctest: +SKIP
    (array([[-4.52...,  4.52...]]), [array([ 0.,  1.])])
    """
    if not isinstance(gbrt, BaseGradientBoosting):
        raise ValueError('gbrt has to be an instance of BaseGradientBoosting')
    if gbrt.estimators_.shape[0] == 0:
        raise ValueError('Call %s.fit before partial_dependence' %
                         gbrt.__class__.__name__)
    if (grid is None and X is None) or (grid is not None and X is not None):
        raise ValueError('Either grid or X must be specified')

    target_variables = np.asarray(target_variables, dtype=np.int32,
                                  order='C').ravel()

    if any([not (0 <= fx < gbrt.n_features) for fx in target_variables]):
        raise ValueError('target_variables must be in [0, %d]'
                         % (gbrt.n_features - 1))

    if X is not None:
        X = check_array(X, dtype=DTYPE, order='C')
        grid, axes = _grid_from_X(X[:, target_variables], percentiles,
                                  grid_resolution)
    else:
        assert grid is not None
        # dont return axes if grid is given
        axes = None
        # grid must be 2d
        if grid.ndim == 1:
            grid = grid[:, np.newaxis]
        if grid.ndim != 2:
            raise ValueError('grid must be 2d but is %dd' % grid.ndim)

    grid = np.asarray(grid, dtype=DTYPE, order='C')
    assert grid.shape[1] == target_variables.shape[0]

    n_trees_per_stage = gbrt.estimators_.shape[1]
    n_estimators = gbrt.estimators_.shape[0]
    pdp = np.zeros((n_trees_per_stage, grid.shape[0],), dtype=np.float64,
                   order='C')
    for stage in range(n_estimators):
        for k in range(n_trees_per_stage):
            tree = gbrt.estimators_[stage, k].tree_
            _partial_dependence_tree(tree, grid, target_variables,
                                     gbrt.learning_rate, pdp[k])

    return pdp, axes

def plot_partial_dependence(gbrt, X, features, feature_names=None,
                            label=None, n_cols=3, grid_resolution=100,
                            percentiles=(0.05, 0.95), n_jobs=1,
                            verbose=0, ax=None, line_kw=None,
                            contour_kw=None, **fig_kw):
    """Partial dependence plots for ``features``.

    The ``len(features)`` plots are arranged in a grid with ``n_cols``
    columns. Two-way partial dependence plots are plotted as contour
    plots.

    Read more in the :ref:`User Guide <partial_dependence>`.

    Parameters
    ----------
    gbrt : BaseGradientBoosting
        A fitted gradient boosting model.
    X : array-like, shape=(n_samples, n_features)
        The data on which ``gbrt`` was trained.
    features : seq of tuples or ints
        If seq[i] is an int or a tuple with one int value, a one-way
        PDP is created; if seq[i] is a tuple of two ints, a two-way
        PDP is created.
    feature_names : seq of str
        Name of each feature; feature_names[i] holds
        the name of the feature with index i.
    label : object
        The class label for which the PDPs should be computed.
        Only if gbrt is a multi-class model. Must be in ``gbrt.classes_``.
    n_cols : int
        The number of columns in the grid plot (default: 3).
    percentiles : (low, high), default=(0.05, 0.95)
        The lower and upper percentile used to create the extreme values
        for the PDP axes.
    grid_resolution : int, default=100
        The number of equally spaced points on the axes.
    n_jobs : int
        The number of CPUs to use to compute the PDs. -1 means 'all CPUs'.
        Defaults to 1.
    verbose : int
        Verbose output during PD computations. Defaults to 0.
    ax : Matplotlib axis object, default None
        An axis object onto which the plots will be drawn.
    line_kw : dict
        Dict with keywords passed to the ``pylab.plot`` call.
        For one-way partial dependence plots.
    contour_kw : dict
        Dict with keywords passed to the ``pylab.plot`` call.
        For two-way partial dependence plots.
    fig_kw : dict
        Dict with keywords passed to the figure() call.
        Note that all keywords not recognized above will be automatically
        included here.

    Returns
    -------
    fig : figure
        The Matplotlib Figure object.
    axs : seq of Axis objects
        A seq of Axis objects, one for each subplot.

    Examples
    --------
    >>> from sklearn.datasets import make_friedman1
    >>> from sklearn.ensemble import GradientBoostingRegressor
    >>> X, y = make_friedman1()
    >>> clf = GradientBoostingRegressor(n_estimators=10).fit(X, y)
    >>> fig, axs = plot_partial_dependence(clf, X, [0, (0, 1)]) #doctest: +SKIP
    ...
    """
    import matplotlib.pyplot as plt
    from matplotlib import transforms
    from matplotlib.ticker import MaxNLocator
    from matplotlib.ticker import ScalarFormatter

    if not isinstance(gbrt, BaseGradientBoosting):
        raise ValueError('gbrt has to be an instance of BaseGradientBoosting')
    if gbrt.estimators_.shape[0] == 0:
        raise ValueError('Call %s.fit before partial_dependence' %
                         gbrt.__class__.__name__)

    # set label_idx for multi-class GBRT
    if hasattr(gbrt, 'classes_') and np.size(gbrt.classes_) > 2:
        if label is None:
            raise ValueError('label is not given for multi-class PDP')
        label_idx = np.searchsorted(gbrt.classes_, label)
        if gbrt.classes_[label_idx] != label:
            raise ValueError('label %s not in ``gbrt.classes_``' % str(label))
    else:
        # regression and binary classification
        label_idx = 0

    X = check_array(X, dtype=DTYPE, order='C')
    if gbrt.n_features != X.shape[1]:
        raise ValueError('X.shape[1] does not match gbrt.n_features')

    if line_kw is None:
        line_kw = {'color': 'green'}
    if contour_kw is None:
        contour_kw = {}

    # convert feature_names to list
    if feature_names is None:
        # if not feature_names use fx indices as name
        feature_names = [str(i) for i in range(gbrt.n_features)]
    elif isinstance(feature_names, np.ndarray):
        feature_names = feature_names.tolist()

    def convert_feature(fx):
        if isinstance(fx, six.string_types):
            try:
                fx = feature_names.index(fx)
            except ValueError:
                raise ValueError('Feature %s not in feature_names' % fx)
        return fx

    # convert features into a seq of int tuples
    tmp_features = []
    for fxs in features:
        if isinstance(fxs, (numbers.Integral,) + six.string_types):
            fxs = (fxs,)
        try:
            fxs = np.array([convert_feature(fx) for fx in fxs], dtype=np.int32)
        except TypeError:
            raise ValueError('features must be either int, str, or tuple '
                             'of int/str')
        if not (1 <= np.size(fxs) <= 2):
            raise ValueError('target features must be either one or two')

        tmp_features.append(fxs)

    features = tmp_features

    names = []
    try:
        for fxs in features:
            l = []
            # explicit loop so "i" is bound for exception below
            for i in fxs:
                l.append(feature_names[i])
            names.append(l)
    except IndexError:
        raise ValueError('features[i] must be in [0, n_features) '
                         'but was %d' % i)

    # compute PD functions
    pd_result = Parallel(n_jobs=n_jobs, verbose=verbose)(
        delayed(partial_dependence)(gbrt, fxs, X=X,
                                    grid_resolution=grid_resolution,
                                    percentiles=percentiles)
        for fxs in features)

    # get global min and max values of PD grouped by plot type
    pdp_lim = {}
    for pdp, axes in pd_result:
        min_pd, max_pd = pdp[label_idx].min(), pdp[label_idx].max()
        n_fx = len(axes)
        old_min_pd, old_max_pd = pdp_lim.get(n_fx, (min_pd, max_pd))
        min_pd = min(min_pd, old_min_pd)
        max_pd = max(max_pd, old_max_pd)
        pdp_lim[n_fx] = (min_pd, max_pd)

    # create contour levels for two-way plots
    if 2 in pdp_lim:
        Z_level = np.linspace(*pdp_lim[2], num=8)

    if ax is None:
        fig = plt.figure(**fig_kw)
    else:
        fig = ax.get_figure()
        fig.clear()

    n_cols = min(n_cols, len(features))
    n_rows = int(np.ceil(len(features) / float(n_cols)))
    axs = []
    for i, fx, name, (pdp, axes) in zip(count(), features, names,
                                        pd_result):
        ax = fig.add_subplot(n_rows, n_cols, i + 1)

        if len(axes) == 1:
            ax.plot(axes[0], pdp[label_idx].ravel(), **line_kw)
        else:
            # make contour plot
            assert len(axes) == 2
            XX, YY = np.meshgrid(axes[0], axes[1])
            Z = pdp[label_idx].reshape(list(map(np.size, axes))).T
            CS = ax.contour(XX, YY, Z, levels=Z_level, linewidths=0.5,
                            colors='k')
            ax.contourf(XX, YY, Z, levels=Z_level, vmax=Z_level[-1],
                        vmin=Z_level[0], alpha=0.75, **contour_kw)
            ax.clabel(CS, fmt='%2.2f', colors='k', fontsize=10, inline=True)

        # plot data deciles + axes labels
        deciles = mquantiles(X[:, fx[0]], prob=np.arange(0.1, 1.0, 0.1))
        trans = transforms.blended_transform_factory(ax.transData,
                                                     ax.transAxes)
        ylim = ax.get_ylim()
        ax.vlines(deciles, [0], 0.05, transform=trans, color='k')
        ax.set_xlabel(name[0])
        ax.set_ylim(ylim)

        # prevent x-axis ticks from overlapping
        ax.xaxis.set_major_locator(MaxNLocator(nbins=6, prune='lower'))
        tick_formatter = ScalarFormatter()
        tick_formatter.set_powerlimits((-3, 4))
        ax.xaxis.set_major_formatter(tick_formatter)

        if len(axes) > 1:
            # two-way PDP - y-axis deciles + labels
            deciles = mquantiles(X[:, fx[1]], prob=np.arange(0.1, 1.0, 0.1))
            trans = transforms.blended_transform_factory(ax.transAxes,
                                                         ax.transData)
            xlim = ax.get_xlim()
            ax.hlines(deciles, [0], 0.05, transform=trans, color='k')
            ax.set_ylabel(name[1])
            # hline erases xlim
            ax.set_xlim(xlim)
        else:
            ax.set_ylabel('Partial dependence')

        if len(axes) == 1:
            ax.set_ylim(pdp_lim[1])
        axs.append(ax)

    fig.subplots_adjust(bottom=0.15, top=0.7, left=0.1, right=0.95, wspace=0.4,
                        hspace=0.3)
    return fig, axs


"""
Django settings for expdj project.
"""

# Build paths inside the project like this: os.path.join(BASE_DIR, ...)
import os
import sys
import tempfile
from datetime import timedelta

import matplotlib
from celery import Celery
from kombu import Exchange, Queue

matplotlib.use('Agg')

BASE_DIR = os.path.dirname(os.path.dirname(__file__))

DOMAIN_NAME = "https://expfactory.org"  # MUST BE HTTPS FOR MECHANICAL TURK
DOMAIN_NAME_HTTP = "http://expfactory.org"  # MUST BE HTTPS FOR MECHANICAL TURK

ADMINS = (('rblair', 'rosswilsonblair@gmail.com'),)

MANAGERS = ADMINS

DEBUG = True
MTURK_ALLOW = False # Allow users to deploy to real Mturk (not just sandbox)
TEMPLATE_DEBUG = False
ALLOWED_HOSTS = ["*"]

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.postgresql_psycopg2',
        'NAME': 'postgres',
        'USER': 'postgres',
        'HOST': 'db',
        'PORT': '5432',
    }
}

# Application definition

INSTALLED_APPS = (
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.sites',
    'django.contrib.sitemaps',
    'django_user_agents',
    'django.contrib.staticfiles',
    'django_extensions',
    'expdj.apps.main',
    'expdj.apps.turk',
    'expdj.apps.experiments',
    'expdj.apps.users',
    'crispy_forms',
    'polymorphic',
    'guardian',
    'dbbackup',
    'djcelery',
    'rest_framework',
    'rest_framework.authtoken',
)

AUTHENTICATION_BACKENDS = (
    'django.contrib.auth.backends.ModelBackend',
    'social.backends.facebook.FacebookOAuth2',
    'social.backends.google.GoogleOAuth2',
    'guardian.backends.ObjectPermissionBackend'
)

SOCIAL_AUTH_PIPELINE = (
    'social.pipeline.social_auth.social_details',
    'social.pipeline.social_auth.social_uid',
    'social.pipeline.social_auth.auth_allowed',
    'social.pipeline.social_auth.social_user',
    'social.pipeline.user.get_username',
    'social.pipeline.social_auth.associate_by_email',  # <--- enable this one
    'social.pipeline.user.create_user',
    'social.pipeline.social_auth.associate_user',
    'social.pipeline.social_auth.load_extra_data',
    'social.pipeline.user.user_details'
)

SOCIAL_AUTH_FACEBOOK_SCOPE = ['email']

MIDDLEWARE_CLASSES = (
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.auth.middleware.SessionAuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
    'django_user_agents.middleware.UserAgentMiddleware',
)

ROOT_URLCONF = 'expdj.urls'

TEMPLATES = [
    {
        'BACKEND': 'django.template.backends.django.DjangoTemplates',
        'DIRS': [os.path.join(BASE_DIR, 'templates')],
        'APP_DIRS': True,
        'OPTIONS': {
            'context_processors': [
                'django.template.context_processors.debug',
                'django.template.context_processors.request',
                'django.contrib.auth.context_processors.auth',
                'django.contrib.messages.context_processors.messages',
            ],
        },
    },
]

TEST_RUNNER = 'django.test.runner.DiscoverRunner'

WSGI_APPLICATION = 'expdj.wsgi.application'

# Internationalization
LANGUAGE_CODE = 'en-us'
TIME_ZONE = 'UTC'
USE_I18N = True
USE_L10N = True
USE_TZ = True
SITE_ID = 1
ANONYMOUS_USER_ID = -1  # django-guardian

# Static files (CSS, JavaScript, Images)
MEDIA_ROOT = './static'
MEDIA_URL = '/static/'
STATIC_ROOT = './assets'
STATIC_URL = '/assets/'
# List of finder classes that know how to find static files in
# various locations.
STATICFILES_FINDERS = (
    'django.contrib.staticfiles.finders.FileSystemFinder',
    'django.contrib.staticfiles.finders.AppDirectoriesFinder',
)

SESSION_SERIALIZER = 'django.contrib.sessions.serializers.PickleSerializer'

SENDFILE_BACKEND = 'sendfile.backends.development'
PRIVATE_MEDIA_REDIRECT_HEADER = 'X-Accel-Redirect'
CRISPY_TEMPLATE_PACK = 'bootstrap3'

CACHES = {
    'default': {
        'BACKEND': 'django.core.cache.backends.locmem.LocMemCache',
    }
}

# Celery config
BROKER_URL = 'redis://redis:6379/0'
CELERY_RESULT_BACKEND = 'djcelery.backends.database:DatabaseBackend'
CELERY_ACCEPT_CONTENT = ['json']
CELERY_TASK_SERIALIZER = 'json'
CELERY_RESULT_SERIALIZER = 'json'
CELERY_DEFAULT_QUEUE = 'default'
CELERY_QUEUES = (
    Queue('default', Exchange('default'), routing_key='default'),
)
CELERY_IMPORTS = ('expdj.apps.turk.tasks', )

# here is how to run a task regularly
# CELERYBEAT_SCHEDULE = {
#    'task name': {
#        'task': 'task_name',
#        'schedule': timedelta(days=1)
#    },
# }

CELERY_TIMEZONE = 'Europe/Berlin'

# REST FRAMEWORK
REST_FRAMEWORK = {
    # Use Django's standard `django.contrib.auth` permissions,
    # or allow read-only access for unauthenticated users.
    'DEFAULT_PERMISSION_CLASSES': [
        'rest_framework.permissions.IsAuthenticatedOrReadOnly',
    ],
    'DEFAULT_AUTHENTICATION_CLASSES': (
        'rest_framework.authentication.BasicAuthentication',
        'rest_framework.authentication.SessionAuthentication',
        'rest_framework.authentication.TokenAuthentication',
    ),
    'DEFAULT_PAGINATION_CLASS': 'rest_framework.pagination.LimitOffsetPagination',
    'PAGE_SIZE': 10,
}

CSRF_COOKIE_SECURE = False
SESSION_COOKIE_SECURE = False

# SSL ENABLED
SECURE_PROXY_SSL_HEADER = ('HTTP_X_FORWARDED_PROTOCOL', 'https')

EXP_REPO = os.path.join(BASE_DIR, 'expdj/experiment_repo')

# Bogus secret key.
try:
    from secrets import *
except ImportError:
    from bogus_secrets import *

# Local settings
try:
    from local_settings import *
except ImportError:
    pass


import pandas as pd
from absl import flags
from absl import app

FLAGS = flags.FLAGS
flags.DEFINE_string('preprocess_file',
                    'NPS_1978-2018_Data.tsv',
                    'file path to tsv file with data to proess',
                    short_name='p')

def convert_nan_for_calculation(value):
    if pd.isna(value):
        return 0
    else:
        return value

def total_jurisdiction_columns_helper(df):
    """calculation to include private facility numbers"""
    df["PVINF_Temp"] = df["PVINF"].apply(convert_nan_for_calculation)
    df["PVOTHF_Temp"] = df["PVOTHF"].apply(convert_nan_for_calculation)
    df["PVINM_Temp"] = df["PVINM"].apply(convert_nan_for_calculation)
    df["PVOTHM_Temp"] = df["PVOTHM"].apply(convert_nan_for_calculation)
    df["Female_Total_Temp"] = df[["JURTOTF", "PVINF_Temp", "PVOTHF_Temp"
                                 ]].sum(axis=1).where(df["PVINCLF"] == 2,
                                                      df["JURTOTF"])
    df["Male_Total_Temp"] = df[["JURTOTM", "PVINM_Temp", "PVOTHM_Temp"
                               ]].sum(axis=1).where(df["PVINCLM"] == 2,
                                                    df["JURTOTM"])
    """calculation to include local facility numbers"""
    df["LFF_Temp"] = df["LFF"].apply(convert_nan_for_calculation)
    df["LFM_Temp"] = df["LFM"].apply(convert_nan_for_calculation)
    df["Female_Total_Temp"] = df[["Female_Total_Temp", "LFF_Temp"
                                 ]].sum(axis=1).where(df["LFINCLF"] == 2,
                                                      df["Female_Total_Temp"])
    df["Male_Total_Temp"] = df[["Male_Total_Temp", "LFM_Temp"
                               ]].sum(axis=1).where(df["LFINCLM"] == 2,
                                                    df["Male_Total_Temp"])
    """calculation to include numbers from local facilities solely to ease crowding"""
    df["LFCRSTF_Temp"] = df["LFCRSTF"].apply(convert_nan_for_calculation)
    df["LFCRSTM_Temp"] = df["LFCRSTM"].apply(convert_nan_for_calculation)
    df["Female_Total_Temp"] = df[["Female_Total_Temp", "LFCRSTF_Temp"
                                 ]].sum(axis=1).where(df["LFCRINCF"] == 2,
                                                      df["Female_Total_Temp"])
    df["Male_Total_Temp"] = df[["Male_Total_Temp", "LFCRSTM_Temp"
                               ]].sum(axis=1).where(df["LFCRINCM"] == 2,
                                                    df["Male_Total_Temp"])
    """calculation to include federal and other state facility numbers"""
    df["FEDF_Temp"] = df["FEDF"].apply(convert_nan_for_calculation)
    df["OTHSTF_Temp"] = df["OTHSTF"].apply(convert_nan_for_calculation)
    df["FEDM_Temp"] = df["FEDM"].apply(convert_nan_for_calculation)
    df["OTHSTM_Temp"] = df["OTHSTM"].apply(convert_nan_for_calculation)
    df["Female_Total_Temp"] = df[[
        "Female_Total_Temp", "FEDF_Temp", "OTHSTF_Temp"
    ]].sum(axis=1).where(df["FACINCLF"] == 2, df["Female_Total_Temp"])
    df["Male_Total_Temp"] = df[["Male_Total_Temp", "FEDM_Temp", "OTHSTM_Temp"
                               ]].sum(axis=1).where(df["FACINCLM"] == 2,
                                                    df["Male_Total_Temp"])

def get_columns(df):
    df_out = {}
    total_jurisdiction_columns_helper(df)
    df_out["GeoId"] = df["GeoId"]
    df_out["YEAR"] = df["YEAR"]
    df_out["Count_Person_Female_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
        "Female_Total_Temp"]
    df_out[
        "Count_Person_Female_Incarcerated_WhiteAlone_MeasuredBasedOnJurisdiction"] = df[
            "WHITEF"]
    df_out[
        "Count_Person_BlackOrAfricanAmericanAlone_Female_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "BLACKF"]
    df_out[
        "Count_Person_Female_HispanicOrLatino_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "HISPF"]
    df_out[
        "Count_Person_AmericanIndianOrAlaskaNativeAlone_Female_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "AIANF"]
    df_out[
        "Count_Person_AsianAlone_Female_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "ASIANF"]
    df_out[
        "Count_Person_Female_Incarcerated_NativeHawaiianOrOtherPacificIslanderAlone_MeasuredBasedOnJurisdiction"] = df[
            "NHPIF"]
    df_out[
        "Count_Person_Female_Incarcerated_TwoOrMoreRaces_MeasuredBasedOnJurisdiction"] = df[
            "TWORACEF"]
    df_out[
        "Count_MortalityEvent_Female_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "DTHTOTF"]
    df_out[
        "Count_MortalityEvent_Female_Incarcerated_JudicialExecution_MeasuredBasedOnJurisdiction"] = df[
            "DTHEXECF"]
    df_out[
        "Count_MortalityEvent_Female_IllnessOrNaturalCause_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "DTHILLNF"]
    df_out[
        "Count_MortalityEvent_AIDS_Female_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "DTHAIDSF"]
    df_out[
        "Count_MortalityEvent_Female_Incarcerated_IntentionalSelf-Harm(Suicide)_MeasuredBasedOnJurisdiction"] = df[
            "DTHSUICF"]
    df_out[
        "Count_MortalityEvent_Accidents(UnintentionalInjuries)_Female_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "DTHACCF"]
    df_out[
        "Count_MortalityEvent_DeathDueToAnotherPerson_Female_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "DTHPERSF"]
    df_out[
        "Count_MortalityEvent_Assault(Homicide)_Female_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "DTHHOMIF"]
    df_out[
        "Count_MortalityEvent_Female_Incarcerated_NPSOtherCauseOfDeath_MeasuredBasedOnJurisdiction"] = df[
            "DTHOTHF"]
    df_out[
        "Count_IncarcerationEvent_AdmittedToPrison_Female_Incarcerated_MaxSentenceGreaterThan1Year_Sentenced_MeasuredBasedOnJurisdiction"] = df[
            "ADTOTF"]
    df_out[
        "Count_IncarcerationEvent_Female_Incarcerated_MaxSentenceGreaterThan1Year_ReleasedFromPrison_Sentenced_MeasuredBasedOnJurisdiction"] = df[
            "RLTOTF"]
    df_out[
        "Count_Person_Female_Incarcerated_MaxSentenceGreaterThan1Year_Sentenced_MeasuredBasedOnJurisdiction"] = df[
            "JURGT1F"]
    df_out[
        "Count_Person_Female_Incarcerated_MaxSentence1YearOrLess_Sentenced_MeasuredBasedOnJurisdiction"] = df[
            "JURLT1F"]
    df_out[
        "Count_Person_Female_Incarcerated_Unsentenced_MeasuredBasedOnJurisdiction"] = df[
            "JURUNSF"]
    df_out[
        "Count_Person_Female_Incarcerated_InState_PrivatelyOperated_MeasuredBasedOnJurisdiction"] = df[
            "PVINF"]
    df_out[
        "Count_Person_Female_Incarcerated_OutOfState_PrivatelyOperated_MeasuredBasedOnJurisdiction"] = df[
            "PVOTHF"]
    df_out[
        "Count_Person_Female_Incarcerated_Local_LocallyOperated_MeasuredBasedOnJurisdiction"] = df[
            "LFF"]
    df_out[
        "Count_Person_FederallyOperated_Female_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "FEDF"]
    df_out[
        "Count_Person_Female_Incarcerated_OutOfState_StateOperated_MeasuredBasedOnJurisdiction"] = df[
            "OTHSTF"]
    df_out[
        "Count_Person_Female_Incarcerated_NotAUSCitizen_StateOperated&FederallyOperated&PrivatelyOperated_MeasuredBasedOnCustody"] = df[
            "NCITZTOTF"]
    df_out[
        "Count_Person_Female_Incarcerated_MaxSentenceGreaterThan1Year_NotAUSCitizen_Sentenced_StateOperated&FederallyOperated&PrivatelyOperated_MeasuredBasedOnCustody"] = df[
            "NCITZGT1F"]
    df_out[
        "Count_Person_Female_Incarcerated_MaxSentence1YearOrLess_NotAUSCitizen_Sentenced_StateOperated&FederallyOperated&PrivatelyOperated_MeasuredBasedOnCustody"] = df[
            "NCITZLE1F"]
    df_out[
        "Count_Person_Female_Incarcerated_NotAUSCitizen_StateOperated&FederallyOperated&PrivatelyOperated_Unsentenced_MeasuredBasedOnCustody"] = df[
            "NCITZUNSF"]
    df_out[
        "Count_Person_Female_Incarcerated_Under18_MeasuredBasedOnCustody"] = df[
            "CUSLT18F"]
    df_out["Count_Person_Incarcerated_Male_MeasuredBasedOnJurisdiction"] = df[
        "Male_Total_Temp"]
    df_out[
        "Count_Person_Incarcerated_Male_WhiteAlone_MeasuredBasedOnJurisdiction"] = df[
            "WHITEM"]
    df_out[
        "Count_Person_BlackOrAfricanAmericanAlone_Incarcerated_Male_MeasuredBasedOnJurisdiction"] = df[
            "BLACKM"]
    df_out[
        "Count_Person_HispanicOrLatino_Incarcerated_Male_MeasuredBasedOnJurisdiction"] = df[
            "HISPM"]
    df_out[
        "Count_Person_AmericanIndianOrAlaskaNativeAlone_Incarcerated_Male_MeasuredBasedOnJurisdiction"] = df[
            "AIANM"]
    df_out[
        "Count_Person_AsianAlone_Incarcerated_Male_MeasuredBasedOnJurisdiction"] = df[
            "ASIANM"]
    df_out[
        "Count_Person_Incarcerated_Male_NativeHawaiianOrOtherPacificIslanderAlone_MeasuredBasedOnJurisdiction"] = df[
            "NHPIM"]
    df_out[
        "Count_Person_Incarcerated_Male_TwoOrMoreRaces_MeasuredBasedOnJurisdiction"] = df[
            "TWORACEM"]
    df_out[
        "Count_MortalityEvent_Incarcerated_Male_MeasuredBasedOnJurisdiction"] = df[
            "DTHTOTM"]
    df_out[
        "Count_MortalityEvent_Incarcerated_JudicialExecution_Male_MeasuredBasedOnJurisdiction"] = df[
            "DTHEXECM"]
    df_out[
        "Count_MortalityEvent_IllnessOrNaturalCause_Incarcerated_Male_MeasuredBasedOnJurisdiction"] = df[
            "DTHILLNM"]
    df_out[
        "Count_MortalityEvent_AIDS_Incarcerated_Male_MeasuredBasedOnJurisdiction"] = df[
            "DTHAIDSM"]
    df_out[
        "Count_MortalityEvent_Incarcerated_IntentionalSelf-Harm(Suicide)_Male_MeasuredBasedOnJurisdiction"] = df[
            "DTHSUICM"]
    df_out[
        "Count_MortalityEvent_Accidents(UnintentionalInjuries)_Incarcerated_Male_MeasuredBasedOnJurisdiction"] = df[
            "DTHACCM"]
    df_out[
        "Count_MortalityEvent_DeathDueToAnotherPerson_Incarcerated_Male_MeasuredBasedOnJurisdiction"] = df[
            "DTHPERSM"]
    df_out[
        "Count_MortalityEvent_Assault(Homicide)_Incarcerated_Male_MeasuredBasedOnJurisdiction"] = df[
            "DTHHOMIM"]
    df_out[
        "Count_MortalityEvent_Incarcerated_Male_NPSOtherCauseOfDeath_MeasuredBasedOnJurisdiction"] = df[
            "DTHOTHM"]
    df_out[
        "Count_IncarcerationEvent_AdmittedToPrison_Incarcerated_Male_MaxSentenceGreaterThan1Year_Sentenced_MeasuredBasedOnJurisdiction"] = df[
            "ADTOTM"]
    df_out[
        "Count_IncarcerationEvent_Incarcerated_Male_MaxSentenceGreaterThan1Year_ReleasedFromPrison_Sentenced_MeasuredBasedOnJurisdiction"] = df[
            "RLTOTM"]
    df_out[
        "Count_Person_Incarcerated_Male_MaxSentenceGreaterThan1Year_Sentenced_MeasuredBasedOnJurisdiction"] = df[
            "JURGT1M"]
    df_out[
        "Count_Person_Incarcerated_Male_MaxSentence1YearOrLess_Sentenced_MeasuredBasedOnJurisdiction"] = df[
            "JURLT1M"]
    df_out[
        "Count_Person_Incarcerated_Male_Unsentenced_MeasuredBasedOnJurisdiction"] = df[
            "JURUNSM"]
    df_out[
        "Count_Person_Incarcerated_InState_Male_PrivatelyOperated_MeasuredBasedOnJurisdiction"] = df[
            "PVINM"]
    df_out[
        "Count_Person_Incarcerated_Male_OutOfState_PrivatelyOperated_MeasuredBasedOnJurisdiction"] = df[
            "PVOTHM"]
    df_out[
        "Count_Person_Incarcerated_Local_LocallyOperated_Male_MeasuredBasedOnJurisdiction"] = df[
            "LFM"]
    df_out[
        "Count_Person_FederallyOperated_Incarcerated_Male_MeasuredBasedOnJurisdiction"] = df[
            "FEDM"]
    df_out[
        "Count_Person_Incarcerated_Male_OutOfState_StateOperated_MeasuredBasedOnJurisdiction"] = df[
            "OTHSTM"]
    df_out[
        "Count_Person_Incarcerated_Male_NotAUSCitizen_StateOperated&FederallyOperated&PrivatelyOperated_MeasuredBasedOnCustody"] = df[
            "NCITZTOTM"]
    df_out[
        "Count_Person_Incarcerated_Male_MaxSentenceGreaterThan1Year_NotAUSCitizen_Sentenced_StateOperated&FederallyOperated&PrivatelyOperated_MeasuredBasedOnCustody"] = df[
            "NCITZGT1M"]
    df_out[
        "Count_Person_Incarcerated_Male_MaxSentence1YearOrLess_NotAUSCitizen_Sentenced_StateOperated&FederallyOperated&PrivatelyOperated_MeasuredBasedOnCustody"] = df[
            "NCITZLE1M"]
    df_out[
        "Count_Person_Incarcerated_Male_NotAUSCitizen_StateOperated&FederallyOperated&PrivatelyOperated_Unsentenced_MeasuredBasedOnCustody"] = df[
            "NCITZUNSM"]
    df_out[
        "Count_Person_Incarcerated_Male_Under18_MeasuredBasedOnCustody"] = df[
            "CUSLT18M"]
    df_out["Count_Person_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
        "Female_Total_Temp"] + df["Male_Total_Temp"]
    df_out[
        "Count_Person_Incarcerated_WhiteAlone_MeasuredBasedOnJurisdiction"] = df[
            "WHITEF"] + df["WHITEM"]
    df_out[
        "Count_Person_BlackOrAfricanAmericanAlone_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "BLACKF"] + df["BLACKM"]
    df_out[
        "Count_Person_HispanicOrLatino_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "HISPF"] + df["HISPM"]
    df_out[
        "Count_Person_AmericanIndianOrAlaskaNativeAlone_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "AIANF"] + df["AIANM"]
    df_out[
        "Count_Person_AsianAlone_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "ASIANF"] + df["ASIANM"]
    df_out[
        "Count_Person_Incarcerated_NativeHawaiianOrOtherPacificIslanderAlone_MeasuredBasedOnJurisdiction"] = df[
            "NHPIF"] + df["NHPIM"]
    df_out[
        "Count_Person_Incarcerated_TwoOrMoreRaces_MeasuredBasedOnJurisdiction"] = df[
            "TWORACEF"] + df["TWORACEM"]
    df_out[
        "Count_MortalityEvent_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "DTHTOTF"] + df["DTHTOTM"]
    df_out[
        "Count_MortalityEvent_Incarcerated_JudicialExecution_MeasuredBasedOnJurisdiction"] = df[
            "DTHEXECF"] + df["DTHEXECM"]
    df_out[
        "Count_MortalityEvent_IllnessOrNaturalCause_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "DTHILLNF"] + df["DTHILLNM"]
    df_out[
        "Count_MortalityEvent_AIDS_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "DTHAIDSF"] + df["DTHAIDSM"]
    df_out[
        "Count_MortalityEvent_Incarcerated_IntentionalSelf-Harm(Suicide)_MeasuredBasedOnJurisdiction"] = df[
            "DTHSUICF"] + df["DTHSUICM"]
    df_out[
        "Count_MortalityEvent_Accidents(UnintentionalInjuries)_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "DTHACCF"] + df["DTHACCM"]
    df_out[
        "Count_MortalityEvent_DeathDueToAnotherPerson_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "DTHPERSF"] + df["DTHPERSM"]
    df_out[
        "Count_MortalityEvent_Assault(Homicide)_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "DTHHOMIF"] + df["DTHHOMIM"]
    df_out[
        "Count_MortalityEvent_Incarcerated_NPSOtherCauseOfDeath_MeasuredBasedOnJurisdiction"] = df[
            "DTHOTHF"] + df["DTHOTHM"]
    df_out[
        "Count_IncarcerationEvent_AdmittedToPrison_Incarcerated_MaxSentenceGreaterThan1Year_Sentenced_MeasuredBasedOnJurisdiction"] = df[
            "ADTOTF"] + df["ADTOTM"]
    df_out[
        "Count_IncarcerationEvent_Incarcerated_MaxSentenceGreaterThan1Year_ReleasedFromPrison_Sentenced_MeasuredBasedOnJurisdiction"] = df[
            "RLTOTF"] + df["RLTOTM"]
    df_out[
        "Count_Person_Incarcerated_MaxSentenceGreaterThan1Year_Sentenced_MeasuredBasedOnJurisdiction"] = df[
            "JURGT1F"] + df["JURGT1M"]
    df_out[
        "Count_Person_Incarcerated_MaxSentence1YearOrLess_Sentenced_MeasuredBasedOnJurisdiction"] = df[
            "JURLT1F"] + df["JURLT1M"]
    df_out[
        "Count_Person_Incarcerated_Unsentenced_MeasuredBasedOnJurisdiction"] = df[
            "JURUNSF"] + df["JURUNSM"]
    df_out[
        "Count_Person_Incarcerated_InState_PrivatelyOperated_MeasuredBasedOnJurisdiction"] = df[
            "PVINF"] + df["PVINM"]
    df_out[
        "Count_Person_Incarcerated_OutOfState_PrivatelyOperated_MeasuredBasedOnJurisdiction"] = df[
            "PVOTHF"] + df["PVOTHM"]
    df_out[
        "Count_Person_Incarcerated_Local_LocallyOperated_MeasuredBasedOnJurisdiction"] = df[
            "LFF"] + df["LFM"]
    df_out[
        "Count_Person_FederallyOperated_Incarcerated_MeasuredBasedOnJurisdiction"] = df[
            "FEDF"] + df["FEDM"]
    df_out[
        "Count_Person_Incarcerated_OutOfState_StateOperated_MeasuredBasedOnJurisdiction"] = df[
            "OTHSTF"] + df["OTHSTM"]
    df_out[
        "Count_Person_Incarcerated_NotAUSCitizen_StateOperated&FederallyOperated&PrivatelyOperated_MeasuredBasedOnCustody"] = df[
            "NCITZTOTF"] + df["NCITZTOTM"]
    df_out[
        "Count_Person_Incarcerated_MaxSentenceGreaterThan1Year_NotAUSCitizen_Sentenced_StateOperated&FederallyOperated&PrivatelyOperated_MeasuredBasedOnCustody"] = df[
            "NCITZGT1F"] + df["NCITZGT1M"]
    df_out[
        "Count_Person_Incarcerated_MaxSentence1YearOrLess_NotAUSCitizen_Sentenced_StateOperated&FederallyOperated&PrivatelyOperated_MeasuredBasedOnCustody"] = df[
            "NCITZLE1F"] + df["NCITZLE1M"]
    df_out[
        "Count_Person_Incarcerated_NotAUSCitizen_StateOperated&FederallyOperated&PrivatelyOperated_Unsentenced_MeasuredBasedOnCustody"] = df[
            "NCITZUNSF"] + df["NCITZUNSM"]
    df_out["Count_Person_Incarcerated_Under18_MeasuredBasedOnCustody"] = df[
        "CUSLT18F"] + df["CUSLT18M"]
    return df_out

def convert_geoId(fips_code):
    """Creates geoId column"""
    return 'geoId/' + str(fips_code).zfill(2)

def convert_missing_value_to_nan(value):
    """codes for missing values are always negative and actual data is always >= 0"""
    if isinstance(value, int) and value < 0:
        return float("nan")
    else:
        return value

def convert_nan_to_empty_cell(value):
    if pd.isna(value):
        return ''
    else:
        return value

def preprocess_df(raw_df):
    """cleans raw_df

    Args:
        raw_data: raw data frame to be used as starting point for cleaning
    """
    df = raw_df.copy()
    df['GeoId'] = df['STATEID'].apply(convert_geoId)

    # convert missing values to NaN for aggregation
    for column_name in list(df.columns):
        df[column_name] = df[column_name].apply(convert_missing_value_to_nan)

    #get columns matching stat var names and add aggregate columns
    df_out = pd.DataFrame(get_columns(df))

    #convert NaN to empty cell
    for column_name in list(df_out.columns):
        df_out[column_name] = df_out[column_name].apply(
            convert_nan_to_empty_cell)

    return df_out

def main(args):
    filename = FLAGS.preprocess_file
    print('Processing {0}'.format(filename))
    df = pd.read_csv(filename, delimiter='\t')
    processed_df = preprocess_df(df)
    processed_df.to_csv(filename.replace('.tsv', '_processed.csv'), index=False)
    print('Done processing {0}'.format(filename))

if __name__ == '__main__':
    app.run(main)


# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""Namespace for ops used in imperative programming."""

from ..ndarray import numpy as _mx_nd_np
from ..random import seed

__all__ = ["randint", "uniform", "normal", "choice", "rand", "multinomial", "multivariate_normal",
           "logistic", "gumbel", "f",
           "laplace",
           "shuffle", "randn", "gamma", "beta", "chisquare", "exponential", "lognormal",
           "weibull", "pareto", "power", "rayleigh",
           "seed"]

def randint(low, high=None, size=None, dtype=None, ctx=None, out=None):
    r"""Return random integers from `low` (inclusive) to `high` (exclusive).

    Return random integers from the "discrete uniform" distribution of
    the specified dtype in the "half-open" interval [`low`, `high`). If
    `high` is None (the default), then results are from [0, `low`).

    Parameters
    ----------
    low : int
        Lowest (signed) integer to be drawn from the distribution (unless
        ``high=None``, in which case this parameter is one above the
        *highest* such integer).
    high : int, optional
        If provided, one above the largest (signed) integer to be drawn
        from the distribution (see above for behavior if ``high=None``).
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  Default is None, in which case a
        single value is returned.
    dtype : dtype, optional
        Desired dtype of the result. All dtypes are determined by their
        name, i.e., 'int64', 'int', etc, so byteorder is not available
        and a specific precision may have different C types depending
        on the platform. The default value is 'np.int'.
    ctx : Context, optional
        Device context of output. Default is current context.
    out : ndarray, optional
        The output ndarray (default is `None`).

    Returns
    -------
    out : ndarray of ints
        `size`-shaped array of random integers from the appropriate
        distribution, or a single such random int if `size` not provided.

    Examples
    --------
    >>> np.random.randint(2, size=10)
    array([1, 0, 0, 0, 1, 1, 0, 0, 1, 0])
    >>> np.random.randint(1, size=10)
    array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])

    Generate a 2 x 4 array of ints between 0 and 4, inclusive:

    >>> np.random.randint(5, size=(2, 4))
    array([[4, 0, 2, 1],
        [3, 2, 2, 0]])
    """
    return _mx_nd_np.random.randint(low, high, size, dtype, ctx, out)

def uniform(low=0.0, high=1.0, size=None, dtype=None, ctx=None, out=None):
    r"""Draw samples from a uniform distribution.

    Samples are uniformly distributed over the half-open interval
    ``[low, high)`` (includes low, but excludes high).  In other words,
    any value within the given interval is equally likely to be drawn
    by `uniform`.

    Parameters
    ----------
    low : float, ndarray, optional
        Lower boundary of the output interval.  All values generated will be
        greater than or equal to low.  The default value is 0.
    high : float, ndarray, optional
        Upper boundary of the output interval.  All values generated will be
        less than high.  The default value is 1.0.
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a scalar tensor containing a single value is returned if
        ``low`` and ``high`` are both scalars. Otherwise,
        ``np.broadcast(low, high).size`` samples are drawn.
    dtype : {'float16', 'float32', 'float64'}, optional
        Data type of output samples.
        When npx.is_np_default_dtype() returns False, default dtype is float32;
        When npx.is_np_default_dtype() returns True, default dtype is float64.
    ctx : Context, optional
        Device context of output. Default is current context.

    Returns
    -------
    out : ndarray
        Drawn samples from the parameterized uniform distribution.

    See Also
    --------
    randint : Discrete uniform distribution, yielding integers.
    rand : Convenience function that accepts dimensions as input, e.g.,
           ``rand(2,2)`` would generate a 2-by-2 array of floats,
           uniformly distributed over ``[0, 1)``.

    Notes
    -----
    The probability density function of the uniform distribution is

    .. math:: p(x) = \frac{1}{b - a}

    anywhere within the interval ``[a, b)``, and zero elsewhere.

    When ``high`` == ``low``, values of ``low`` will be returned.
    If ``high`` < ``low``, the results are officially undefined
    and may eventually raise an error, i.e. do not rely on this
    function to behave when passed arguments satisfying that
    inequality condition.
    """
    return _mx_nd_np.random.uniform(low, high, size=size, ctx=ctx, dtype=dtype, out=out)

def normal(loc=0.0, scale=1.0, size=None, dtype=None, ctx=None, out=None):
    r"""Draw random samples from a normal (Gaussian) distribution.

    Samples are distributed according to a normal distribution parametrized
    by *loc* (mean) and *scale* (standard deviation).

    Parameters
    ----------
    loc : float, optional
        Mean (centre) of the distribution.
    scale : float, optional
        Standard deviation (spread or "width") of the distribution.
    size : int or tuple of ints, optional
        Output shape. If the given shape is, e.g., `(m, n, k)`, then `m * n * k`
        samples are drawn. If size is `None` (default), a scalar tensor containing
        a single value is returned if loc and scale are both scalars. Otherwise,
        ``np.broadcast(low, high).size`` samples are drawn.
    dtype : {'float16', 'float32', 'float64'}, optional
        Data type of output samples.
        When npx.is_np_default_dtype() returns False, default dtype is float32;
        When npx.is_np_default_dtype() returns True, default dtype is float64.
    ctx : Context, optional
        Device context of output, default is current context.
    out : ``ndarray``, optional
        Store output to an existing ``ndarray``.

    Returns
    -------
    out : ndarray
        Drawn samples from the parameterized `normal distribution` [1]_.

    Notes
    -----
    The probability density for the Gaussian distribution is

    .. math:: p(x) = \frac{1}{\sqrt{ 2 \pi \sigma^2 }}
                     e^{ - \frac{ (x - \mu)^2 } {2 \sigma^2} },

    where :math:`\mu` is the mean and :math:`\sigma` the standard
    deviation. The square of the standard deviation, :math:`\sigma^2`,
    is called the variance.

    The function has its peak at the mean, and its "spread" increases with
    the standard deviation (the function reaches 0.607 times its maximum at
    :math:`x + \sigma` and :math:`x - \sigma` [2]_).  This implies that
    `numpy.random.normal` is more likely to return samples lying close to
    the mean, rather than those far away.

    References
    ----------
    .. [1] Wikipedia, "Normal distribution",
           https://en.wikipedia.org/wiki/Normal_distribution
    .. [2] P. R. Peebles Jr., "Central Limit Theorem" in "Probability,
           Random Variables and Random Signal Principles", 4th ed., 2001,
           pp. 51, 51, 125.

    Examples
    --------
    >>> mu, sigma = 0, 0.1 # mean and standard deviation
    >>> s = np.random.normal(mu, sigma, 1000)

    Verify the mean and the variance:

    >>> np.abs(mu - np.mean(s)) < 0.01
    array(True)
    """
    return _mx_nd_np.random.normal(loc, scale, size, dtype, ctx, out)

def lognormal(mean=0.0, sigma=1.0, size=None, dtype=None, ctx=None, out=None):
    r"""Draw samples from a log-normal distribution.

    Draw samples from a `log-normal distribution` [1]_ with specified mean,
    standard deviation, and array shape. Note that the mean and standard
    deviation are not the values for the distribution itself, but of the
    underlying normal distribution it is derived from.

    Parameters
    ----------
    mean : float or array_like of floats, optional
        Mean value of the underlying normal distribution. Default is 0.
    sigma : float or array_like of floats, optional
        Standard deviation of the underlying normal distribution. Must be
        non-negative. Default is 1.
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a single value is returned if ``mean`` and ``sigma`` are both scalars.
        Otherwise, ``np.broadcast(mean, sigma).size`` samples are drawn.
    dtype : {'float16', 'float32', 'float64'}, optional
        Data type of output samples. Default is 'float32'
    ctx : Context, optional
        Device context of output. Default is current context.
    out : ``ndarray``, optional
        Store output to an existing ``ndarray``.

    Returns
    -------
    out : ndarray or scalar
        Drawn samples from the parameterized log-normal distribution.

    Notes
    -----
    A variable `x` has a log-normal distribution if `log(x)` is normally
    distributed.  The `probability density function for the log-normal
    distribution` [2]_ is:

    .. math:: p(x) = \frac{1}{\sigma x \sqrt{2\pi}}
                    e^{(-\frac{(ln(x)-\mu)^2}{2\sigma^2})}

    where :math:`\mu` is the mean and :math:`\sigma` is the standard
    deviation of the normally distributed logarithm of the variable.
    A log-normal distribution results if a random variable is the *product*
    of a large number of independent, identically-distributed variables in
    the same way that a normal distribution results if the variable is the
    *sum* of a large number of independent, identically-distributed
    variables.

    References
    ----------
    .. [1] Limpert, E., Stahel, W. A., and Abbt, M., "Log-normal
           Distributions across the Sciences: Keys and Clues,"
           BioScience, Vol. 51, No. 5, May, 2001.
           https://stat.ethz.ch/~stahel/lognormal/bioscience.pdf
    .. [2] Reiss, R.D. and Thomas, M., "Statistical Analysis of Extreme
           Values," Basel: Birkhauser Verlag, 2001, pp. 31-32.

    Examples
    --------
    Draw samples from the distribution:
    >>> mu, sigma = 3., 1. # mean and standard deviation
    >>> s = np.random.lognormal(mu, sigma, 1000)
    """
    return _mx_nd_np.random.lognormal(mean, sigma, size, dtype, ctx, out)

def logistic(loc=0.0, scale=1.0, size=None, ctx=None, out=None):
    r"""Draw samples from a logistic distribution.

    Samples are drawn from a logistic distribution with specified
    parameters, loc (location or mean, also median), and scale (>0).

    Parameters
    ----------
    loc : float or array_like of floats, optional
        Parameter of the distribution. Default is 0.
    scale : float or array_like of floats, optional
        Parameter of the distribution. Must be non-negative.
        Default is 1.
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a single value is returned if ``loc`` and ``scale`` are both scalars.
        Otherwise, ``np.broadcast(loc, scale).size`` samples are drawn.
    ctx : Context, optional
        Device context of output, default is current context.
    out : ``ndarray``, optional
        Store output to an existing ``ndarray``.

    Returns
    -------
    out : ndarray or scalar
        Drawn samples from the parameterized logistic distribution.

    Examples
    --------
    Draw samples from the distribution:
    >>> loc, scale = 10, 1
    >>> s = np.random.logistic(loc, scale, 10000)
    >>> import matplotlib.pyplot as plt
    >>> count, bins, ignored = plt.hist(s, bins=50)
    #   plot against distribution
    >>> def logist(x, loc, scale):
    ...     return np.exp((loc-x)/scale)/(scale*(1+np.exp((loc-x)/scale))**2)
    >>> lgst_val = logist(bins, loc, scale)
    >>> plt.plot(bins, lgst_val * count.max() / lgst_val.max())
    >>> plt.show()
    """
    return _mx_nd_np.random.logistic(loc, scale, size, ctx, out)

def gumbel(loc=0.0, scale=1.0, size=None, ctx=None, out=None):
    r"""Draw samples from a Gumbel distribution.

    Draw samples from a Gumbel distribution with specified location and
    scale.

    Parameters
    ----------
    loc : float or array_like of floats, optional
        The location of the mode of the distribution. Default is 0.
    scale : float or array_like of floats, optional
        The scale parameter of the distribution. Default is 1. Must be non-
        negative.
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a single value is returned if ``loc`` and ``scale`` are both scalars.
        Otherwise, ``np.broadcast(loc, scale).size`` samples are drawn.
    ctx : Context, optional
        Device context of output, default is current context.
    out : ``ndarray``, optional
        Store output to an existing ``ndarray``.

    Returns
    -------
    out : ndarray or scalar
        Drawn samples from the parameterized Gumbel distribution.

    Examples
    --------
    Draw samples from the distribution:
    >>> mu, beta = 0, 0.1 # location and scale
    >>> s = np.random.gumbel(mu, beta, 1000)
    Display the histogram of the samples, along with
    the probability density function:
    >>> import matplotlib.pyplot as plt
    >>> count, bins, ignored = plt.hist(s, 30, density=True)
    >>> plt.plot(bins, (1/beta)*np.exp(-(bins - mu)/beta)
    ...          * np.exp( -np.exp( -(bins - mu) /beta) ),
    ...          linewidth=2, color='r')
    >>> plt.show()
    Show how an extreme value distribution can arise from a Gaussian process
    and compare to a Gaussian:
    >>> means = []
    >>> maxima = []
    >>> for i in range(0,1000) :
    ...    a = np.random.normal(mu, beta, 1000)
    ...    means.append(a.mean())
    ...    maxima.append(a.max())
    >>> count, bins, ignored = plt.hist(maxima, 30, density=True)
    >>> beta = np.std(maxima) * np.sqrt(6) / np.pi
    >>> mu = np.mean(maxima) - 0.57721*beta
    >>> plt.plot(bins, (1/beta)*np.exp(-(bins - mu)/beta)
    ...          * np.exp(-np.exp(-(bins - mu)/beta)),
    ...          linewidth=2, color='r')
    >>> plt.plot(bins, 1/(beta * np.sqrt(2 * np.pi))
    ...          * np.exp(-(bins - mu)**2 / (2 * beta**2)),
    ...          linewidth=2, color='g')
    >>> plt.show()
    """
    return _mx_nd_np.random.gumbel(loc, scale, size, ctx, out)

def multinomial(n, pvals, size=None, **kwargs):
    r"""
    Draw samples from a multinomial distribution.
    The multinomial distribution is a multivariate generalisation of the binomial distribution.
    Take an experiment with one of ``p`` possible outcomes. An example of such an experiment is throwing a dice,
    where the outcome can be 1 through 6. Each sample drawn from the distribution represents n such experiments.
    Its values, ``X_i = [X_0, X_1, ..., X_p]``, represent the number of times the outcome was ``i``.

    Parameters
    ----------
    n : int
        Number of experiments.
    pvals : sequence of floats, length p
        Probabilities of each of the p different outcomes. These should sum to 1.
    size : int or tuple of ints, optional
        Output shape. If the given shape is, e.g., ``(m, n, k)``, then ``m * n * k`` samples
        are drawn. Default is None, in which case a single value is returned.

    Returns
    -------
    out : ndarray
        The drawn samples, of shape size, if that was provided. If not, the shape is ``(N,)``.
        In other words, each entry ``out[i,j,...,:]`` is an N-dimensional value drawn from the distribution.

    Examples
    --------
    Throw a dice 1000 times, and 1000 times again:

    >>> np.random.multinomial(1000, [1/6.]*6, size=2)
    array([[164, 161, 179, 158, 150, 188],
           [178, 162, 177, 143, 163, 177]])

    A loaded die is more likely to land on number 6:

    >>> np.random.multinomial(100, [1/7.]*5 + [2/7.])
    array([19, 14, 12, 11, 21, 23])
    >>> np.random.multinomial(100, [1.0 / 3, 2.0 / 3])
    array([32, 68])
    """
    return _mx_nd_np.random.multinomial(n, pvals, size, **kwargs)

# pylint: disable=unused-argument
def multivariate_normal(mean, cov, size=None, check_valid=None, tol=None):
    """
    multivariate_normal(mean, cov, size=None, check_valid=None, tol=None)

    Draw random samples from a multivariate normal distribution.

    The multivariate normal, multinormal or Gaussian distribution is a
    generalization of the one-dimensional normal distribution to higher
    dimensions.  Such a distribution is specified by its mean and
    covariance matrix.  These parameters are analogous to the mean
    (average or "center") and variance (standard deviation, or "width,"
    squared) of the one-dimensional normal distribution.

    This operator is a little different from the one in official NumPy.
    The official NumPy operator only accepts 1-D ndarray as mean and 2-D ndarray as cov,
    whereas the operator in MXNet np supports batch operation and auto-broadcasting.

    Both `mean` and `cov` may have any number of leading dimensions, which correspond
    to a batch shape. They are not necessarily assumed to have the same batch shape,
    just ones which can be broadcasted.

    Parameters
    ----------
    mean : K-D ndarray, of shape (..., N)
        Mean of the N-dimensional distribution.
    cov : (K+1)-D ndarray, of shape (..., N, N)
        Covariance matrix of the distribution. The last two dimensions must be symmetric and
        positive-semidefinite for proper sampling.
    size : int or tuple of ints, optional
        Given a shape of, for example, ``(m,n,k)``,
        ``m*n*k`` identically distributed batchs of samples are
        generated, and packed in an `m`-by-`n`-by-`k` arrangement.
        If no shape is specified, a batch of (`N`-D) sample is returned.
    check_valid : { 'warn', 'raise', 'ignore' }, optional
        Behavior when the covariance matrix is not positive semidefinite.
        (Not supported)
    tol : float, optional
        Tolerance when checking the singular values in covariance matrix.
        cov is cast to double before the check.
        (Not supported)

    Returns
    -------
    out : ndarray
        The input shape of `mean` and `cov` should satisfy the requirements of broadcasting.
        If the parameter `size` is not provided,
        the output shape is ``np.broadcast(mean.shape, cov.shape[:-1])``.
        Otherwise, the output shape is ``size + np.broadcast(mean.shape, cov.shape[:-1])``

    Examples
    --------
    >>> mean = np.array([1, 2])
    >>> cov = np.array([[1, 0], [0, 1]])
    >>> x = np.random.multivariate_normal(mean, cov, (3, 3))
    >>> x.shape
    (3, 3, 2)

    The following is probably true, given that 0.6 is roughly twice the
    standard deviation:

    >>> list((x[0,0,:] - mean) < 0.6)
    [True, True] # random

    # Performs autobroadcasting when the batch shape of
    # `mean` and `cov` is different but compatible.

    >>> mean = np.zeros((3,2)) # shape (3, 2)
    >>> cov = np.array([[1, 0], [0, 100]]) # shape (2, 2)
    >>> x = np.random.multivariate_normal(mean, cov)
    >>> x
    array([[-1.6115597 , -8.726251  ],
           [ 2.2425299 ,  2.8104177 ],
           [ 0.36229908, -8.386591  ]])
    """
    return _mx_nd_np.random.multivariate_normal(mean, cov, size=size, check_valid=None, tol=None)

def choice(a, size=None, replace=True, p=None, ctx=None, out=None):
    r"""Generates a random sample from a given 1-D array

    Parameters
    -----------
    a : 1-D array-like or int
        If an ndarray, a random sample is generated from its elements.
        If an int, the random sample is generated as if a were np.arange(a)
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  Default is None, in which case a
        single value is returned.
    replace : boolean, optional
        Whether the sample is with or without replacement
    p : 1-D array-like, optional
        The probabilities associated with each entry in a.
        If not given the sample assumes a uniform distribution over all
        entries in a.
    ctx : Context, optional
        Device context of output. Default is current context.

    Returns
    --------
    samples : ndarray
        The generated random samples

    Examples
    ---------
    Generate a uniform random sample from np.arange(5) of size 3:

    >>> np.random.choice(5, 3)
    array([0, 3, 4])
    >>> #This is equivalent to np.random.randint(0,5,3)

    Generate a non-uniform random sample from np.arange(5) of size 3:

    >>> np.random.choice(5, 3, p=[0.1, 0, 0.3, 0.6, 0])
    array([3, 3, 0])

    Generate a uniform random sample from np.arange(5) of size 3 without
    replacement:

    >>> np.random.choice(5, 3, replace=False)
    array([3,1,0])
    >>> #This is equivalent to np.random.permutation(np.arange(5))[:3]

    Generate a non-uniform random sample from np.arange(5) of size
    3 without replacement:

    >>> np.random.choice(5, 3, replace=False, p=[0.1, 0, 0.3, 0.6, 0])
    array([2, 3, 0])
    """
    return _mx_nd_np.random.choice(a, size, replace, p, ctx, out)

def rayleigh(scale=1.0, size=None, ctx=None, out=None):
    r"""Draw samples from a Rayleigh distribution.

    The :math:`\chi` and Weibull distributions are generalizations of the
    Rayleigh.

    Parameters
    ----------
    scale : float, optional
        Scale, also equals the mode. Must be non-negative. Default is 1.
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a single value is returned if ``scale`` is a scalar.  Otherwise,
        ``np.array(scale).size`` samples are drawn.
    ctx : Context, optional
        Device context of output, default is current context.
    out : ``ndarray``, optional
        Store output to an existing ``ndarray``.

    Returns
    -------
    out : ndarray or scalar
        Drawn samples from the parameterized Rayleigh distribution.
    """
    return _mx_nd_np.random.rayleigh(scale, size, ctx, out)

def rand(*size, **kwargs):
    r"""Random values in a given shape.

    Create an array of the given shape and populate it with random
    samples from a uniform distribution over [0, 1).

    Parameters
    ----------
    d0, d1, ..., dn : int, optional
        The dimensions of the returned array, should be all positive.
        If no argument is given a single Python float is returned.

    Returns
    -------
    out : ndarray
       Random values.

    Examples
    --------
    >>> np.random.rand(3,2)
    array([[ 0.14022471,  0.96360618],  #random
           [ 0.37601032,  0.25528411],  #random
           [ 0.49313049,  0.94909878]]) #random
    """
    output_shape = ()
    for s in size:
        output_shape += (s,)
    return _mx_nd_np.random.uniform(0, 1, size=output_shape, **kwargs)

def exponential(scale=1.0, size=None, ctx=None, out=None):
    r"""Draw samples from an exponential distribution.

    Parameters
    ----------
    scale : float or array_like of floats
        The scale parameter, :math:`\beta = 1/\lambda`. Must be
        non-negative.
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a single value is returned if ``scale`` is a scalar.  Otherwise,
        ``np.array(scale).size`` samples are drawn.
    ctx : Context, optional
        Device context of output, default is current context.
    out : ``ndarray``, optional
        Store output to an existing ``ndarray``.

    Returns
    -------
    out : ndarray or scalar
        Drawn samples from the parameterized exponential distribution.
    """
    return _mx_nd_np.random.exponential(scale, size=size, ctx=ctx, out=out)

def weibull(a, size=None, ctx=None, out=None):
    r"""Draw samples from a 1-parameter Weibull distribution with given parameter a
    via inversion.

    Parameters
    ----------
    a : float or array_like of floats
        Shape of the distribution. Must be non-negative.
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a single value is returned if ``a`` is a scalar. Otherwise,
        ``np.array(a).size`` samples are drawn.

    Returns
    -------
    out : ndarray or scalar
        Drawn samples from the 1-parameter Weibull distribution.

    Examples
    --------
    >>> np.random.weibull(a=5)
    array(0.9553641)
    >>> np.random.weibull(a=5, size=[2,3])
    array([[1.0466299 , 1.1320982 , 0.98415005],
          [1.1430776 , 0.9532727 , 1.1344457 ]])
    >>> np.random.weibull(a=np.array([2,3])
    array([0.98843634, 1.0125613 ])
    The Weibull distribution is one of a class of Generalized Extreme
    Value (GEV) distributions. This class includes the Gumbel and Frechet
    distributions.
    The probability density for the Weibull distribution is
    f(x) = \frac{a}{\lambda}(\frac{x}{\lambda})^{a-1}e^{-(x/\lambda)^a},
    where a is the shape and \lambda the scale. The generated 1-parameter Weibull
    sample has the scale parameter \lambda = 1.
    The Weibull distribution is commonly used in reliability engineering to
    model time to failure, in modeling particle sizes, in information retrieval
    to model dwell time on pages, in quantitative finance to model risk etc.
    """
    return _mx_nd_np.random.weibull(a, size=size, ctx=ctx, out=out)

def pareto(a, size=None, ctx=None, out=None):
    r"""Draw samples from a Pareto II or Lomax distribution with specified shape a.

    Parameters
    ----------
    a : float or array_like of floats
            Shape of the distribution. Must be > 0.
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a single value is returned if ``a`` is a scalar. Otherwise,
        ``np.array(a).size`` samples are drawn.

    Returns
    -------
    out : ndarray or scalar
        Drawn samples from the Pareto distribution.

    Examples
    --------
    >>> np.random.pareto(a=5)
    array(0.12749612)
    >>> mx.numpy.random.pareto(a=5, size=[2,3])
    array([[0.06933999, 0.0344373 , 0.10654891],
            [0.0311172 , 0.12911797, 0.03370714]])
    >>> np.random.pareto(a=np.array([2,3])
    array([0.26636696, 0.15685666])
    The probability density for the Pareto distribution is f(x) = \frac{am^a}{x^{a+1}}
    where a is the shape and m the scale. Here m is assumed 1. The Pareto distribution
    is a power law distribution. Pareto created it to describe the wealth in the economy.
    """
    return _mx_nd_np.random.pareto(a, size=size, ctx=ctx, out=out)

def power(a, size=None, ctx=None, out=None):
    r"""Draw samples in [0, 1] from a power distribution with given parameter a.

    Parameters
    ----------
    a : float or array_like of floats
        Shape of the distribution. Must be > 0.
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a single value is returned if ``a`` is a scalar. Otherwise,
        ``np.array(a).size`` samples are drawn.

    Returns
    -------
    out : ndarray or scalar
        Drawn samples from the power distribution.

    Examples
    --------
    >>> np.random.power(a=5)
    array(0.8602478)
    >>> np.random.power(a=5, size=[2,3])
    array([[0.988391  , 0.5153122 , 0.9383134 ],
           [0.9078098 , 0.87819266, 0.730635]])
    >>> np.random.power(a=np.array([2,3])
    array([0.7499419 , 0.88894516])
    The probability density function is f(x; a) = ax^{a-1}, 0 \le x \le 1, a>0.
    The power distribution is just the inverse of the Pareto distribution and
    a special case of the Beta distribution.
    """
    return _mx_nd_np.random.power(a, size=size, ctx=ctx, out=out)

def shuffle(x):
    """
    Modify a sequence in-place by shuffling its contents.

    This function only shuffles the array along the first axis of a
    multi-dimensional array. The order of sub-arrays is changed but
    their contents remain the same.

    Parameters
    ----------
    x: ndarray
        The array or list to be shuffled.

    Examples
    --------
    >>> arr = np.arange(10)
    >>> np.random.shuffle(arr)
    >>> arr
    array([5., 1., 0., 6., 7., 3., 9., 8., 4., 2.])  # random

    Multi-dimensional arrays are only shuffled along the first axis:

    >>> arr = np.arange(9).reshape((3, 3))
    >>> np.random.shuffle(arr)
    >>> arr
    array([[6., 7., 8.], # random
           [3., 4., 5.],
           [0., 1., 2.]])
    """
    _mx_nd_np.random.shuffle(x)

def gamma(shape, scale=1.0, size=None, dtype=None, ctx=None, out=None):
    """Draw samples from a Gamma distribution.

    Samples are drawn from a Gamma distribution with specified parameters,
    `shape` (sometimes designated "k") and `scale` (sometimes designated
    "theta"), where both parameters are > 0.

    The Gamma distribution is often used to model the times to failure of
    electronic components, and arises naturally in processes for which the
    waiting times between Poisson distributed events are relevant.

    Parameters
    ----------
    shape : float or array_like of floats
        The shape of the gamma distribution. Should be greater than zero.
    scale : float or array_like of floats, optional
        The scale of the gamma distribution. Should be greater than zero.
        Default is equal to 1.
    dtype : {'float16', 'float32', 'float64'}, optional
        Data type of output samples. Default is 'float32'.
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a single value is returned if ``shape`` and ``scale`` are both scalars.
        Otherwise, ``np.broadcast(shape, scale).size`` samples are drawn.
    ctx : Context, optional
        Device context of output. Default is current context.

    Returns
    -------
    out : ndarray or scalar
        Drawn samples from the parameterized gamma distribution.
    """
    return _mx_nd_np.random.gamma(shape, scale, size, dtype, ctx, out)

def beta(a, b, size=None, dtype=None, ctx=None):
    r"""Draw samples from a Beta distribution.

    The Beta distribution is a special case of the Dirichlet distribution,
    and is related to the Gamma distribution.  It has the probability
    distribution function

    .. math:: f(x; a,b) = \frac{1}{B(\alpha, \beta)} x^{\alpha - 1}
                                                     (1 - x)^{\beta - 1},

    where the normalisation, B, is the beta function,

    .. math:: B(\alpha, \beta) = \int_0^1 t^{\alpha - 1}
                                 (1 - t)^{\beta - 1} dt.

    It is often seen in Bayesian inference and order statistics.

    Parameters
    ----------
    a : float or array_like of floats
        Alpha, positive (>0).
    b : float or array_like of floats
        Beta, positive (>0).
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a single value is returned if ``a`` and ``b`` are both scalars.
        Otherwise, ``np.broadcast(a, b).size`` samples are drawn.
    dtype : {'float16', 'float32', 'float64'}, optional
        Data type of output samples. Default is 'float32'.
        Dtype 'float32' or 'float64' is strongly recommended,
        since lower precision might lead to out of range issue.
    ctx : Context, optional
        Device context of output. Default is current context.

    Notes
    -----
    To use this operator with scalars as input, please run
    ``npx.set_np()`` first.

    Returns
    -------
    out : ndarray or scalar
        Drawn samples from the parameterized beta distribution.
    """
    return _mx_nd_np.random.beta(a, b, size=size, dtype=dtype, ctx=ctx)

def f(dfnum, dfden, size=None, ctx=None):
    r"""Draw samples from an F distribution.

    Samples are drawn from an F distribution with specified parameters,
    `dfnum` (degrees of freedom in numerator) and `dfden` (degrees of
    freedom in denominator), where both parameters must be greater than
    zero.

    The random variate of the F distribution (also known as the
    Fisher distribution) is a continuous probability distribution
    that arises in ANOVA tests, and is the ratio of two chi-square
    variates.

    Parameters
    ----------
    dfnum : float or ndarray of floats
        Degrees of freedom in numerator, must be > 0.
    dfden : float or ndarray of float
        Degrees of freedom in denominator, must be > 0.
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a single value is returned if ``dfnum`` and ``dfden`` are both scalars.
        Otherwise, ``np.broadcast(dfnum, dfden).size`` samples are drawn.
    ctx : Context, optional
        Device context of output. Default is current context.

    Returns
    -------
    out : ndarray or scalar
        Drawn samples from the parameterized Fisher distribution.

    Examples
    --------
    An example from Glantz[1], pp 47-40:

    Two groups, children of diabetics (25 people) and children from people
    without diabetes (25 controls). Fasting blood glucose was measured,
    case group had a mean value of 86.1, controls had a mean value of
    82.2. Standard deviations were 2.09 and 2.49 respectively. Are these
    data consistent with the null hypothesis that the parents diabetic
    status does not affect their children's blood glucose levels?
    Calculating the F statistic from the data gives a value of 36.01.

    Draw samples from the distribution:

    >>> dfnum = 1. # between group degrees of freedom
    >>> dfden = 48. # within groups degrees of freedom
    >>> s = np.random.f(dfnum, dfden, 1000)

    The lower bound for the top 1% of the samples is :

    >>> np.sort(s)[-10]
    7.61988120985 # random

    So there is about a 1% chance that the F statistic will exceed 7.62,
    the measured value is 36, so the null hypothesis is rejected at the 1%
    level.
    """
    return _mx_nd_np.random.f(dfnum, dfden, size=size, ctx=ctx)

def chisquare(df, size=None, dtype=None, ctx=None):
    r"""Draw samples from a chi-square distribution.

    When `df` independent random variables, each with standard normal
    distributions (mean 0, variance 1), are squared and summed, the
    resulting distribution is chi-square (see Notes).  This distribution
    is often used in hypothesis testing.

    Parameters
    ----------
    df : float or ndarray of floats
         Number of degrees of freedom, must be > 0.
    size : int or tuple of ints, optional
        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
        ``m * n * k`` samples are drawn.  If size is ``None`` (default),
        a single value is returned if ``df`` is a scalar.  Otherwise,
        ``np.array(df).size`` samples are drawn.
    dtype : {'float16', 'float32', 'float64'}, optional
        Data type of output samples. Default is 'float32'.
    ctx : Context, optional
        Device context of output. Default is current context.

    Returns
    -------
    out : ndarray or scalar
        Drawn samples from the parameterized `chi-square distribution` [1]_.

    Raises
    ------
    ValueError
        When `df` <= 0 or when an inappropriate `size`
        is given.

    Notes
    -----
    The variable obtained by summing the squares of `df` independent,
    standard normally distributed random variables:

    .. math:: Q = \sum_{i=0}^{\mathtt{df}} X^2_i

    is chi-square distributed, denoted

    .. math:: Q \sim \chi^2_k.

    The probability density function of the chi-squared distribution is

    .. math:: p(x) = \frac{(1/2)^{k/2}}{\Gamma(k/2)}
                     x^{k/2 - 1} e^{-x/2},

    where :math:`\Gamma` is the gamma function,

    .. math:: \Gamma(x) = \int_0^{-\infty} t^{x - 1} e^{-t} dt.

    References
    ----------
    .. [1] NIST "Engineering Statistics Handbook"
           https://www.itl.nist.gov/div898/handbook/eda/section3/eda3666.htm

    Examples
    --------
    >>> np.random.chisquare(2,4)
    array([ 1.89920014,  9.00867716,  3.13710533,  5.62318272]) # random
    """
    return _mx_nd_np.random.chisquare(df, size=size, dtype=dtype, ctx=ctx)

def randn(*size, **kwargs):
    r"""Return a sample (or samples) from the "standard normal" distribution.
    If positive, int_like or int-convertible arguments are provided,
    `randn` generates an array of shape ``(d0, d1, ..., dn)``, filled
    with random floats sampled from a univariate "normal" (Gaussian)
    distribution of mean 0 and variance 1 (if any of the :math:`d_i` are
    floats, they are first converted to integers by truncation). A single
    float randomly sampled from the distribution is returned if no
    argument is provided.
    This is a convenience function.  If you want an interface that takes a
    tuple as the first argument, use `numpy.random.standard_normal` instead.
    Parameters
    ----------
    d0, d1, ..., dn : int, optional
        The dimensions of the returned array, should be all positive.
        If no argument is given a single Python float is returned.
    Returns
    -------
    Z : ndarray
        A ``(d0, d1, ..., dn)``-shaped array of floating-point samples from
        the standard normal distribution, or a single such float if
        no parameters were supplied.
    Notes
    -----
    For random samples from :math:`N(\mu, \sigma^2)`, use:
    ``sigma * np.random.randn(...) + mu``
    Examples
    --------
    >>> np.random.randn()
    2.1923875335537315 #random
    Two-by-four array of samples from N(3, 6.25):
    >>> 2.5 * np.random.randn(2, 4) + 3
    array([[-4.49401501,  4.00950034, -1.81814867,  7.29718677],  #random
        [ 0.39924804,  4.68456316,  4.99394529,  4.84057254]]) #random
    """
    output_shape = ()
    for s in size:
        output_shape += (s,)
    return _mx_nd_np.random.normal(0, 1, size=output_shape, **kwargs)

def laplace(loc=0.0, scale=1.0, size=None, dtype=None, ctx=None, out=None):
    r"""Draw random samples from a Laplace distribution.

    Samples are distributed according to a Laplace distribution parametrized
    by *loc* (mean) and *scale* (the exponential decay).

    Parameters
    ----------
    loc : float, The position of the distribution peak.

    scale : float, the exponential decay.

    size : int or tuple of ints, optional. Output shape.
        If the given shape is, e.g., (m, n, k), then m * n * k samples are drawn.
        Default is None, in which case a single value is returned.

    dtype : {'float16', 'float32', 'float64'}, optional
        Data type of output samples. Default is 'float32'
    ctx : Context, optional
        Device context of output. Default is current context.
    out : ``ndarray``, optional
        Store output to an existing ``ndarray``.

    Returns
    -------
    out : ndarray
        Drawn samples from the parameterized Laplace distribution.
    """
    return _mx_nd_np.random.laplace(loc, scale, size, dtype, ctx, out)