#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
A module of methods to create the pandas DataFrame representing the "data"
array.

"""
import six
import gc
import os
import warnings
import numpy as np
import pandas as pd
import itertools
from multiprocessing import Queue, Process
import multiprocessing
import psutil
import time
if six.PY3:
    from functools import reduce
from collections import OrderedDict
idx = pd.IndexSlice

# "Aspect" is my term for the articulation points at a given time frame
# So if you track 49 skeleton points per frame, then x and y have
# 49 aspects per frame, but ox is only recorded once per frame, or
# even once across the whole video.
elements_with_aspect = ['x', 'y']
elements_without_aspect = ['ox', 'oy', 'cx', 'cy', 'head', 'ventral']
basic_data_keys = elements_with_aspect + elements_without_aspect
supported_data_keys = basic_data_keys + ['id', 't']
MIN_FRAMES_FOR_MULTIPROCESSING = 3000  # This is somewhat arbitrary
USE_MULTIPROCESSING = False    # If False, multiprocessing is never used


def get_mask(arr, desired_key):
    """
    In an array of dicts, obtain a mask on arr of elements having
    the desired key

    Parameters
    ----------
    arr: list-like array of dicts
    desired_key: str

    Returns
    ----------
    Boolean numpy array of same size as arr

    """
    # Find elements having a time aspect
    desired_indexes = [i for (i, s) in enumerate(arr) if desired_key in s]
    desired_indexes = np.array(desired_indexes)

    # Transform our list of indexes into a mask on the data array
    return np.bincount(desired_indexes, minlength=len(arr)).astype(bool)


def df_upsert(src, dest):
    """
    Append src to dest, doing an "update/insert" where
    conflicts cause an AssertionError to be raised.

    Return the new dest

    src always corresponds to a data segment, which is always for exactly
    one worm.  dest is our target, containing possibly many worms.  So
    we have three scenarios:

    1. src is a new worm, not in dest.  thus basically do a FULL OUTER JOIN.
    2. OTHERWISE src is a worm already in dest. in this case rows can be:
       - new timepoints.  so UNION
       - revisions to old timepoints
        In this last case if any of the revisions are changes, error.
        If not, do not insert any duplicate rows, but any new rows, UNION.

    Parameters
    -----------
    src, dest: pandas DataFrames
        src is the one to be added to dest
        dest is the one that is modified in place

    # TODO: concatenate using a list comprehension as calling it
    # iteratively like this causes a performance hit (see
    # http://pandas.pydata.org/pandas-docs/stable/merging.html#concatenating-objects
    # "Note It is worth noting however, that concat (and therefore append)
    # makes a full copy of the data, and that constantly reusing this
    # function can create a signifcant performance hit. If you need
    # to use the operation over several datasets, use a list comprehension."

    """
    # Add all columns for worms that don't currently exist in dest
    # but have entries on the same timeframes
    if not any(src.columns.isin(dest.columns)):
        # Scenario 1: src is a new worm.  So OUTER JOIN
        dest = pd.merge(dest, src, left_index=True,
                        right_index=True, how='outer')
    else:
        # Scenario 2: src is worm already in dest.

        # Handle the timepoints NOT shared with dest
        if any(~src.index.isin(dest.index)):
            # Add all rows that don't currently exist in dest
            dest = pd.concat([dest, src[~src.index.isin(dest.index)]], axis=0)

        # Handle the timepoints are shared with dest
        if any(src.index.isin(dest.index)):
            # Obtain a sliced version of dest, showing only
            # the columns and rows shared with the src
            dest_sliced = \
                dest.loc[dest.index.isin(src.index),
                         dest.columns.isin(src.columns)]

            src_sliced = \
                src.loc[src.index.isin(dest.index),
                        src.columns.isin(dest.columns)]

            # Sort our slices so they will be lined up for comparison
            dest_sliced.sort_index(inplace=True)
            src_sliced.sort_index(inplace=True)

            # Obtain a mask of the conflicts in the current segment
            # as compared with all previously loaded data.  That is:
            # NaN NaN = False
            # NaN 2   = False
            # 2   2   = False
            # 2   3   = True
            # 2   NaN = True
            data_conflicts = (pd.notnull(dest_sliced) &
                              (dest_sliced != src_sliced))

            if data_conflicts.any().any():
                raise AssertionError("Data from this segment conflicted "
                                     "with previously loaded data:\n",
                                     data_conflicts)

            # Replace any rows that do exist with the src version
            dest.update(src_sliced)

    # Sort the time series indices
    dest.sort_index(axis=0, inplace=True)

    return dest


def convert_origin(df):
    """
    Offset the coordinates and centroid by the offsets if available.

    This code will ensure that there is no offset besides a centroid,
    if the centroid exists.  If the centroid exists, all x and y coordinates
    are relative to this centroid.

    In pseudocode:

    For each worm and time frame:
        If 'ox' is not NaN:
            Add 'ox' to 'x'
            If 'cx' is not NaN:
                Add 'ox' to 'cx'
                Subtract 'cx' from 'x'

        (Also do the same for y)

    After this is done, drop the offset columns 'ox', and 'oy'
    entirely from the dataframe.

    Parameters
    ------------
    df: Pandas DataFrame

    Returns
    ------------
    None.  Modifies `df` in place.

    """
    offset_keys = ['ox', 'oy']
    centroid_keys = ['cx', 'cy']
    coord_keys = ['x', 'y']

    for worm_id in df.columns.get_level_values(0).unique():
        cur_worm = df.loc[:, (worm_id)]

        for offset, centroid, coord in zip(offset_keys,
                                           centroid_keys, coord_keys):

            # Note: This code block uses `x` as the stylized example for
            # variable naming purposes, but be assured that the enclosing
            # `for` loop loops through both `x` and `y`.

            if offset in cur_worm.columns.get_level_values(0):
                # Consider offset is 0 if not available in a certain frame
                ox_column = cur_worm.loc[:, (offset)].fillna(0)

                # Shift our 'x' values by offset
                all_x_columns = cur_worm.loc[:, (coord)]
                ox_affine_change = (np.array(ox_column) *
                                    np.ones(all_x_columns.shape))
                all_x_columns += ox_affine_change

                if centroid in cur_worm.columns.get_level_values(0):
                    cx_column = cur_worm.loc[:, (centroid)]
                    # Shift the centroid by the offset
                    cx_column += ox_column

                    # Now make the centroid our new offset, since the rule
                    # is that if the offset exists, the centroid is not
                    # the offset, but we want it to be.
                    cx_affine_change = (np.array(cx_column) *
                                        np.ones(all_x_columns.shape))
                    all_x_columns -= cx_affine_change

                    # Now assign these values back to the passed dataframe df
                    df.loc[:, (worm_id, centroid)] = cx_column.values

                # Now assign these values back to the passed dataframe df
                df.loc[:, (worm_id, coord)] = all_x_columns.values

                # Now reset our 'ox' values to zero.
                if offset in cur_worm.columns.get_level_values(0):
                    df.loc[:, (worm_id, offset)] = np.zeros(ox_column.shape)

    # Drop the offset columns entirely from the dataframe.
    # This is so DataFrames with and without offsets
    # will show as comparing identically.
    for offset_key in offset_keys:
        df.drop(offset_key, axis=1, level='key', inplace=True)

    # Because of a known issue in Pandas
    # (https://github.com/pydata/pandas/issues/2770), the dropped columns
    # remain in the "levels" attribute of MultiIndex, even if they don't
    # appear in the "labels" and are thus not 'observed'.
    # The workaround is to reconstitute the MultiIndex from tuples:
    df.columns = pd.MultiIndex.from_tuples(df.columns.values,
                                           names=df.columns.names)


def reverse_backwards_worms(df, coord_keys=['x', 'y']):
    """
    Reverse all worms in all time frames with head == 'R':

    - Reverse the coordinates
    - Change head to 'L'

    Parameters
    ------------
    df: Pandas DataFrame
    coord_keys: list of strings
        The corresponding coordinates to be reversed

    Returns
    ------------
    None.  Modifies `df` in place.

    """
    # TODO: this method is not fully implemented
    return

    if 'head' not in df.columns.get_level_values(level='key'):
        return

    mask_of_times_to_reverse = df.loc[:, idx[:, 'head', :]] == 'R'

    for worm_id in df.columns.get_level_values(level='id').unique():
        cur_worm = df.loc[:, idx[worm_id, :, :]]
        cur_keys = cur_worm.columns.get_level_values(level='key')

        if 'head' not in cur_keys:
            # Nothing to do
            continue

        cur_mask = mask_of_times_to_reverse.loc[:, idx[worm_id, 'head', 0]]

        for key in coord_keys:
            if key not in cur_keys:
                continue

            thing_to_reverse = df.loc[cur_mask, idx[worm_id, 'x', :]]

            # TODO: I need to reverse "thing_to_reverse", accounting for
            # aspect size.  [::-1] just reverses the row order.
            pass

    pass


def parse_data(data):
    """
    Parse the array of entries conforming to the WCON schema definition
    for "_data" in the root object.  The canonical example is
    that of worm "skeleton" (midline) information over time.

    This could be the standard "data" array from the root object, or
    some custom array that needs to be processed

    Note that all elements are required to have "id", "t", and "x" and "y"
    entries.

    Returns
    --------
    An ordered dict of DataFrames, one DataFrame per worm id.

    """
    # If data is single-valued, wrap it in a list so it will be just
    # a special case of the array case.
    if isinstance(data, dict):
        data = [data]

    data = np.array(data)

    # We only care about data that have all the mandatory fields.
    # the Custom Feature Type 2 objects are suppresed by this filter:
    is_time_series_mask = get_mask(data, 't')
    has_id_mask = get_mask(data, 'id')
    data = data[is_time_series_mask & has_id_mask]

    # Clean up and validate all time-series data segments
    _validate_time_series_data(data)

    df_odict = _obtain_time_series_data_frame(data)

    return df_odict


def _obtain_time_series_data_frame(time_series_data):
    """
    Obtain a time-series pandas DataFrame

    Parameters
    ------------
    time_series_data: list
        All entries must be lists of dicts, all of which
        must have 'id' and 't'

    Returns
    ----------
    ordered dictionary of pandas DataFrames
        - keys are the worm_ids
        - Time-series data goes into this Pandas DataFrame
        The dataframe will have t as index, and multilevel columns
        with id at the first level and all other keys at second level.

    """

    # Convert each segment into a dataframe
    data_segment_dfs = []

    # Consider only time-series data stamped with an id:
    for data_segment in time_series_data:
        # Add this data_segment to a Pandas dataframe
        worm_id = data_segment['id']
        segment_keys = np.array([k for k in data_segment.keys()
                                 if k not in ['t', 'id']])

        cur_timeframes = np.array(data_segment['t']).astype(float)

        # Create our column names as the cartesian product of
        # the segment's keys and the id of the segment
        # Only elements articulated across the skeleton, etc, of the worm
        # have the "aspect" key though
        cur_elements_with_aspect = \
            [k for k in elements_with_aspect if k in segment_keys]
        cur_elements_without_aspect = ['aspect_size'] + \
            [k for k in elements_without_aspect if k in segment_keys]

        # We want to be able to fit the largest aspect size in our
        # DataFrame
        max_aspect_size = int(max([k[0] for k in data_segment['aspect_size']]))

        key_combos = list(itertools.product([worm_id],
                                            cur_elements_with_aspect,
                                            range(max_aspect_size)))
        key_combos.extend(list(itertools.product([worm_id],
                                                 cur_elements_without_aspect,
                                                 [0])))

        column_type_names = ['id', 'key', 'aspect']
        cur_columns = pd.MultiIndex.from_tuples(key_combos,
                                                names=column_type_names)

        # e.g. if this segment has only 'x', 'y', that's what we'll be
        # looking to add to our dataframe data staging in the next step
        cur_data_keys = cur_elements_with_aspect + \
            cur_elements_without_aspect

        # We must pad the timeframes where the data doesn't have maximal
        # aspect or else the concatenation step below will fail.
        for k in cur_elements_with_aspect:
            for i in range(len(cur_timeframes)):
                data_segment[k][i] = (
                    data_segment[k][i] +
                    [np.NaN] * (max_aspect_size - len(data_segment[k][i])))

        num_timeframes = len(cur_timeframes)

        # Stage the data for addition to our DataFrame.
        # Shape KxI where K is the number of keys and
        #                 I is the number of "aspects"
        cur_data = _stage_dataframe_data(num_timeframes,
                                         data_segment, cur_data_keys)

        cur_df = pd.DataFrame(cur_data, columns=cur_columns)

        cur_df.index = cur_timeframes
        cur_df.index.names = 't'

        # We want the index (time) to be in order.
        cur_df.sort_index(axis=0, inplace=True)

        # Apparently multiindex must be sorted to work properly:
        cur_df.sort_index(axis=1, inplace=True)

        # If we don't do this, for very large files (>50 MB) the memory
        # footprint grows until the program crashes
        # gc.collect()

        data_segment_dfs.append(cur_df)

    # Our DataFrame to return
    df_odict = OrderedDict()

    # Obtain a list of lists of the worm ids
    worm_ids = [cur_df.columns.levels[0].values for cur_df in data_segment_dfs]

    # To see if any worms appear more than once, flatten the list and
    # check for duplicates
    worm_ids_flat = [item for sublist in worm_ids for item in sublist]

    # If worm_ids are mutually distinct in each dataframe, then
    if len(set(worm_ids_flat)) == len(worm_ids):
        df_odict = OrderedDict(zip(worm_ids_flat, data_segment_dfs))
        # l = data_segment_dfs[0]
        # memory_use_log = [l.memory_usage().sum()]
        # memory_use_indiv = [l.memory_usage().sum()]
        # for r in data_segment_dfs[1:]:
        #   l = pd.merge(l, r, left_index=True, right_index=True, how='outer')
        # time_df = l
    else:
        # Otherwise we have to iterate
        # TODO: we could do the above time shortcut with all available
        #       mutually distinct segments, rather than only doing it if
        #       ALL segments are mutually distinct
        for cur_df in data_segment_dfs:
            worm_id = cur_df.columns.levels[0].values[0]
            # Add the segment to our main DataFrame
            if worm_id not in df_odict:
                # if it's not there, INSERT
                df_odict[worm_id] = cur_df
            else:
                # if it is already there, UPDATE
                df_odict[worm_id] = df_upsert(src=cur_df,
                                              dest=df_odict[worm_id])

    for worm_id in df_odict:
        # We want the index (time) to be in order.
        df_odict[worm_id].sort_index(axis=0, inplace=True)
        # The columns also need to be in order.
        df_odict[worm_id].sort_index(axis=1, inplace=True)

        # We have to do this because somehow with entire worms who have 'head'
        # or 'ventral' columns, all their columns (even the numeric ones)
        # have dtype = object !  We want dtype=float64 (or some float anyway)
        # (Ignore the FutureWarning here about convert_objects being
        # deprecated, because to_numeric only acts on Series and I don't know
        # which ones we need to convert.)
        with warnings.catch_warnings():
            warnings.filterwarnings(action="ignore", category=FutureWarning)
            df_odict[worm_id] = \
                df_odict[worm_id].convert_objects(convert_numeric=True)

        # If 'head' or 'ventral' is NaN, we must specify '?' since
        # otherwise, when saving this object, to specify "no value" we would
        # have to create a new data segment with no 'head' column at all.
        # But since both 'head' and 'ventral' have this issue, we might need
        # segments with just 'head', just 'ventral', and both 'head' and
        # 'ventral', and none.  That's too unwieldy for a rare use case.
        # So instead we treat nan as '?'.

        # We must replace NaN with None, otherwise the JSON encoder will
        # save 'NaN' as the string and this will get rejected by our schema
        # on any subsequent loads
        # Note we can't use .fillna(None) due to this issue:
        # https://github.com/pydata/pandas/issues/1972
        df_keys = set(df_odict[worm_id].columns.get_level_values('key'))
        for k in ['head', 'ventral']:
            if k in df_keys:
                cur_slice = df_odict[worm_id].loc[:, idx[:, k, :]]
                df_odict[worm_id].loc[:, idx[:, k, :]] = \
                    cur_slice.fillna(value=np.nan)

        # Make sure aspect_size is a float, since only floats are nullable:
        if 'aspect_size' in df_keys:
            df_odict[worm_id].loc[:, idx[:, 'aspect_size', :]] = \
                df_odict[worm_id].loc[:, idx[:, 'aspect_size', :]] \
                .astype(float)

    return sort_odict(df_odict)


def __pivot_data_for_staging(data_segment, cur_data_keys, index_range,
                             core_index=None, queue=None):
    """
    Transform the WCON data segment into one properly arranged for adding to
    a pandas DataFrame

    i.e. Pivot the data properly for the dataframe, for a certain range
    of frames f where i_start <= f < i_end, where
    i_start = min(index_range) and i_end = max(index_range).

    Parameters
    ----------
    data_segment: list of lists of floats
        Taken directly from the WCON file being loaded
    cur_data_keys: list of 1-character strings
        The keys, e.g. 'x', 'y', that need to be obtained
    index_range: range object
        The timeframes across which the transformation is to be applied
    core_index: int
        None if this is not a sub-process
        Otherwise, this is the core that the sub-process should be pinned to
        (there is a problem with subprocesses using multiple cores when
        numpy is involved, so we must explicitly pin them to separate cores)
    queue: multiprocessing.Queue object
        None if this is not a sub-process
        Otherwise, this is a sub-process, and we should append our result
        to queue rather than return a value

    """
    # DEBUG: This code definitely gets used across multiple cores
    #        use this to verify multiprocessing
    # myList = np.random.randint(0, 10, 3333333/4)
    # finalSum = sum(myList)

    # Set process affinity if we are multiprocessing
    if core_index:
        cur_pid = os.getpid()
        p = psutil.Process(cur_pid)
        # print("Current pid:" + str(cur_pid))
        # print("Parent pid:" + str(os.getppid()))
        # print("PSUTIL CPU affinity:" + str(p.cpu_affinity()))
        p.cpu_affinity([core_index])
        # print("PSUTIL CPU affinity:" + str(p.cpu_affinity()))

    # Pivot the data
    cur_data_slice = np.array(
        [np.concatenate(
            [data_segment[k][i] for k in cur_data_keys]
        ) for i in index_range]
    )

    # Add to the queue only if we are multiprocessing
    if queue is None:
        return cur_data_slice
    else:
        # We put a duple into the queue, the first entry
        # being the correct final order of the data, the
        # second entry being the data we just calculated
        queue.put((core_index, cur_data_slice))


def _stage_dataframe_data(num_timeframes, data_segment, cur_data_keys):
    """
    Transform the WCON data segment into one properly arranged for adding to
    a pandas DataFrame.

    As this is an "embarassingly parallel" problem since each frame can be
    processed independently, we split the frames across multiple processes
    if
    independent

    (This version attempts to use multiple process)

    """
    num_cores = multiprocessing.cpu_count()

    # Don't use multiple processes unless we have many frames of data,
    # and unless we have multiple CPUs to spread the work around
    if ((not USE_MULTIPROCESSING) or
            num_timeframes < MIN_FRAMES_FOR_MULTIPROCESSING or num_cores <= 1):
        return __pivot_data_for_staging(data_segment, cur_data_keys,
                                        range(num_timeframes))

    # Create a Queue to share results
    q = Queue()

    # Create 4 sub-processes to do the work
    splits = list(range(0, num_timeframes, num_timeframes // num_cores))
    # If the number of frames is not divisible by the step length,
    # range will fall short of ranging up to the num_timeframes
    # so here we overwrite the last endpoint
    if not num_timeframes % num_cores:
        # If num_timeframes is divisible by num_cores, add an entry
        splits.append(num_timeframes)
    else:
        # Otherwise, we already have enough entries, so just extend
        # the final one
        splits[-1] = num_timeframes

    processes = []
    for i in range(num_cores):
        processes.append(Process(target=__pivot_data_for_staging, args=(
                                 data_segment, cur_data_keys,
                                 range(splits[i], splits[i + 1]),
                                 i, q)))

    for i in range(len(processes)):
        processes[i].start()

    data_pieces = {}
    # Grab 4 values from the queue, one for each process
    # They are calculated by subprocesses and hence might
    # arrive in the queue out of order, so we collect them into
    # a dictionary keyed with their correct final order.
    for i in range(len(processes)):
        # Set block=True to block until we get a result
        data_piece = q.get(True)
        # data_piece is a tuple with the first entry providing the order
        data_pieces[data_piece[0]] = data_piece[1]

    # Now put the pieces together in the correct order
    cur_data = []
    for i in range(len(processes)):
        cur_data.extend(data_pieces[i])

    # Combine all the processes' data into one list
    cur_data = np.array(cur_data)

    for i in range(len(processes)):
        processes[i].join()

    return cur_data


def _validate_time_series_data(time_series_data):
    """
    Validate and standardise all time-series data segments in the
    data dictionary provided

    Parameters
    -----------
    data_series_data: array
        Dictionary extracted from JSON, but only the entries with
        a time series.
        All elements must be lists of dicts, all of which
        must have 't' entries

    """
    canonical_elements = ['id', 't', 'x', 'y', 'cx', 'cy', 'ox', 'oy',
                          'head', 'ventral', 'aspect_size']

    for (data_segment_index, data_segment) in enumerate(time_series_data):
        # Filter the data_segment to ignore non-canonical elements
        if six.PY3:
            data_segment = {k: v for (k, v) in data_segment.items()
                            if k in canonical_elements}
        else:
            data_segment = {k: v for (k, v) in data_segment.iteritems()
                            if k in canonical_elements}

        segment_keys = [k for k in data_segment.keys() if k != 'id']

        # If one axis is present, the other must be as well
        assert(not(('cx' in segment_keys) ^ ('cy' in segment_keys)))
        assert(not(('ox' in segment_keys) ^ ('oy' in segment_keys)))

        """
        We require elements to be wrapped in arrays.  They may come in
        various singleton formats, so we must convert.  The final state
        must be:

        "t": [1.5], "x": [[6, 7, 8]], "ox": [[8.2]]

        or for multiple timeframes:

        "t": [1.5, 1.6], "x": [[6, 7, 8], [6.2, 7.2, 8.2]],
        "ox": [[8.2], [8.2]]

        The data might arrive in the following formats:

                                   elements with aspect
                            singleton      array               array of arrays

        TIME singleton  "t":1.5, "x":6  "x":1.5, "x": [6,7,8]    NOT ALLOWED

        TIME array        NOT ALLOWED   "t":[1.5, 1.8]         "t":[1.5, 1.8]
                                        "x":[6.3, 6.2]         "x":[[6,7,8],
                                                                    [8,9,10]]

                                   elements without aspect
                            singleton      array               array of arrays
        TIME singleton  "t":1.5, "ox":6  "x":1.5, "ox": [8]    NOT ALLOWED

        TIME array      "t":[1.5,1.8]    "t":[1.5, 1.8]        NOT ALLOWED
                        "ox":8           "ox":[8, 8.2]

        So in all but one of these cases, we must wrap the data in more
        brackets.

        """
        # HANDLE TIME ('t')
        time_is_singleton = not isinstance(data_segment['t'], list)
        if time_is_singleton:
            data_segment['t'] = [data_segment['t']]
        num_timeframes = len(data_segment['t'])

        # HANDLE ALL OTHER KEYS (besides 'id' and 't')
        for subkey in elements_with_aspect + elements_without_aspect:
            if subkey in segment_keys:
                if not isinstance(data_segment[subkey], list):
                    # SINGLETON CASE
                    if time_is_singleton:
                        data_segment[subkey] = [[data_segment[subkey]]]
                    else:
                        if subkey in elements_without_aspect:
                            # Broadcast aspectless elements to be
                            # length n = subelement_length if it's
                            # just being shown once right now
                            # (e.g. for 'ox', 'oy', etc.)
                            data_segment[subkey] = [[x] for x in [
                                data_segment[subkey]] * num_timeframes]
                        else:
                            raise Exception(
                                "Error with element '%s' in data "
                                "segment %s: time is array but "
                                "element is singleton." %
                                (subkey, str(data_segment)))
                elif not isinstance(data_segment[subkey][0], list):
                    # ARRAY CASE
                    if time_is_singleton:
                        data_segment[subkey] = [data_segment[subkey]]
                    else:
                        data_segment[subkey] = [[x] for x in
                                                data_segment[subkey]]
                else:
                    # ARRAY OF ARRAYS CASE
                    if time_is_singleton:
                        raise Exception("Error with element '%s' in data "
                                        "segment %s: time is singleton but "
                                        "element is an array of arrays." %
                                        (subkey, str(data_segment)))
                    else:
                        if subkey in elements_without_aspect:
                            # We could allow this case but it makes no sense
                            # for the file to have aspectless data
                            # double-wrapped so let's reject it.
                            raise Exception("Error with element '%s' in data "
                                            "segment %s: element is "
                                            "aspectless but "
                                            "element is an array of arrays." %
                                            (subkey, str(data_segment)))
                        else:
                            # This is the one case where all is good.
                            pass

        # Validate that all elements have the same number of timeframes
        element_timeframes = [len(data_segment[subkey])
                              for subkey in segment_keys]

        # Now we can assure ourselves that num_timeframes is
        # well-defined; if not, raise an error.
        if len(set(element_timeframes)) > 1:
            raise AssertionError("Error: Elements must have all have "
                                 "the same number of timeframes.")

        # In each time, the aspect size could change, so we need to keep
        # track of it since the dataframe will ultimately have columns
        # for the maximal aspect size and it won't otherwise be possible
        # to determine what the aspect size is in each timeframe

        # First let's validate that the aspect size is identical
        # across data elements in each time frame:
        aspect_size_over_time = []
        for t in range(num_timeframes):
            # The x and y arrays for element i of the data segment
            # must have the same length
            try:
                cur_aspect_sizes = [len(data_segment[k][t]) for k in
                                    elements_with_aspect]
            except TypeError as err:
                raise TypeError("In the following data segment, an "
                                "element with aspect (x, y, etc.) was not "
                                "double-wrapped in arrays, even "
                                "though time ('t') was. {0}".format(err))

            if len(set(cur_aspect_sizes)) > 1:
                raise AssertionError(
                    "Error: Aspects x and y, etc. must have same "
                    "length for data segment " + str(data_segment_index) +
                    " and time index " + str(data_segment['t'][t]))
            else:
                aspect_size_over_time.append([cur_aspect_sizes[0]])

        # We need aspect_size to be float rather than int since it will
        # be in a DataFrame that may be compared with others and so we
        # want to force all numeric dtypes to be float so the comparison
        # won't fail simply because e.g. int 1 isn't equal to float 1
        # More information about this:
        # Pandas DataFrames are stored internally as a series of "blocks".
        # You can see these blocks by looking at time_df._data, for instance.
        # We want to avoid any of the columns being stored within an IntBlock,
        # because it will then fail to be .equals() another DataFrame with the
        # same data but stored within a FloatBlock.
        # See http://stackoverflow.com/questions/17141828/ and
        # http://stackoverflow.com/questions/19912611/
        aspect_size_over_time = np.array(aspect_size_over_time, dtype=float)

        data_segment['aspect_size'] = aspect_size_over_time

        # Update the data segment
        time_series_data[data_segment_index] = data_segment


"""
===============================================================================
SAVING DATA
===============================================================================
"""

precision = 2  # TODO


def _data_segment_as_odict(worm_id, df_segment):
    """
    Convert a pandas dataframe into an ordered_dictionary.  This is a support
    method of data_as_array.

    """
    data_segment = [("id", worm_id)]

    # We only care about time indices for which we have some "aspect" or
    # else which have some non-aspect
    df_segment = df_segment[(~df_segment.isnull()).any(axis=1)]

    # We must make the array "jagged" according to the "aspect size",
    # since aspect size may differ frame-by-frame
    worm_aspect_size = df_segment.loc[:, ('aspect_size')]

    data_segment.append(('t', list(df_segment.index)))

    keys_used = [k for k in df_segment.columns.get_level_values('key')
                 if k != 'aspect_size']
    # There is great duplication in df_segment.columns so we must
    # filter to just the unique entries (e.g. ['x', 'y'])
    keys_used = sorted(set(keys_used))

    # e.g. ox, oy, head, ventral
    for key in [k for k in keys_used if k in elements_without_aspect]:
        cur_segment_slice = df_segment.loc[:, idx[key, 0]]
        # We must replace NaN with None, otherwise the JSON encoder
        # will save 'NaN' as the string and this will get rejected
        # by our schema on any subsequent loads
        # Note we can't use .fillna(None) due to this issue:
        # https://github.com/pydata/pandas/issues/1972
        if key in ['head', 'ventral']:
            cur_segment_slice = \
                cur_segment_slice.where(pd.notnull(cur_segment_slice),
                                        None)
        cur_list = list(np.array(cur_segment_slice))
        data_segment.append((key, cur_list))

    num_spine_points = worm_aspect_size.loc[:, 0].astype(int).tolist()

    # e.g. x, y
    for key in [k for k in keys_used if k in elements_with_aspect]:
        non_jagged_array = df_segment.loc[:, (key)]

        jagged_array = non_jagged_array.values.tolist()

        num_time_points = len(jagged_array)

        for i in range(num_time_points):
            cur_aspect_size = num_spine_points[i]

            jagged_array[i] = jagged_array[i][:cur_aspect_size]

        data_segment.append((key, jagged_array))

    return OrderedDict(data_segment)


def data_as_array(df):
    """
    Convert a pandas dataframe into an array of objects conforming
    to the WCON standard.

    Parameters
    ------------
    df: Pandas DataFrame

    Returns
    ------------
    A list of dictionaries
        Conforms to WCON standard

    """
    arr = []

    # We'd like to use df.to_dict(),
    # but you'll first have to simplify the multiindex, by taking slices
    # (across time) and then interating over those slices
    for worm_id in set(df.columns.get_level_values('id')):
        try:
            # Downconvert any numpy data types (which are not
            # JSON-serializable) into native Python data types
            worm_id = worm_id.item()
        except AttributeError:
            # If the item wasn't a numpy data type, .item() won't be available
            # just ignore this error
            pass

        df_segment = df.xs(worm_id, level='id', axis=1)

        arr.append(_data_segment_as_odict(worm_id, df_segment))

    return arr


def get_sorted_ordered_dict(v):
    """
    Recursively sort all levels of a potentially nested dict,
    list, or value.

    From http://stackoverflow.com/questions/22721579/

    Parameters
    -----------
    v: a potentially nested dict, list, or value

    Returns
    -----------
    A potentially nested OrderedDict object
        All keys at each nesting level are sorted alphabetically

    """
    if isinstance(v, dict):
        # If it's a dict, order it
        ov = OrderedDict()
        for k, k_v in sorted(v.items()):
            ov[k] = get_sorted_ordered_dict(k_v)

    elif isinstance(v, list):
        # If it's an array, retain original order, but
        # also recursively order each element
        # in case later nestings contain dicts

        ov = list(v)  # Copy
        for i in range(len(ov)):
            ov[i] = get_sorted_ordered_dict(ov[i])

    else:
        # Otherwise it's a "leaf" and no further ordering is
        # required
        ov = v

    return ov


def sort_odict(odict, sort_as_strings=True):
    """
    Returns another odict with the items sorted.

    """
    if six.PY3:
        to_sort = odict.items()
    else:
        to_sort = odict.iteritems()

    if sort_as_strings:
        return OrderedDict(sorted(to_sort, key=lambda k: str(k)))
    else:
        return OrderedDict(sorted(to_sort))
