from copy import deepcopy

import numpy as np

from sklearn.base import clone
from sklearn.ensemble import AdaBoostClassifier
from sklearn.ensemble._base import _set_random_states
from sklearn.utils import _safe_indexing

from ..under_sampling.base import BaseUnderSampler
from ..under_sampling import RandomUnderSampler
from ..pipeline import make_pipeline
from ..utils import Substitution, check_target_type
from ..utils._docstring import _random_state_docstring
from ..utils._validation import _deprecate_positional_args

@Substitution(
    sampling_strategy=BaseUnderSampler._sampling_strategy_docstring,
    random_state=_random_state_docstring,
)
class RUSBoostClassifier(AdaBoostClassifier):
    """Random under-sampling integrated in the learning of AdaBoost.

    During learning, the problem of class balancing is alleviated by random
    under-sampling the sample at each iteration of the boosting algorithm.

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

    .. versionadded:: 0.4

    Parameters
    ----------
    base_estimator : estimator object, default=None
        The base estimator from which the boosted ensemble is built.
        Support for sample weighting is required, as well as proper
        ``classes_`` and ``n_classes_`` attributes. If ``None``, then
        the base estimator is ``DecisionTreeClassifier(max_depth=1)``.

    n_estimators : int, default=50
        The maximum number of estimators at which boosting is terminated.
        In case of perfect fit, the learning procedure is stopped early.

    learning_rate : float, default=1.0
        Learning rate shrinks the contribution of each classifier by
        ``learning_rate``. There is a trade-off between ``learning_rate`` and
        ``n_estimators``.

    algorithm : {{'SAMME', 'SAMME.R'}}, default='SAMME.R'
        If 'SAMME.R' then use the SAMME.R real boosting algorithm.
        ``base_estimator`` must support calculation of class probabilities.
        If 'SAMME' then use the SAMME discrete boosting algorithm.
        The SAMME.R algorithm typically converges faster than SAMME,
        achieving a lower test error with fewer boosting iterations.

    {sampling_strategy}

    replacement : bool, default=False
        Whether or not to sample randomly with replacement or not.

    {random_state}

    Attributes
    ----------
    base_estimator_ : estimator
        The base estimator from which the ensemble is grown.

    estimators_ : list of classifiers
        The collection of fitted sub-estimators.

    samplers_ : list of RandomUnderSampler
        The collection of fitted samplers.

    pipelines_ : list of Pipeline
        The collection of fitted pipelines (samplers + trees).

    classes_ : ndarray of shape (n_classes,)
        The classes labels.

    n_classes_ : int
        The number of classes.

    estimator_weights_ : ndarray of shape (n_estimator,)
        Weights for each estimator in the boosted ensemble.

    estimator_errors_ : ndarray of shape (n_estimator,)
        Classification error for each estimator in the boosted
        ensemble.

    feature_importances_ : ndarray of shape (n_features,)
        The feature importances if supported by the ``base_estimator``.

    See Also
    --------
    BalancedBaggingClassifier : Bagging classifier for which each base
        estimator is trained on a balanced bootstrap.

    BalancedRandomForestClassifier : Random forest applying random-under
        sampling to balance the different bootstraps.

    EasyEnsembleClassifier : Ensemble of AdaBoost classifier trained on
        balanced bootstraps.

    References
    ----------
    .. [1] Seiffert, C., Khoshgoftaar, T. M., Van Hulse, J., & Napolitano, A.
       "RUSBoost: A hybrid approach to alleviating class imbalance." IEEE
       Transactions on Systems, Man, and Cybernetics-Part A: Systems and Humans
       40.1 (2010): 185-197.

    Examples
    --------
    >>> from imblearn.ensemble import RUSBoostClassifier
    >>> from sklearn.datasets import make_classification
    >>>
    >>> X, y = make_classification(n_samples=1000, n_classes=3,
    ...                            n_informative=4, weights=[0.2, 0.3, 0.5],
    ...                            random_state=0)
    >>> clf = RUSBoostClassifier(random_state=0)
    >>> clf.fit(X, y)  # doctest: +ELLIPSIS
    RUSBoostClassifier(...)
    >>> clf.predict(X)  # doctest: +ELLIPSIS
    array([...])
    """

    @_deprecate_positional_args
    def __init__(
        self,
        base_estimator=None,
        *,
        n_estimators=50,
        learning_rate=1.0,
        algorithm="SAMME.R",
        sampling_strategy="auto",
        replacement=False,
        random_state=None,
    ):
        super().__init__(
            base_estimator=base_estimator,
            n_estimators=n_estimators,
            learning_rate=learning_rate,
            algorithm=algorithm,
            random_state=random_state,
        )
        self.sampling_strategy = sampling_strategy
        self.replacement = replacement

    def fit(self, X, y, sample_weight=None):
        """Build a boosted classifier from the training set (X, y).

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            The training input samples. Sparse matrix can be CSC, CSR, COO,
            DOK, or LIL. DOK and LIL are converted to CSR.

        y : array-like of shape (n_samples,)
            The target values (class labels).

        sample_weight : array-like of shape (n_samples,), default=None
            Sample weights. If None, the sample weights are initialized to
            ``1 / n_samples``.

        Returns
        -------
        self : object
            Returns self.
        """
        check_target_type(y)
        self.samplers_ = []
        self.pipelines_ = []
        super().fit(X, y, sample_weight)
        return self

    def _validate_estimator(self):
        """Check the estimator and the n_estimator attribute, set the
        `base_estimator_` attribute."""
        super()._validate_estimator()

        self.base_sampler_ = RandomUnderSampler(
            sampling_strategy=self.sampling_strategy,
            replacement=self.replacement,
        )

    def _make_sampler_estimator(self, append=True, random_state=None):
        """Make and configure a copy of the `base_estimator_` attribute.
        Warning: This method should be used to properly instantiate new
        sub-estimators.
        """
        estimator = clone(self.base_estimator_)
        estimator.set_params(**{p: getattr(self, p) for p in self.estimator_params})
        sampler = clone(self.base_sampler_)

        if random_state is not None:
            _set_random_states(estimator, random_state)
            _set_random_states(sampler, random_state)

        if append:
            self.estimators_.append(estimator)
            self.samplers_.append(sampler)
            self.pipelines_.append(
                make_pipeline(deepcopy(sampler), deepcopy(estimator))
            )

        return estimator, sampler

    def _boost_real(self, iboost, X, y, sample_weight, random_state):
        """Implement a single boost using the SAMME.R real algorithm."""
        estimator, sampler = self._make_sampler_estimator(random_state=random_state)

        X_res, y_res = sampler.fit_resample(X, y)
        sample_weight_res = _safe_indexing(sample_weight, sampler.sample_indices_)
        estimator.fit(X_res, y_res, sample_weight=sample_weight_res)

        y_predict_proba = estimator.predict_proba(X)

        if iboost == 0:
            self.classes_ = getattr(estimator, "classes_", None)
            self.n_classes_ = len(self.classes_)

        y_predict = self.classes_.take(np.argmax(y_predict_proba, axis=1), axis=0)

        # Instances incorrectly classified
        incorrect = y_predict != y

        # Error fraction
        estimator_error = np.mean(np.average(incorrect, weights=sample_weight, axis=0))

        # Stop if classification is perfect
        if estimator_error <= 0:
            return sample_weight, 1.0, 0.0

        # Construct y coding as described in Zhu et al [2]:
        #
        #    y_k = 1 if c == k else -1 / (K - 1)
        #
        # where K == n_classes_ and c, k in [0, K) are indices along the second
        # axis of the y coding with c being the index corresponding to the true
        # class label.
        n_classes = self.n_classes_
        classes = self.classes_
        y_codes = np.array([-1.0 / (n_classes - 1), 1.0])
        y_coding = y_codes.take(classes == y[:, np.newaxis])

        # Displace zero probabilities so the log is defined.
        # Also fix negative elements which may occur with
        # negative sample weights.
        proba = y_predict_proba  # alias for readability
        np.clip(proba, np.finfo(proba.dtype).eps, None, out=proba)

        # Boost weight using multi-class AdaBoost SAMME.R alg
        estimator_weight = (
            -1.0
            * self.learning_rate
            * ((n_classes - 1.0) / n_classes)
            * (y_coding * np.log(y_predict_proba)).sum(axis=1)
        )

        # Only boost the weights if it will fit again
        if not iboost == self.n_estimators - 1:
            # Only boost positive weights
            sample_weight *= np.exp(
                estimator_weight * ((sample_weight > 0) | (estimator_weight < 0))
            )

        return sample_weight, 1.0, estimator_error

    def _boost_discrete(self, iboost, X, y, sample_weight, random_state):
        """Implement a single boost using the SAMME discrete algorithm."""
        estimator, sampler = self._make_sampler_estimator(random_state=random_state)

        X_res, y_res = sampler.fit_resample(X, y)
        sample_weight_res = _safe_indexing(sample_weight, sampler.sample_indices_)
        estimator.fit(X_res, y_res, sample_weight=sample_weight_res)

        y_predict = estimator.predict(X)

        if iboost == 0:
            self.classes_ = getattr(estimator, "classes_", None)
            self.n_classes_ = len(self.classes_)

        # Instances incorrectly classified
        incorrect = y_predict != y

        # Error fraction
        estimator_error = np.mean(np.average(incorrect, weights=sample_weight, axis=0))

        # Stop if classification is perfect
        if estimator_error <= 0:
            return sample_weight, 1.0, 0.0

        n_classes = self.n_classes_

        # Stop if the error is at least as bad as random guessing
        if estimator_error >= 1.0 - (1.0 / n_classes):
            self.estimators_.pop(-1)
            self.samplers_.pop(-1)
            self.pipelines_.pop(-1)
            if len(self.estimators_) == 0:
                raise ValueError(
                    "BaseClassifier in AdaBoostClassifier "
                    "ensemble is worse than random, ensemble "
                    "can not be fit."
                )
            return None, None, None

        # Boost weight using multi-class AdaBoost SAMME alg
        estimator_weight = self.learning_rate * (
            np.log((1.0 - estimator_error) / estimator_error) + np.log(n_classes - 1.0)
        )

        # Only boost the weights if I will fit again
        if not iboost == self.n_estimators - 1:
            # Only boost positive weights
            sample_weight *= np.exp(estimator_weight * incorrect * (sample_weight > 0))

        return sample_weight, estimator_weight, estimator_error


__author__ = 'JMartinez'

import numpy as np, pandas as pd, os
from synthicity.utils import misc
import pysal as py

class elasticity_model(object):

    def __init__(self, dset):
        self.zones = dset.zones
        self.buildings_far = pd.merge(dset.buildings, dset.fars, left_on='far_id', right_index=True)

    def estimate_elasticity(self, zones):
        dummies = pd.get_dummies(zones.county)
        zones = pd.concat([zones, dummies], axis=1)
        zones['avg_far'] = self.buildings_far.groupby('zone_id').far.mean() #use far_x because Xavier's code adds far to buildings

        #zones = zones[zones.residential_sqft_zone>0]

        #wrook = py.queen_from_shapefile('C:/users/jmartinez/documents/Test Zones/zones_prj_res2.shp')
        wqueen = py.queen_from_shapefile(os.path.join(misc.data_dir(),'shapefiles\\zones.shp'))
        w = py.weights.weights.W(wqueen.neighbors, wqueen.weights)
        x = zones[['zonal_pop','mean_income']]
        x = x.apply(np.log1p)

        x['ln_jobs_within_30min'] = zones['ln_jobs_within_30min']
        x['zone_contains_park'] = zones['zone_contains_park']
        x['Arapahoe'] = zones['Arapahoe']
        x['Boulder'] = zones['Boulder']
        x['Broomfield'] = zones['Broomfield']
        x['Clear Creek'] = zones['Clear Creek']
        x['Denver'] = zones['Denver']
        x['Douglas'] = zones['Douglas']
        x['Elbert'] = zones['Elbert']
        x['Gilpin'] = zones['Gilpin']
        x['Jefferson'] = zones['Jefferson']
        x['Weld'] = zones['Weld']
        x=x.fillna(0)
        x = x.as_matrix()

        imat = zones[['ln_avg_nonres_unit_price_zone','avg_far']]
        imat = imat.fillna(0)
        imat = imat.as_matrix()

        yend = zones['ln_avg_unit_price_zone']
        yend = yend.fillna(0)
        yend = yend.as_matrix()
        yend = np.reshape(yend,(zones.shape[0],1))

        y = zones['residential_sqft_zone']
        y = y.fillna(0)
        y = y.apply(np.log1p)
        y = y.as_matrix()
        y = np.reshape(y,(zones.shape[0],1))

        imat_names = ['non_res_price','avg_far']
        x_names = ['zonal_pop', 'mean_income', 'ln_jobs_within_30min', 'zone_contains_park','Arapahoe','Boulder','Broomfield','Clear Creek','Denver','Douglas','Elbert','Gilpin','Jefferson','Weld']
        yend_name = ['ln_avg_unit_price_zone']
        y_name = 'residential_sqft_zone'

        reg_2sls = py.spreg.twosls_sp.GM_Lag(y, x, yend=yend, q=imat, w=w, w_lags=2, robust ='white', name_x = x_names, name_q = imat_names, name_y = y_name, name_yend = yend_name)

        demand_elasticity = np.absolute(reg_2sls.betas[15])
        demand_elasticity = 1/demand_elasticity[0]
        #
        return demand_elasticity

    def estimate_non_res_elasticity(self,zones):

        dummies = pd.get_dummies(zones.county)
        zones = pd.concat([zones, dummies], axis=1)
        zones['avg_far'] = self.buildings_far.groupby('zone_id').far.mean() #use far_x because Xavier's code adds far to buildings
        #zones = zones[zones.non_residential_sqft_zone>0]

        ####spatial weights matrix#####
        #zones = zones.reset_index()
        #zone_coord = zones[['zone_id','zonecentroid_x', 'zonecentroid_y']]

        #zone_coord = zone_coord.as_matrix()

        wqueen = py.queen_from_shapefile(os.path.join(misc.data_dir(),'shapefiles\\zones.shp'))
        #w = py.weights.Distance.DistanceBand(zone_coord, threshold = 50000, binary = False)
        #w.transform ='r'
        #w = py.weights.weights.W(w.neighbors, w.weights)
        w = py.weights.weights.W(wqueen.neighbors, wqueen.weights)
        x = zones[['zonal_emp','residential_units_zone']]
        x = x.apply(np.log1p)
        #x['ln_emp_aggsector_within_5min'] = zones['ln_emp_aggsector_within_5min']
        #x['zone_contains_park'] = zones['zone_contains_park']
        x['percent_younghead'] = zones['percent_younghead']
        x['Arapahoe'] = zones['Arapahoe']
        x['Boulder'] = zones['Boulder']
        x['Broomfield'] = zones['Broomfield']
        x['Clear Creek'] = zones['Clear Creek']
        x['Denver'] = zones['Denver']
        x['Douglas'] = zones['Douglas']
        x['Elbert'] = zones['Elbert']
        x['Gilpin'] = zones['Gilpin']
        x['Jefferson'] = zones['Jefferson']
        x['Weld'] = zones['Weld']
        x=x.fillna(0)
        x = x.as_matrix()

        imat = zones[['ln_avg_unit_price_zone','avg_far']]
        imat = imat.fillna(0)
        imat = imat.as_matrix()

        yend = zones['ln_avg_nonres_unit_price_zone']
        yend = yend.fillna(0)
        yend = yend.as_matrix()
        yend = np.reshape(yend,(zones.shape[0],1))

        y = zones['non_residential_sqft_zone']
        y = y.fillna(0)
        y = y.apply(np.log1p)
        y = y.as_matrix()
        y = np.reshape(y,(zones.shape[0],1))

        imat_names = ['res_price','avg_far']
        x_names = ['zonal_emp', 'residential_units_zone', 'percent_younghead','Arapahoe','Boulder','Broomfield','Clear Creek', 'Denver', 'Douglas','Elbert','Gilpin','Jefferson','Weld']
        yend_name = ['ln_avg_nonres_unit_price_zone']
        y_name = 'non_residential_sqft_zone'

        reg_2sls = py.spreg.twosls_sp.GM_Lag(y, x, yend=yend, q=imat, w=w, w_lags=2,robust ='white', name_x = x_names, name_q = imat_names, name_y = y_name, name_yend = yend_name)

        #
        # ######estimation
        # x = zones[['zonal_emp','residential_units_zone']]
        # x = x.apply(np.log1p)
        # #x['ln_emp_aggsector_within_5min'] = zones['ln_emp_aggsector_within_5min']
        # #x['zone_contains_park'] = zones['zone_contains_park']
        # x['percent_younghead'] = zones['percent_younghead']
        # x=x.fillna(0)
        # x = x.as_matrix()
        #
        # imat = zones[['ln_avg_unit_price_zone','ln_avg_land_value_per_sqft_zone','median_year_built']]
        # imat = imat.fillna(0)
        # imat = imat.as_matrix()
        #
        # yend = zones['ln_avg_nonres_unit_price_zone']
        # yend = yend.fillna(0)
        # yend = yend.as_matrix()
        # yend = np.reshape(yend,(zones.shape[0],1))
        #
        # y = zones['non_residential_sqft_zone']
        # y = y.fillna(0)
        # y = y.apply(np.log1p)
        # y = y.as_matrix()
        # y = np.reshape(y,(zones.shape[0],1))
        #
        #
        # imat_names = ['res_price','land_value','median_year_built']
        # x_names = ['zonal_emp', 'residential_units_zone', 'percent_younghead']
        # yend_name = ['ln_avg_nonres_unit_price_zone']
        # y_name = 'non_residential_sqft_zone'
        #
        # reg_2sls = py.spreg.twosls_sp.GM_Lag(y, x, yend=yend, q=imat, w=w, robust ='white', name_x = x_names, name_q = imat_names, name_y = y_name, name_yend = yend_name)
        #
        #
        demand_elasticity = np.absolute(reg_2sls.betas[14])
        demand_elasticity = 1/demand_elasticity[0]
        #
        return demand_elasticity


#!/bin/python2

from __future__ import division

import subprocess
from time import sleep
import os

# for generating sound files
import numpy as np
import matplotlib.pyplot as plt
import scipy.io.wavfile
import scipy.signal as sig
import scipy.stats as stats

master_volume = 1
sounds = {
    'A':{'filename':'tick1.wav',
         'volume': .9,
         'freq': 10,
         'length': .01,
         'quality': 1,
         'tone':'sawtooth',
         'a': 2,
         'b': 10,},
    'W':{'filename':'tick2.wav',
         'volume': .8,
         'freq': 5,
         'length': .01,
         'quality': .8,
         'tone':'sawtooth',
         'a': 2,
         'b': 5,},
    'M':{'filename':'tick3.wav',
         'volume': .8,
         'freq': 10,
         'length': .05,
         'quality': .95,
         'tone':'sawtooth',
         'a': 2,
         'b': 5,},
    'S':{'filename':'tick4.wav',
         'volume': .4,
         'freq': 50,
         'length': .04,
         'quality': .6,
         'tone':'sawtooth',
         'a': 2,
         'b': 5,},
    'U':{'filename':'tick5.wav',
         'volume': .5,
         'freq': 40,
         'length': .02,
         'quality': .9,
         'tone':'sawtooth',
         'a': 2,
         'b': 5,},
}

def construct_sound(params, plot_sound=False):
    print "constructing sound: %s" % params['filename']
    rate = 44100
    N = int(rate*params['length'])
    time = range(N)

    if params['tone'] == 'sawtooth':
        raw = sig.sawtooth(np.linspace(0,params['freq'],N))
    elif params['tone'] == 'sine':
        # not succesfully tested, try plotting
        raw = np.sin(np.linspace(0,params['freq'],N))
    else:
        raise RuntimeError('unknown tone type')

    noise = np.random.uniform(-1, 1, N) # 44100 random samples between -1 and 1
    envelope = stats.beta(params['a'],params['b']).pdf([n/N for n in time])

    data = raw*params['quality'] + noise*(1-params['quality'])
    data *= envelope
    save_wav(data, params['filename'], params['volume'])
    if plot_sound:
        figure()
        plt.plot(time, raw)
        plt.plot(time, envelope)
        plt.plot(time, data)
        plt.show()

def save_wav(data, filename, volume=1):
    total_volume = volume * master_volume
    if total_volume > 1 or total_volume < 0:
        raise RuntimeError('volume out of range')
    scaled_data = np.int16(data/np.max(np.abs(data)) * total_volume * 32767)
    scipy.io.wavfile.write(filename, 44100, scaled_data)

def main():
    data = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8] * 10000
    save_wav(data, "test.wav")
    # for code in sounds.keys():
    #     construct_sound(sounds[code])

main()


import matplotlib.pyplot as plt
import numpy as np
import sys
import math
from scipy.signal import triang
from scipy.fftpack import fft, fftshift

M = 127
N = 128
hM1 = int(math.floor((M+1)/2))
hM2 = int(math.floor(M/2))
x = triang(M)
fftbuffer = np.zeros(N)
fftbuffer[:hM1] = x[hM2:]
fftbuffer[N-hM2:] = x[:hM2]
X = fftshift(fft(fftbuffer))
mX = abs(X)
pX = np.unwrap(np.angle(X))

plt.figure(1, figsize=(9.5, 4))
plt.subplot(311)
plt.title('x[n]')
plt.plot(np.arange(-hM2, hM1, 1.0), x, 'b', lw=1.5)
plt.axis([-hM2, hM1, 0, 1])

plt.subplot(323)
plt.title('real(X)')
plt.plot(np.arange(-N/2, N/2, 1.0), np.real(X), 'r', lw=1.5)
plt.axis([-N/2, N/2, min(np.real(X)), max(np.real(X))])

plt.subplot(324)
plt.title('im(X)')
plt.plot(np.arange(-N/2, N/2, 1.0), np.imag(X), 'c', lw=1.5)
plt.axis([-N/2, N/2, -1, 1])

plt.subplot(325)
plt.title('abs(X)')
plt.plot(np.arange(-N/2, N/2, 1.0), mX, 'r', lw=1.5)
plt.axis([-N/2,N/2,min(mX),max(mX)])

plt.subplot(326)
plt.title('angle(X)')
plt.plot(np.arange(-N/2, N/2, 1.0), pX, 'c', lw=1.5)
plt.axis([-N/2, N/2, -1, 1])

plt.tight_layout()
plt.savefig('symmetry-real-even.png')
plt.show()


from os import path
from setuptools import setup, find_packages

# read the contents of your README file
this_directory = path.abspath(path.dirname(__file__))
# with open(path.join(this_directory, 'README.md'), encoding='utf-8') as f:
with open(path.join(this_directory, 'README.md'), encoding='utf-8') as f:
    long_description = f.read()

setup(
    name="OpenPIV",
    version='0.23.6',
    packages=find_packages(),
    include_package_data=True,
    long_description=long_description,
    long_description_content_type='text/markdown',
    setup_requires=[
        'setuptools',
    ],
    install_requires=[
        'numpy',
        'imageio',
        'matplotlib>=3',
        'scikit-image',
        'scipy',
        'natsort',
        'GitPython',
        'pytest',
        'tqdm'
    ],
    classifiers=[
        # PyPI-specific version type. The number specified here is a magic
        # constant
        # with no relation to this application's version numbering scheme.
        # *sigh*
        'Development Status :: 4 - Beta',

        # Sublist of all supported Python versions.
        'Programming Language :: Python :: 3.7',
        'Programming Language :: Python :: 3.8',

        # Sublist of all supported platforms and environments.
        'Operating System :: MacOS :: MacOS X',
        'Operating System :: Microsoft :: Windows',
        'Operating System :: POSIX',

        # Miscellaneous metadata.
        'Intended Audience :: Science/Research',
        'License :: OSI Approved :: GNU General Public License v3 (GPLv3)',
        'Natural Language :: English',
        'Operating System :: OS Independent',
        'Topic :: Scientific/Engineering',
    ],
    # long_description=long_description,
    # long_description_content_type='text/markdown'
)


# coding: utf-8

# Author: Johannes Schönberger
#
# License: BSD 3 clause

import numpy as np

from ..base import BaseEstimator, MetaEstimatorMixin, RegressorMixin, clone
from ..utils import check_random_state, check_array, check_consistent_length
from ..utils.random import sample_without_replacement
from .base import LinearRegression

_EPSILON = np.spacing(1)

def _dynamic_max_trials(n_inliers, n_samples, min_samples, probability):
    """Determine number trials such that at least one outlier-free subset is
    sampled for the given inlier/outlier ratio.

    Parameters
    ----------
    n_inliers : int
        Number of inliers in the data.

    n_samples : int
        Total number of samples in the data.

    min_samples : int
        Minimum number of samples chosen randomly from original data.

    probability : float
        Probability (confidence) that one outlier-free sample is generated.

    Returns
    -------
    trials : int
        Number of trials.

    """
    inlier_ratio = n_inliers / float(n_samples)
    nom = max(_EPSILON, 1 - probability)
    denom = max(_EPSILON, 1 - inlier_ratio ** min_samples)
    if nom == 1:
        return 0
    if denom == 1:
        return float('inf')
    return abs(float(np.ceil(np.log(nom) / np.log(denom))))

class RANSACRegressor(BaseEstimator, MetaEstimatorMixin, RegressorMixin):
    """RANSAC (RANdom SAmple Consensus) algorithm.

    RANSAC is an iterative algorithm for the robust estimation of parameters
    from a subset of inliers from the complete data set. More information can
    be found in the general documentation of linear models.

    A detailed description of the algorithm can be found in the documentation
    of the ``linear_model`` sub-package.

    Parameters
    ----------
    base_estimator : object, optional
        Base estimator object which implements the following methods:

         * `fit(X, y)`: Fit model to given training data and target values.
         * `score(X, y)`: Returns the mean accuracy on the given test data,
           which is used for the stop criterion defined by `stop_score`.
           Additionally, the score is used to decide which of two equally
           large consensus sets is chosen as the better one.

        If `base_estimator` is None, then
        ``base_estimator=sklearn.linear_model.LinearRegression()`` is used for
        target values of dtype float.

        Note that the current implementation only supports regression
        estimators.

    min_samples : int (>= 1) or float ([0, 1]), optional
        Minimum number of samples chosen randomly from original data. Treated
        as an absolute number of samples for `min_samples >= 1`, treated as a
        relative number `ceil(min_samples * X.shape[0]`) for
        `min_samples < 1`. This is typically chosen as the minimal number of
        samples necessary to estimate the given `base_estimator`. By default a
        ``sklearn.linear_model.LinearRegression()`` estimator is assumed and
        `min_samples` is chosen as ``X.shape[1] + 1``.

    residual_threshold : float, optional
        Maximum residual for a data sample to be classified as an inlier.
        By default the threshold is chosen as the MAD (median absolute
        deviation) of the target values `y`.

    is_data_valid : callable, optional
        This function is called with the randomly selected data before the
        model is fitted to it: `is_data_valid(X, y)`. If its return value is
        False the current randomly chosen sub-sample is skipped.

    is_model_valid : callable, optional
        This function is called with the estimated model and the randomly
        selected data: `is_model_valid(model, X, y)`. If its return value is
        False the current randomly chosen sub-sample is skipped.
        Rejecting samples with this function is computationally costlier than
        with `is_data_valid`. `is_model_valid` should therefore only be used if
        the estimated model is needed for making the rejection decision.

    max_trials : int, optional
        Maximum number of iterations for random sample selection.

    stop_n_inliers : int, optional
        Stop iteration if at least this number of inliers are found.

    stop_score : float, optional
        Stop iteration if score is greater equal than this threshold.

    stop_probability : float in range [0, 1], optional
        RANSAC iteration stops if at least one outlier-free set of the training
        data is sampled in RANSAC. This requires to generate at least N
        samples (iterations)::

            N >= log(1 - probability) / log(1 - e**m)

        where the probability (confidence) is typically set to high value such
        as 0.99 (the default) and e is the current fraction of inliers w.r.t.
        the total number of samples.

    residual_metric : callable, optional
        Metric to reduce the dimensionality of the residuals to 1 for
        multi-dimensional target values ``y.shape[1] > 1``. By default the sum
        of absolute differences is used::

            lambda dy: np.sum(np.abs(dy), axis=1)

    random_state : integer or numpy.RandomState, optional
        The generator used to initialize the centers. If an integer is
        given, it fixes the seed. Defaults to the global numpy random
        number generator.

    Attributes
    ----------
    estimator_ : object
        Best fitted model (copy of the `base_estimator` object).

    n_trials_ : int
        Number of random selection trials until one of the stop criteria is
        met. It is always ``<= max_trials``.

    inlier_mask_ : bool array of shape [n_samples]
        Boolean mask of inliers classified as ``True``.

    References
    ----------
    .. [1] http://en.wikipedia.org/wiki/RANSAC
    .. [2] http://www.cs.columbia.edu/~belhumeur/courses/compPhoto/ransac.pdf
    .. [3] http://www.bmva.org/bmvc/2009/Papers/Paper355/Paper355.pdf
    """

    def __init__(self, base_estimator=None, min_samples=None,
                 residual_threshold=None, is_data_valid=None,
                 is_model_valid=None, max_trials=100,
                 stop_n_inliers=np.inf, stop_score=np.inf,
                 stop_probability=0.99, residual_metric=None,
                 random_state=None):

        self.base_estimator = base_estimator
        self.min_samples = min_samples
        self.residual_threshold = residual_threshold
        self.is_data_valid = is_data_valid
        self.is_model_valid = is_model_valid
        self.max_trials = max_trials
        self.stop_n_inliers = stop_n_inliers
        self.stop_score = stop_score
        self.stop_probability = stop_probability
        self.residual_metric = residual_metric
        self.random_state = random_state

    def fit(self, X, y):
        """Fit estimator using RANSAC algorithm.

        Parameters
        ----------
        X : array-like or sparse matrix, shape [n_samples, n_features]
            Training data.

        y : array-like, shape = [n_samples] or [n_samples, n_targets]
            Target values.

        Raises
        ------
        ValueError
            If no valid consensus set could be found. This occurs if
            `is_data_valid` and `is_model_valid` return False for all
            `max_trials` randomly chosen sub-samples.

        """
        X = check_array(X, accept_sparse='csr')
        y = check_array(y, ensure_2d=False)
        if y.ndim == 1:
            y = y.reshape(-1, 1)
        check_consistent_length(X, y)

        if self.base_estimator is not None:
            base_estimator = clone(self.base_estimator)
        else:
            base_estimator = LinearRegression()

        if self.min_samples is None:
            # assume linear model by default
            min_samples = X.shape[1] + 1
        elif 0 < self.min_samples < 1:
            min_samples = np.ceil(self.min_samples * X.shape[0])
        elif self.min_samples >= 1:
            if self.min_samples % 1 != 0:
                raise ValueError("Absolute number of samples must be an "
                                 "integer value.")
            min_samples = self.min_samples
        else:
            raise ValueError("Value for `min_samples` must be scalar and "
                             "positive.")
        if min_samples > X.shape[0]:
            raise ValueError("`min_samples` may not be larger than number "
                             "of samples ``X.shape[0]``.")

        if self.stop_probability < 0 or self.stop_probability > 1:
            raise ValueError("`stop_probability` must be in range [0, 1].")

        if self.residual_threshold is None:
            # MAD (median absolute deviation)
            residual_threshold = np.median(np.abs(y - np.median(y)))
        else:
            residual_threshold = self.residual_threshold

        if self.residual_metric is None:
            residual_metric = lambda dy: np.sum(np.abs(dy), axis=1)
        else:
            residual_metric = self.residual_metric

        random_state = check_random_state(self.random_state)

        try:  # Not all estimator accept a random_state
            base_estimator.set_params(random_state=random_state)
        except ValueError:
            pass

        n_inliers_best = 0
        score_best = np.inf
        inlier_mask_best = None
        X_inlier_best = None
        y_inlier_best = None

        # number of data samples
        n_samples = X.shape[0]
        sample_idxs = np.arange(n_samples)

        n_samples, _ = X.shape

        for self.n_trials_ in range(1, self.max_trials + 1):

            # choose random sample set
            subset_idxs = sample_without_replacement(n_samples, min_samples,
                                                     random_state=random_state)
            X_subset = X[subset_idxs]
            y_subset = y[subset_idxs]

            # check if random sample set is valid
            if (self.is_data_valid is not None
                    and not self.is_data_valid(X_subset, y_subset)):
                continue

            # fit model for current random sample set
            base_estimator.fit(X_subset, y_subset)

            # check if estimated model is valid
            if (self.is_model_valid is not None and not
                    self.is_model_valid(base_estimator, X_subset, y_subset)):
                continue

            # residuals of all data for current random sample model
            y_pred = base_estimator.predict(X)
            if y_pred.ndim == 1:
                y_pred = y_pred[:, None]

            residuals_subset = residual_metric(y_pred - y)

            # classify data into inliers and outliers
            inlier_mask_subset = residuals_subset < residual_threshold
            n_inliers_subset = np.sum(inlier_mask_subset)

            # less inliers -> skip current random sample
            if n_inliers_subset < n_inliers_best:
                continue

            # extract inlier data set
            inlier_idxs_subset = sample_idxs[inlier_mask_subset]
            X_inlier_subset = X[inlier_idxs_subset]
            y_inlier_subset = y[inlier_idxs_subset]

            # score of inlier data set
            score_subset = base_estimator.score(X_inlier_subset,
                                                y_inlier_subset)

            # same number of inliers but worse score -> skip current random
            # sample
            if (n_inliers_subset == n_inliers_best
                    and score_subset < score_best):
                continue

            # save current random sample as best sample
            n_inliers_best = n_inliers_subset
            score_best = score_subset
            inlier_mask_best = inlier_mask_subset
            X_inlier_best = X_inlier_subset
            y_inlier_best = y_inlier_subset

            # break if sufficient number of inliers or score is reached
            if (n_inliers_best >= self.stop_n_inliers
                    or score_best >= self.stop_score
                    or self.n_trials_
                       >= _dynamic_max_trials(n_inliers_best, n_samples,
                                              min_samples,
                                              self.stop_probability)):
                break

        # if none of the iterations met the required criteria
        if inlier_mask_best is None:
            raise ValueError(
                "RANSAC could not find valid consensus set, because"
                " either the `residual_threshold` rejected all the samples or"
                " `is_data_valid` and `is_model_valid` returned False for all"
                " `max_trials` randomly ""chosen sub-samples. Consider "
                "relaxing the ""constraints.")

        # estimate final model using all inliers
        base_estimator.fit(X_inlier_best, y_inlier_best)

        self.estimator_ = base_estimator
        self.inlier_mask_ = inlier_mask_best
        return self

    def predict(self, X):
        """Predict using the estimated model.

        This is a wrapper for `estimator_.predict(X)`.

        Parameters
        ----------
        X : numpy array of shape [n_samples, n_features]

        Returns
        -------
        y : array, shape = [n_samples] or [n_samples, n_targets]
            Returns predicted values.
        """
        return self.estimator_.predict(X)

    def score(self, X, y):
        """Returns the score of the prediction.

        This is a wrapper for `estimator_.score(X, y)`.

        Parameters
        ----------
        X : numpy array or sparse matrix of shape [n_samples, n_features]
            Training data.

        y : array, shape = [n_samples] or [n_samples, n_targets]
            Target values.

        Returns
        -------
        z : float
            Score of the prediction.
        """
        return self.estimator_.score(X, y)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Methods
------------
reject_duplicates

Classes
------------
WCONWorms

"""
import six
import warnings
from collections import OrderedDict
from six import StringIO
from os import path
import os
import shutil
import json
import jsonschema
import zipfile
import numpy as np
import pandas as pd
idx = pd.IndexSlice

from .wcon_data import parse_data, convert_origin
from .wcon_data import df_upsert, data_as_array
from .wcon_data import get_sorted_ordered_dict
from .wcon_data import reverse_backwards_worms, sort_odict
from .measurement_unit import MeasurementUnit

class WCONWorms():
    """
    A set of worm tracker data for one or more worms, as specified by
    the WCON standard.

    Attributes
    -------------
    units: dict
        May be empty, but is never None since 'units' is required
        to be specified.
    metadata: dict
        If 'metadata' was not specified, metadata is None.
        The values in this dict might be nested into further dicts or other
        data types.
    _data: dictionary of Pandas DataFrames  [private]
    num_worms: int                              [property]
    data_as_dict: dict                          [property]
    data: DataFrame if num_worms == 1 else dict of DataFrames [property]

    [Note: the "files" key is not persisted unless the .load
           factory method is used.]

    Public-Facing Methods
    -------------
    load_from_file   (JSON_path)                [class method]
    save_to_file     (JSON_path, pretty_print)
    to_canon                                    [property]
    __add__                                     [use "+"]
    __eq__                                      [use "=="]

    Usage
    -------------
    # From a string literal:
    from io import StringIO
    w2 = WCONWorms.load(StringIO('{"units":{"t":"s","x":"mm","y":"mm"}, '
                                  '"data":[]}'))

    # WCONWorms.load_from_file accepts any valid WCON, but .save_to_file
    # output is always "canonical" WCON, which makes specific choices about
    # how to arrange and format the WCON file.  This way the functional
    # equality of any two WCON files can be tested by this:

        w1 = WCONWorms.load_from_file('file1.wcon')
        w2 = WCONWorms.load_from_file('file2.wcon')

        assert(w1 == w2)

        # or:

        w1.save_to_file('file1.wcon')
        w2.save_to_file('file2.wcon')

        import filecmp
        assert(filecmp.cmp('file1.wcon', file2.wcon'))

    Custom WCON versions
    --------------------

    Any top-level key other than the basic:

    - files
    - units
    - metadata
    - data

    ... is ignored.  Handling them requires subclassing WCONWorms.

    """
    """
    ================================================================
    Properties
    ================================================================
    """

    basic_keys = ['files', 'units', 'metadata', 'data']

    @property
    def num_worms(self):
        try:
            return self._num_worms
        except AttributeError:
            self._num_worms = len(self.worm_ids)
            return self._num_worms

    @property
    def worm_ids(self):
        try:
            return self._worm_ids
        except AttributeError:
            self._worm_ids = list(self._data.keys())
            return self._worm_ids

    @property
    def data(self):
        """
        Return all worms as one giant DataFrame.  Since this can
        be inefficient for sparse multiworm data, it is only "lazily"
        calculated, i.e. once requested, not at object initialization

        """
        try:
            return self._data_df
        except AttributeError:
            if self.num_worms == 0:
                self._data_df = None
            else:
                # Get a list of all dfs
                dfs = list(self._data.values())
                l = dfs[0]
                # Merge all the worm dfs together into one
                for r in dfs[1:]:
                    l = pd.merge(l, r, left_index=True, right_index=True,
                                 how='outer')
                self._data_df = l

            return self._data_df

    @property
    def data_as_odict(self):
        """
        Return the native ordered-dict-of-DataFrames, the cheapest
        option for sparse multiworm data

        """
        return self._data

    @property
    def schema(self):
        try:
            return self._schema

        except AttributeError:
            # Only load _schema if this method gets called.  Once
            # it's loaded, though, persist it in memory and don't lose it
            here = path.abspath(path.dirname(__file__))

            with open(path.join(here, "wcon_schema.json"), "r") as f:
                self._schema = json.loads(f.read())

            # Now that the schema has been loaded, we can try again
            return self._schema

    @classmethod
    def validate_from_schema(cls, wcon_string):
        jsonschema.validate(json.load(StringIO(wcon_string)), cls().schema)

    @property
    def canonical_units(self):
        """
        A dictionary of canonical versions of the unit for all quantities

        """
        return {k: self.units[k].canonical_unit for k in self.units.keys()}

    @property
    def as_ordered_dict(self):
        """
        Return a representation of the worm as an OrderedDict.  This is most
        useful when saving to a file.

        Returns the canonical version of the data, with units in
        canonical form, and the data converted to canonical form.

        The three keys are:

        - 'units'
        - 'metadata'
        - 'data'

        """
        # Not strictly required by JSON but nice to order the four top-level
        # keys so we use OrderedDict here instead of dict.
        ord_dict = OrderedDict()

        # A dictionary of the canonical unit strings for all quantities except
        # aspect_size, which is generated at runtime.
        units_obj = {k: self.units[k].canonical_unit_string
                     for k in self.units.keys() if k != 'aspect_size'}
        # Sort the units so that every time we save this file, it produces
        # exactly the same output.  Not required in the JSON standard, but
        # nice to have.
        units_obj = get_sorted_ordered_dict(units_obj)
        ord_dict.update({'units': units_obj})

        # The only optional object is "metadata" since "files" is not
        # necessary since we don't currently support saving to more than
        # one chunk.
        if self.metadata:
            # Again, sort the metadata (recursively) so that the same file
            # is produced each time that can stand up to diffing
            metadata_obj = get_sorted_ordered_dict(self.metadata)
            ord_dict.update({'metadata': metadata_obj})

        canonical = self.to_canon
        if canonical._data == {}:
            data_arr = []
        else:
            src = canonical.data_as_odict
            data_arr = []
            for worm_id in src:
                data_arr.extend(data_as_array(src[worm_id]))

        ord_dict.update({'data': data_arr})

        return ord_dict

    """
    ================================================================
    Comparison Methods
    ================================================================
    """
    @classmethod
    def are_units_equal(cls, w1, w2):
        """
        Returns
        ---------
        boolean
            True if w1.units == w2.units, with the only conversion being
            between units that mean the same thing
            (e.g. 'mm' and 'millimetres')
            False otherwise

        """
        if set(w1.units.keys()) != set(w2.units.keys()):
            return False

        for k in w1.units.keys():
            if w1.units[k] != w2.units[k]:
                return False

        return True

    @classmethod
    def is_metadata_equal(cls, w1, w2):
        """
        Returns
        ----------
        boolean
            True if w1.metadata == w2.metadata

        """
        return w1.metadata == w2.metadata

    @classmethod
    def is_data_equal(cls, w1, w2, convert_units=True):
        """
        Parameters
        -------------
        w1, w2: WCONWorms objects
            The objects whose .data attributes will be compared
        convert_units: bool
            If True, the data will first be converted to a standard form
            so that if one worm uses millimetres and the other metres, the
            data can still be properly compared

        TODO:
            Add a "threshold" parameter so that perfect equality is not
            the only option

        """
        # import pdb; pdb.set_trace()
        if w1.num_worms != w2.num_worms:
            return False

        if convert_units:
            d1 = w1.to_canon._data
            d2 = w2.to_canon._data
        else:
            d1 = w1._data
            d2 = w2._data

        for worm_id in w1.worm_ids:
            try:
                df1 = d1[worm_id]
            except KeyError:
                df1 = None
            try:
                df2 = d2[worm_id]
            except KeyError:
                df2 = None

            if (df1 is None) ^ (df2 is None):
                # If one is None but the other is not (XOR), data is not equal
                return False
            elif df1 is None and df2 is None:
                # If both None, they are equal
                continue

            if not pd_equals(df1, df2):
                return False

        return True

    def __eq__(self, other):
        """
        Comparison operator (overloaded)

        Equivalent to .is_data_equal and .is_metadata_equal

        Units are converted

        Special units are not considered

        """
        return (WCONWorms.is_data_equal(self, other) and
                WCONWorms.is_metadata_equal(self, other))

    def __ne__(self, other):
        return not self.__eq__(other)

    def __add__(self, other):
        """
        Addition operator (overloaded)

        """
        return self.merge(self, other)

    @property
    def is_canon(self):
        """
        Returns whether all units are already in their canonical forms.

        """
        for data_key in self.units:
            mu = self.units[data_key]
            if mu.unit_string != mu.canonical_unit_string:
                return False

        return True

    @property
    def to_canon(self):
        """
        Return a new WCONWorms object, with the same .metadata, but with
        .units and .data changed so they are in standard form.

        """
        w = WCONWorms()
        w.metadata = self.metadata
        w.units = self.canonical_units

        # Corner case
        if self._data == {}:
            w._data = OrderedDict({})
            return w

        w._data = OrderedDict()
        for worm_id in self.worm_ids:
            w._data[worm_id] = self._data[worm_id].copy()

        # Go through each "units" key
        for data_key in self.units:
            mu = self.units[data_key]

            # Don't bother to "convert" units that are already in their
            # canonical form.
            if mu.unit_string == mu.canonical_unit_string:
                continue

            tmu = self.units['t']
            for worm_id in w.worm_ids:

                try:
                    # Apply across all worm ids and all aspects
                    mu_slice = \
                        w._data[worm_id].loc[:, idx[:, data_key, :]].copy()

                    w._data[worm_id].loc[:, idx[:, data_key, :]] = \
                        mu_slice.applymap(mu.to_canon)
                except KeyError:
                    # Just ignore cases where there are "units" entries but no
                    # corresponding data
                    pass

                # Special case: change the dataframe index, i.e. the time units
                if tmu.unit_string != tmu.canonical_unit_string:
                    # Create a function that can be applied elementwise to the
                    # index values
                    t_converter = np.vectorize(tmu.to_canon)
                    new_index = t_converter(w._data[worm_id].index.values)

                    w._data[worm_id].set_index(new_index, inplace=True)

        return w

    @classmethod
    def merge(cls, w1, w2):
        """
        Merge two worm groups, in their standard forms.

        Units can differ, but not in their standard forms.

        Metadata must be identical.

        Data can overlap, as long as it does not clash.

        Clashes are checked at a low level of granularity:
        e.g. if two worms have different metadata but the individual metadata
        entries do not conflict, this method will still fail and raise an
        AssertionError.

        """
        if not cls.is_metadata_equal(w1, w2):
            raise AssertionError("Metadata conflicts between worms to be "
                                 "merged.")

        w1c = w1.to_canon
        w2c = w2.to_canon

        for worm_id in w2c.worm_ids:
            if worm_id in w1c.worm_ids:
                try:
                    # Try to upsert w2c's data into w1c.  If we cannot
                    # without an error being raised, the data clashes.
                    w1c._data[worm_id] = df_upsert(w1c._data[worm_id],
                                                   w2c._data[worm_id])
                except AssertionError as err:
                    raise AssertionError("Data conflicts between worms to "
                                         "be merged on worm {0}: {1}"
                                         .format(str(worm_id), err))
            else:
                # The worm isn't in the 1st group, so just add it
                w1c._data[worm_id] = w2c._data[worm_id]

        # Sort w1c's list of worms
        w1c._data = sort_odict(w1c._data)

        # Create a fresh WCONWorms object to reset all the lazily-evaluated
        # properties that may change, such as num_worms, in the merged worm
        merged_worm = WCONWorms()
        merged_worm._data = w1c._data
        merged_worm.metadata = w2c.metadata
        merged_worm.units = w1c.units

        return merged_worm

    """
    ================================================================
    Load / save methods
    ================================================================
    """

    @classmethod
    def validate_filename(cls, JSON_path, is_zipped):
        """
        Perform simple checks on the file path

        JSON_path: str
            The path to the file to be evaluated

        is_zipped: bool
            Whether or not the path is for a zip archive

        """
        assert(isinstance(JSON_path, six.string_types))
        assert(len(JSON_path) > 0)

        if is_zipped:
            if JSON_path[-4:].upper() != '.ZIP':
                raise Exception("A zip archive like %s must have an "
                                "extension ending in '.zip'" % JSON_path)
            else:
                # delete the '.zip' part so the rest can be validated
                JSON_path = JSON_path[:-4]

        warning_message = (' is either less than 5 characters,'
                           'consists of only the extension ".WCON", or '
                           'does not end in ".WCON", the recommended '
                           'file extension.')

        if len(JSON_path) <= 5 or JSON_path[-5:].upper() != '.WCON':
            if is_zipped:
                warnings.warn('Zip file ends properly in .zip, but the '
                              'prefix' + warning_message)
            else:
                warnings.warn('The file name ' + warning_message)

    def save_to_file(self, JSON_path, pretty_print=False,
                     compress_file=False, num_chunks=1):
        """
        Save this object to the path specified.  The object
        will be serialized as a WCON JSON text file.

        Parameters
        -----------
        JSON_path: str
            The path to save this object to.  A warning is raised if the path
            does not end in ".WCON"
        pretty_print: bool
            If True, adds newlines and spaces to make the file more human-
            readable.  Otherwise, the JSON output will use as few characters
            as possible.
        compress_file: bool
            If True, saves a compressed version of the WCON JSON text file
        num_chunks: int
            The number of chunks to break this object into.  If
            num_chunks > 1 then num_chunks files will be created.
            Filenames will have "_1", "_2", etc., added
            to the end of the filename after the last path separator
            (e.g. "/") and then, before the last "." (if any)

        """
        if num_chunks > 1:
            raise NotImplementedError("Saving a worm to more than one chunk "
                                      "has not yet been implemented")

        self.validate_filename(JSON_path, compress_file)

        with open(JSON_path, 'w') as outfile:
            json.dump(self.as_ordered_dict, outfile,
                      indent=4 if pretty_print else None)

        if compress_file:
            # Zip the file to a TEMP file, then rename to the original,
            # overwriting it with the zipped archive.
            zf = zipfile.ZipFile(JSON_path + '.TEMP',
                                 'w', zipfile.ZIP_DEFLATED)
            zf.write(JSON_path)
            zf.close()
            os.rename(JSON_path + '.TEMP', JSON_path)

    @classmethod
    def load_from_file(cls, JSON_path,
                       load_prev_chunks=True,
                       load_next_chunks=True,
                       validate_against_schema=True):
        """
        Factory method returning a merged WCONWorms instance of the file
        located at JSON_path and all related "chunks" as specified in the
        "files" element of the file.

        Uses recursion if there are multiple chunks.

        Parameters
        -------------
        JSON_path: str
            A file path to a file that can be opened
        validate_against_schema: bool
            If True, validate before trying to load the file, otherwise don't.
            jsonschema.validate takes 99% of the compute time for large files
            so use with caution.
        load_prev_chunks: bool
            If a "files" key is present, load the previous chunks and merge
            them with this one.  If not present, return only the current
            file's worm.
        load_next_chunks: bool
            If a "files" key is present, load the next chunks and merge
            them with this one.  If not present, return only the current
            file's worm.

        """
        print("Loading file: " + JSON_path)

        is_zipped = zipfile.is_zipfile(JSON_path)

        cls.validate_filename(JSON_path, is_zipped)

        # Check if the specified file is compressed
        if is_zipped:
            zf = zipfile.ZipFile(JSON_path, 'r')

            zf_namelist = zf.namelist()
            if len(zf_namelist) <= 0:
                raise Exception("Filename %s is a zip archive, which is fine, "
                                "but the archive does not contain any files.")
            elif len(zf_namelist) == 1:
                # Just one file is in the archive.
                print("The file is a zip archive with one file.  Attempting "
                      "to uncompress and then load.")
                wcon_bytes = zf.read(zf.namelist()[0])
                wcon_string = wcon_bytes.decode("utf-8")
                infile = StringIO(wcon_string)
                w_current = cls.load(infile, validate_against_schema)
            else:
                print("The zip archive contains multiple files.  We will "
                      "extract to a temporary folder and then try to load "
                      "the first file in the archive, then delete the "
                      "temporary folder.")
                # Note: the first file is all we should need since we assume
                #       the files in the archive are linked together using
                #       their respective JSON "files" entries

                # Make a temporary archive folder
                cur_path = os.path.abspath(os.path.dirname(JSON_path))
                archive_path = os.path.join(cur_path, '_zip_archive')
                if os.path.exists(archive_path):
                    raise Exception("Archive path %s already exists!"
                                    % archive_path)
                else:
                    os.makedirs(archive_path)

                # Extract zip archive to temporary folder
                for name in zf_namelist:
                    zf.extract(name, archive_path)
                zf.close()

                # Call load_from_file on the first file
                first_path = os.path.join(archive_path, zf_namelist[0])
                w = cls.load_from_file(first_path)

                # Delete the temporary folder
                shutil.rmtree(archive_path, ignore_errors=True)

                return w
        else:
            # The file is not a zip file, so assume it's just plaintext JSON
            with open(JSON_path, 'r') as infile:
                w_current = cls.load(infile, validate_against_schema)

        # CASE 1: NO "files" OBJECT, hence no multiple files.  We are done.
        w_cur = w_current
        if w_current.files is None:
            return w_current
        elif (('next' not in w_cur.files) and ('prev' not in w.cur.files)):
            # CASE 2: "files" object exists but no prev/next, assume nothing is
            # there
            return w_current
        else:
            # The merge operations below will blast away the .files attribute
            # so we need to save a local copy
            current_files = w_current.files

        # OTHERWISE, CASE 2: MULTIPLE FILES

        # The schema guarantees that if "files" is present,
        # "current", will exist.  Also, that "current" is not
        # null and whose corresponding value is a string at least one
        # character in length.
        cur_ext = current_files['current']

        # e.g. cur_filename = 'filename_2.wcon'
        # cur_ext = '_2', prefix = 'filename', suffix = '.wcon'
        cur_filename = JSON_path
        name_offset = cur_filename.find(cur_ext)
        if name_offset == -1:
            raise AssertionError(
                'Mismatch between the filename given in the file "' +
                cur_ext +
                '" and the file we loaded from "' +
                cur_filename +
                '".')
        path_string = cur_filename[:name_offset]

        load_chunks = {'prev': load_prev_chunks,
                       'next': load_next_chunks}

        for direction in ['prev', 'next']:
            # If we are supposed to load the previous chunks, and one exists,
            # load it and merge it with the current chunk
            # Same with the "next" chunks
            if (load_chunks[direction] and
                    current_files is not None and
                    current_files[direction] is not None and
                    len(current_files[direction]) > 0):
                cur_load_prev_chunks = (direction == 'prev')
                cur_load_next_chunks = (direction == 'next')

                new_file_name = path_string + current_files[direction][0]

                w_new = cls.load_from_file(new_file_name,
                                           cur_load_prev_chunks,
                                           cur_load_next_chunks,
                                           validate_against_schema)
                w_current = w_current + w_new

        # If no merging took place, we'll still need to delete the "files"
        # attribute if it's present (i.e. if both "prev" and "next" were null):
        if hasattr(w_current, "files"):
            del(w_current.files)

        return w_current

    @classmethod
    def load(cls, JSON_stream, validate_against_schema=True):
        """
        Factory method to create a WCONWorms instance

        This does NOT load chunks, because a file stream does not
        have a file name.  In order to load chunks, you must invoke the
        factory method load_from_file.  You will be passing it a file path
        from which it can find the other files/chunks.

        Parameters
        -------------
        JSON_stream: a text stream implementing .read()
            e.g. an object inheriting from TextIOBase
        validate_against_schema: bool
            If True, validate before trying to load the file, otherwise don't.
            jsonschema.validate takes 99% of the compute time for large files
            so use with caution.

        """
        w = cls()

        serialized_data = JSON_stream.read()

        # Load the whole JSON file into a nested dict.  Any duplicate
        # keys raise an exception since we've hooked in reject_duplicates
        root = json.loads(serialized_data, object_pairs_hook=reject_duplicates)

        # ===================================================
        # BASIC TOP-LEVEL VALIDATION AGAINST THE SCHEMA

        # Validate the raw file against the WCON schema
        if validate_against_schema:
            jsonschema.validate(root, w.schema)

        # ===================================================
        # HANDLE THE REQUIRED ELEMENTS: 'units', 'data'

        w.units = root['units']

        for key in w.units:
            w.units[key] = MeasurementUnit.create(w.units[key])

        # The only data key without units should be aspect_size, since it's
        # generated during the construction of the pandas dataframe
        # it is a dimensionless quantity
        w.units['aspect_size'] = MeasurementUnit.create('')

        if len(root['data']) > 0:
            w._data = parse_data(root['data'])

            # Shift the coordinates by the amount in the offsets 'ox' and 'oy'
            for worm_id in w.worm_ids:
                convert_origin(w._data[worm_id])

                # Any worms with head=='R' should have their
                # coordinates reversed and head reset to 'L'
                reverse_backwards_worms(w._data[worm_id])
        else:
            # "data": {}
            w._data = OrderedDict({})

        # Raise error if there are any data keys without units
        units_keys = set(w.units.keys())

        for worm_id in w._data:
            df = w._data[worm_id]
            if df is None:
                data_keys = set()
            else:
                data_keys = set(df.columns.get_level_values(1))

            # "head" and "ventral" don't require units.
            keys_missing_units = data_keys - \
                units_keys - set(['head', 'ventral'])
            if keys_missing_units != set():
                raise AssertionError('In worm ' + str(worm_id) + ', the '
                                     'following data keys are missing '
                                     'entries in the "units" object: ' +
                                     str(keys_missing_units))

        # ===================================================
        # HANDLE THE OPTIONAL ELEMENTS: 'files', 'metadata'

        if 'files' in root:
            w.files = root['files']

            # Handle the case of a single 'next' or 'prev' entry, by
            # wrapping it in an array, so we can reliably assume that
            # entries are always wrapped in arrays.
            for direction in ['next', 'prev']:
                if hasattr(w.files, direction):
                    if isinstance(getattr(w.files, direction), str):
                        setattr(w.files, direction,
                                [getattr(w.files, direction)])
        else:
            w.files = None

        if 'metadata' in root:
            w.metadata = root['metadata']
        else:
            w.metadata = None

        return w

def pd_equals(df1, df2):
    """
    I don't use DataFrame.equals because it returned False for no
    apparent reason with one of the centroid unit tests

    """
    if not df1.columns.identical(df2.columns):
        return False

    if not df1.index.identical(df2.index):
        return False

    try:
        pd.util.testing.assert_frame_equal(df1, df2)
    except AssertionError:
        return False

    return True

def reject_duplicates(ordered_pairs):
    """Reject duplicate keys."""
    unique_dict = {}
    for key, val in ordered_pairs:
        if key in unique_dict:
            raise KeyError("Duplicate key: %r" % (key,))
        else:
            unique_dict[key] = val

    return unique_dict