import numpy
import SimpleITK as sitk

from radiomics import base, cShape, deprecated

class RadiomicsShape(base.RadiomicsFeaturesBase):
  r"""
  In this group of features we included descriptors of the three-dimensional size and shape of the ROI. These features
  are independent from the gray level intensity distribution in the ROI and are therefore only calculated on the
  non-derived image and mask.

  Unless otherwise specified, features are derived from the approximated shape defined by the triangle mesh. To build
  this mesh, vertices (points) are first defined as points halfway on an edge between a voxel included in the ROI and
  one outside the ROI. By connecting these vertices a mesh of connected triangles is obtained, with each triangle
  defined by 3 adjacent vertices, which shares each side with exactly one other triangle.

  This mesh is generated using a marching cubes algorithm. In this algorithm, a 2x2 cube is moved through the mask
  space. For each position, the corners of the cube are then marked 'segmented' (1) or 'not segmented' (0). Treating the
  corners as specific bits in a binary number, a unique cube-index is obtained (0-255). This index is then used to
  determine which triangles are present in the cube, which are defined in a lookup table.

  These triangles are defined in such a way, that the normal (obtained from the cross product of vectors describing 2
  out of 3 edges) are always oriented in the same direction. For PyRadiomics, the calculated normals are always pointing
  outward. This is necessary to obtain the correct signed volume used in calculation of ``MeshVolume``.

  Let:

  - :math:`N_v` represent the number of voxels included in the ROI
  - :math:`N_f` represent the number of faces (triangles) defining the Mesh.
  - :math:`V` the volume of the mesh in mm\ :sup:`3`, calculated by :py:func:`getMeshVolumeFeatureValue`
  - :math:`A` the surface area of the mesh in mm\ :sup:`2`, calculated by :py:func:`getMeshSurfaceAreaFeatureValue`

  References:

  - Lorensen WE, Cline HE. Marching cubes: A high resolution 3D surface construction algorithm. ACM SIGGRAPH Comput
    Graph `Internet <http://portal.acm.org/citation.cfm?doid=37402.37422>`_. 1987;21:163-9.
  """

  def __init__(self, inputImage, inputMask, **kwargs):
    assert inputMask.GetDimension() == 3, 'Shape features are only available in 3D. If 2D, use shape2D instead'
    super(RadiomicsShape, self).__init__(inputImage, inputMask, **kwargs)

  def _initVoxelBasedCalculation(self):
    raise NotImplementedError('Shape features are not available in voxel-based mode')

  def _initSegmentBasedCalculation(self):

    self.pixelSpacing = numpy.array(self.inputImage.GetSpacing()[::-1])

    # Pad inputMask to prevent index-out-of-range errors
    self.logger.debug('Padding the mask with 0s')

    cpif = sitk.ConstantPadImageFilter()

    padding = numpy.tile(1, 3)
    try:
      cpif.SetPadLowerBound(padding)
      cpif.SetPadUpperBound(padding)
    except TypeError:
      # newer versions of SITK/python want a tuple or list
      cpif.SetPadLowerBound(padding.tolist())
      cpif.SetPadUpperBound(padding.tolist())

    self.inputMask = cpif.Execute(self.inputMask)

    # Reassign self.maskArray using the now-padded self.inputMask
    self.maskArray = (sitk.GetArrayFromImage(self.inputMask) == self.label)
    self.labelledVoxelCoordinates = numpy.where(self.maskArray != 0)

    self.logger.debug('Pre-calculate Volume, Surface Area and Eigenvalues')

    # Volume, Surface Area and eigenvalues are pre-calculated

    # Compute Surface Area and volume
    self.SurfaceArea, self.Volume, self.diameters = cShape.calculate_coefficients(self.maskArray, self.pixelSpacing)

    # Compute eigenvalues and -vectors
    Np = len(self.labelledVoxelCoordinates[0])
    coordinates = numpy.array(self.labelledVoxelCoordinates, dtype='int').transpose((1, 0))  # Transpose equals zip(*a)
    physicalCoordinates = coordinates * self.pixelSpacing[None, :]
    physicalCoordinates -= numpy.mean(physicalCoordinates, axis=0)  # Centered at 0
    physicalCoordinates /= numpy.sqrt(Np)
    covariance = numpy.dot(physicalCoordinates.T.copy(), physicalCoordinates)
    self.eigenValues = numpy.linalg.eigvals(covariance)

    # Correct machine precision errors causing very small negative eigen values in case of some 2D segmentations
    machine_errors = numpy.bitwise_and(self.eigenValues < 0, self.eigenValues > -1e-10)
    if numpy.sum(machine_errors) > 0:
      self.logger.warning('Encountered %d eigenvalues < 0 and > -1e-10, rounding to 0', numpy.sum(machine_errors))
      self.eigenValues[machine_errors] = 0

    self.eigenValues.sort()  # Sort the eigenValues from small to large

    self.logger.debug('Shape feature class initialized')

  def getMeshVolumeFeatureValue(self):
    r"""
    **1. Mesh Volume**

    .. math::
      V_i = \displaystyle\frac{Oa_i \cdot (Ob_i \times Oc_i)}{6} \text{ (1)}

      V = \displaystyle\sum^{N_f}_{i=1}{V_i} \text{ (2)}

    The volume of the ROI :math:`V` is calculated from the triangle mesh of the ROI.
    For each face :math:`i` in the mesh, defined by points :math:`a_i, b_i` and :math:`c_i`, the (signed) volume
    :math:`V_f` of the tetrahedron defined by that face and the origin of the image (:math:`O`) is calculated. (1)
    The sign of the volume is determined by the sign of the normal, which must be consistently defined as either facing
    outward or inward of the ROI.

    Then taking the sum of all :math:`V_i`, the total volume of the ROI is obtained (2)

    .. note::
      For more extensive documentation on how the volume is obtained using the surface mesh, see the IBSI document,
      where this feature is defined as ``Volume``.
    """
    return self.Volume

  def getVoxelVolumeFeatureValue(self):
    r"""
    **2. Voxel Volume**

    .. math::
      V_{voxel} = \displaystyle\sum^{N_v}_{k=1}{V_k}

    The volume of the ROI :math:`V_{voxel}` is approximated by multiplying the number of voxels in the ROI by the volume
    of a single voxel :math:`V_k`. This is a less precise approximation of the volume and is not used in subsequent
    features. This feature does not make use of the mesh and is not used in calculation of other shape features.

    .. note::
      Defined in IBSI as ``Approximate Volume``.
    """
    z, y, x = self.pixelSpacing
    Np = len(self.labelledVoxelCoordinates[0])
    return Np * (z * x * y)

  def getSurfaceAreaFeatureValue(self):
    r"""
    **3. Surface Area**

    .. math::
      A_i = \frac{1}{2}|\text{a}_i\text{b}_i \times \text{a}_i\text{c}_i| \text{ (1)}

      A = \displaystyle\sum^{N_f}_{i=1}{A_i} \text{ (2)}

    where:

    :math:`\text{a}_i\text{b}_i` and :math:`\text{a}_i\text{c}_i` are edges of the :math:`i^{\text{th}}` triangle in the
    mesh, formed by vertices :math:`\text{a}_i`, :math:`\text{b}_i` and :math:`\text{c}_i`.

    To calculate the surface area, first the surface area :math:`A_i` of each triangle in the mesh is calculated (1).
    The total surface area is then obtained by taking the sum of all calculated sub-areas (2).

    .. note::
      Defined in IBSI as ``Surface Area``.
    """
    return self.SurfaceArea

  def getSurfaceVolumeRatioFeatureValue(self):
    r"""
    **4. Surface Area to Volume ratio**

    .. math::
      \textit{surface to volume ratio} = \frac{A}{V}

    Here, a lower value indicates a more compact (sphere-like) shape. This feature is not dimensionless, and is
    therefore (partly) dependent on the volume of the ROI.
    """
    return self.SurfaceArea / self.Volume

  def getSphericityFeatureValue(self):
    r"""
    **5. Sphericity**

    .. math::
      \textit{sphericity} = \frac{\sqrt[3]{36 \pi V^2}}{A}

    Sphericity is a measure of the roundness of the shape of the tumor region relative to a sphere. It is a
    dimensionless measure, independent of scale and orientation. The value range is :math:`0 < sphericity \leq 1`, where
    a value of 1 indicates a perfect sphere (a sphere has the smallest possible surface area for a given volume,
    compared to other solids).

    .. note::
      This feature is correlated to Compactness 1, Compactness 2 and Spherical Disproportion. In the default
      parameter file provided in the ``pyradiomics/examples/exampleSettings`` folder, Compactness 1 and Compactness 2
      are therefore disabled.
    """
    return (36 * numpy.pi * self.Volume ** 2) ** (1.0 / 3.0) / self.SurfaceArea

  @deprecated
  def getCompactness1FeatureValue(self):
    r"""
    **6. Compactness 1**

    .. math::
      \textit{compactness 1} = \frac{V}{\sqrt{\pi A^3}}

    Similar to Sphericity, Compactness 1 is a measure of how compact the shape of the tumor is relative to a sphere
    (most compact). It is therefore correlated to Sphericity and redundant. It is provided here for completeness.
    The value range is :math:`0 < compactness\ 1 \leq \frac{1}{6 \pi}`, where a value of :math:`\frac{1}{6 \pi}`
    indicates a perfect sphere.

    By definition, :math:`compactness\ 1 = \frac{1}{6 \pi}\sqrt{compactness\ 2} =
    \frac{1}{6 \pi}\sqrt{sphericity^3}`.

    .. note::
      This feature is correlated to Compactness 2, Sphericity and Spherical Disproportion.
      Therefore, this feature is marked, so it is not enabled by default (i.e. this feature will not be enabled if no
      individual features are specified (enabling 'all' features), but will be enabled when individual features are
      specified, including this feature). To include this feature in the extraction, specify it by name in the enabled
      features.
    """
    return self.Volume / (self.SurfaceArea ** (3.0 / 2.0) * numpy.sqrt(numpy.pi))

  @deprecated
  def getCompactness2FeatureValue(self):
    r"""
    **7. Compactness 2**

    .. math::
      \textit{compactness 2} = 36 \pi \frac{V^2}{A^3}

    Similar to Sphericity and Compactness 1, Compactness 2 is a measure of how compact the shape of the tumor is
    relative to a sphere (most compact). It is a dimensionless measure, independent of scale and orientation. The value
    range is :math:`0 < compactness\ 2 \leq 1`, where a value of 1 indicates a perfect sphere.

    By definition, :math:`compactness\ 2 = (sphericity)^3`

    .. note::
      This feature is correlated to Compactness 1, Sphericity and Spherical Disproportion.
      Therefore, this feature is marked, so it is not enabled by default (i.e. this feature will not be enabled if no
      individual features are specified (enabling 'all' features), but will be enabled when individual features are
      specified, including this feature). To include this feature in the extraction, specify it by name in the enabled
      features.
    """
    return (36.0 * numpy.pi) * (self.Volume ** 2.0) / (self.SurfaceArea ** 3.0)

  @deprecated
  def getSphericalDisproportionFeatureValue(self):
    r"""
    **8. Spherical Disproportion**

    .. math::
      \textit{spherical disproportion} = \frac{A}{4\pi R^2} = \frac{A}{\sqrt[3]{36 \pi V^2}}

    Where :math:`R` is the radius of a sphere with the same volume as the tumor, and equal to
    :math:`\sqrt[3]{\frac{3V}{4\pi}}`.

    Spherical Disproportion is the ratio of the surface area of the tumor region to the surface area of a sphere with
    the same volume as the tumor region, and by definition, the inverse of Sphericity. Therefore, the value range is
    :math:`spherical\ disproportion \geq 1`, with a value of 1 indicating a perfect sphere.

    .. note::
      This feature is correlated to Compactness 2, Compactness2 and Sphericity.
      Therefore, this feature is marked, so it is not enabled by default (i.e. this feature will not be enabled if no
      individual features are specified (enabling 'all' features), but will be enabled when individual features are
      specified, including this feature). To include this feature in the extraction, specify it by name in the enabled
      features.
    """
    return self.SurfaceArea / (36 * numpy.pi * self.Volume ** 2) ** (1.0 / 3.0)

  def getMaximum3DDiameterFeatureValue(self):
    r"""
    **9. Maximum 3D diameter**

    Maximum 3D diameter is defined as the largest pairwise Euclidean distance between tumor surface mesh
    vertices.

    Also known as Feret Diameter.
    """
    return self.diameters[3]

  def getMaximum2DDiameterSliceFeatureValue(self):
    r"""
    **10. Maximum 2D diameter (Slice)**

    Maximum 2D diameter (Slice) is defined as the largest pairwise Euclidean distance between tumor surface mesh
    vertices in the row-column (generally the axial) plane.
    """
    return self.diameters[0]

  def getMaximum2DDiameterColumnFeatureValue(self):
    r"""
    **11. Maximum 2D diameter (Column)**

    Maximum 2D diameter (Column) is defined as the largest pairwise Euclidean distance between tumor surface mesh
    vertices in the row-slice (usually the coronal) plane.
    """
    return self.diameters[1]

  def getMaximum2DDiameterRowFeatureValue(self):
    r"""
    **12. Maximum 2D diameter (Row)**

    Maximum 2D diameter (Row) is defined as the largest pairwise Euclidean distance between tumor surface mesh
    vertices in the column-slice (usually the sagittal) plane.
    """
    return self.diameters[2]

  def getMajorAxisLengthFeatureValue(self):
    r"""
    **13. Major Axis Length**

    .. math::
      \textit{major axis} = 4 \sqrt{\lambda_{major}}

    This feature yield the largest axis length of the ROI-enclosing ellipsoid and is calculated using the largest
    principal component :math:`\lambda_{major}`.

    The principal component analysis is performed using the physical coordinates of the voxel centers defining the ROI.
    It therefore takes spacing into account, but does not make use of the shape mesh.
    """
    if self.eigenValues[2] < 0:
      self.logger.warning('Major axis eigenvalue negative! (%g)', self.eigenValues[2])
      return numpy.nan
    return numpy.sqrt(self.eigenValues[2]) * 4

  def getMinorAxisLengthFeatureValue(self):
    r"""
    **14. Minor Axis Length**

    .. math::
      \textit{minor axis} = 4 \sqrt{\lambda_{minor}}

    This feature yield the second-largest axis length of the ROI-enclosing ellipsoid and is calculated using the largest
    principal component :math:`\lambda_{minor}`.

    The principal component analysis is performed using the physical coordinates of the voxel centers defining the ROI.
    It therefore takes spacing into account, but does not make use of the shape mesh.
    """
    if self.eigenValues[1] < 0:
      self.logger.warning('Minor axis eigenvalue negative! (%g)', self.eigenValues[1])
      return numpy.nan
    return numpy.sqrt(self.eigenValues[1]) * 4

  def getLeastAxisLengthFeatureValue(self):
    r"""
    **15. Least Axis Length**

    .. math::
      \textit{least axis} = 4 \sqrt{\lambda_{least}}

    This feature yield the smallest axis length of the ROI-enclosing ellipsoid and is calculated using the largest
    principal component :math:`\lambda_{least}`. In case of a 2D segmentation, this value will be 0.

    The principal component analysis is performed using the physical coordinates of the voxel centers defining the ROI.
    It therefore takes spacing into account, but does not make use of the shape mesh.
    """
    if self.eigenValues[0] < 0:
      self.logger.warning('Least axis eigenvalue negative! (%g)', self.eigenValues[0])
      return numpy.nan
    return numpy.sqrt(self.eigenValues[0]) * 4

  def getElongationFeatureValue(self):
    r"""
    **16. Elongation**

    Elongation shows the relationship between the two largest principal components in the ROI shape.
    For computational reasons, this feature is defined as the inverse of true elongation.

    .. math::
      \textit{elongation} = \sqrt{\frac{\lambda_{minor}}{\lambda_{major}}}

    Here, :math:`\lambda_{\text{major}}` and :math:`\lambda_{\text{minor}}` are the lengths of the largest and second
    largest principal component axes. The values range between 1 (where the cross section through the first and second
    largest principal moments is circle-like (non-elongated)) and 0 (where the object is a maximally elongated: i.e. a 1
    dimensional line).

    The principal component analysis is performed using the physical coordinates of the voxel centers defining the ROI.
    It therefore takes spacing into account, but does not make use of the shape mesh.
    """
    if self.eigenValues[1] < 0 or self.eigenValues[2] < 0:
      self.logger.warning('Elongation eigenvalue negative! (%g, %g)', self.eigenValues[1], self.eigenValues[2])
      return numpy.nan
    return numpy.sqrt(self.eigenValues[1] / self.eigenValues[2])

  def getFlatnessFeatureValue(self):
    r"""
    **17. Flatness**

    Flatness shows the relationship between the largest and smallest principal components in the ROI shape.
    For computational reasons, this feature is defined as the inverse of true flatness.

    .. math::
      \textit{flatness} = \sqrt{\frac{\lambda_{least}}{\lambda_{major}}}

    Here, :math:`\lambda_{\text{major}}` and :math:`\lambda_{\text{least}}` are the lengths of the largest and smallest
    principal component axes. The values range between 1 (non-flat, sphere-like) and 0 (a flat object, or single-slice
    segmentation).

    The principal component analysis is performed using the physical coordinates of the voxel centers defining the ROI.
    It therefore takes spacing into account, but does not make use of the shape mesh.
    """
    if self.eigenValues[0] < 0 or self.eigenValues[2] < 0:
      self.logger.warning('Elongation eigenvalue negative! (%g, %g)', self.eigenValues[0], self.eigenValues[2])
      return numpy.nan
    return numpy.sqrt(self.eigenValues[0] / self.eigenValues[2])


import json, os, cv2
from pathlib import Path
import numpy as np
import pandas as pd
from tqdm import tqdm
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from sklearn.cluster import KMeans
import matplotlib.colors as mcolors

from src.base import BaseInput
from src.topcam import Topcam
from src.utils.auxiliary import flatten_series, find_index_in_list
from src.utils.path import find

class AvoidanceTrial(BaseInput):
    def __init__(self, s_input, path_input, metadata_input):
        self.s = s_input # series from dataframe of all trials
        self.likelihood_thresh = 0.99
        self.dist_across_arena = 30.48 # cm between bottom-right and bottom-left pillar
        self.path = path_input
        self.camname = 'top1'
        self.shared_metadata = metadata_input

        self.num_clusters_to_use = self.shared_metadata[self.s['date']][self.s['animal']][str(self.s['task'])]['num_positions']
        self.trial_path = os.path.join(*[self.path, str(self.s['date']), self.s['animal'],str(self.s['task'])])

        self.generic_camconfig = {
            'internals': {
                'follow_strict_naming': False,
                'likelihood_threshold': 0.99
            }
        }

    def convert_pxls_to_dist(self):
        x_cols = [i for i in self.data.columns.values if '_x' in i]
        y_cols = [i for i in self.data.columns.values if '_y' in i]
        for i in range(len(x_cols)):
            self.data[x_cols[i]+'_cm'] = self.data.loc[:,x_cols[i]] / self.pxls2cm
            self.data[y_cols[i]+'_cm'] = self.data.loc[:,y_cols[i]] / self.pxls2cm

    #def setup_trial(self, c, odd):

    def make_task_df(self):
        num_odd_trials = np.min([len(self.s['poke1_ts']), len(self.s['poke2_ts'])])
        df1 = pd.DataFrame([])
        count = -1
        for c in range(num_odd_trials):
            # odd
            count += 1
            df1.at[count, 'first_poke'] = self.s['poke1_ts'][c]
            df1.at[count, 'second_poke'] = self.s['poke2_ts'][c]
            time = self.s['top1_ts']; time = time[time > df1.loc[count,'first_poke']]; time = time[time < df1.loc[count,'second_poke']]
            df1.at[count, 'trial_timestamps'] = time.astype(object)
            start_stop_inds = (int(np.where([self.s['top1_ts']==time[0]])[1]), int(np.where([self.s['top1_ts']==time[-1]])[1]))
            for pos in list(self.positions['point_loc'].values):
                df1.at[count, pos] = np.array(self.positions.loc[start_stop_inds[0]:start_stop_inds[1], pos]).astype(object)
            df1.at[count, 'len'] = start_stop_inds[1] - start_stop_inds[0]
            # even
            count += 1
            if c+1 < len(self.s['poke1_ts']):
                df1.at[count, 'first_poke'] = self.s['poke2_ts'][c]
                df1.at[count, 'second_poke'] = self.s['poke1_ts'][c+1]
                time = self.s['top1_ts']; time = time[time > df1.loc[count,'first_poke']]; time = time[time < df1.loc[count,'second_poke']]
                vidframes = np.array(list(find_index_in_list(list(self.s['top1_ts']), list(time))))
                df1.at[count, 'recording_timestamps'] = self.s['top1_ts'].astype(object)
                df1.at[count, 'trial_timestamps'] = time.astype(object)
                df1.at[count, 'trial_vidframes'] = vidframes.astype(object)
                start_stop_inds = (int(np.where([self.s['top1_ts']==time[0]])[1]), int(np.where([self.s['top1_ts']==time[-1]])[1]))
                for pos in list(self.positions['point_loc'].values):
                    df1.at[count, pos] = np.array(self.positions.loc[start_stop_inds[0]:start_stop_inds[1], pos]).astype(object)
                df1.at[count, 'len'] = start_stop_inds[1] - start_stop_inds[0]
        df1['animal'] = self.s['animal']; df1['date'] = self.s['date']; df1['task'] = self.s['task']
        self.data = df1

    #def drop_lickport_times(self, ):

    def get_median_trace(self, df):
        fake_time = np.linspace(0,1,100)
        all_nose_positions = np.zeros([len(df), 2, 100])
        count = 0
        for ind, row in df.iterrows():
            xT = np.linspace(0,1,len(row['nose_x'])); yT = np.linspace(0,1,len(row['nose_y']))
            all_nose_positions[count,0,:] = interp1d(xT, row['nose_x_cm'], bounds_error=False)(fake_time)
            all_nose_positions[count,1,:] = interp1d(yT, row['nose_y_cm'], bounds_error=False)(fake_time)
            count += 1
        median_trace = np.nanmedian(all_nose_positions, axis=0)
        for ind, row in df.iterrows():
            df.at[ind,'median_x_cm'] = median_trace[0,:].astype(object); df.at[ind,'median_y_cm'] = median_trace[1,:].astype(object)
        return df

    def distance_from_nose(self, row, target, xy_together=False):
        if not xy_together:
            x_dist = np.abs(row[target+'_x_cm'] - row['nose_x_cm'])
            y_dist = np.abs(row[target+'_y_cm'] - row['nose_y_cm'])
        elif xy_together:
            x_dist = np.abs(row[target][0] - row['nose_x_cm'])
            y_dist = np.abs(row[target][1] - row['nose_y_cm'])
        length = len(x_dist)
        dist = np.zeros([length])
        for i in range(length):
            dist[i] = np.sqrt(x_dist[i]**2 + y_dist[i]**2)
        return dist

    def angle_from_nose(self, row, target, xy_together=False):
        if not xy_together:
            x_dist = np.abs(row[target+'_x_cm'] - row['nose_x_cm'])
            y_dist = np.abs(row[target+'_y_cm'] - row['nose_y_cm'])
        elif xy_together:
            x_dist = np.abs(row[target][0] - row['nose_x_cm'])
            y_dist = np.abs(row[target][1] - row['nose_y_cm'])
        length = len(x_dist)
        ang = np.zeros([length])
        for i in range(length):
            try:
                ang[i] = (y_dist[i]/x_dist[i])
            except ZeroDivisionError:
                ang[i] = np.nan
        return ang%np.pi

    def get_head_angle(self):
        for ind, row in self.data.iterrows():
            leftear_x = row['leftear_x_cm']
            leftear_y = row['leftear_y_cm']
            rightear_x = row['rightear_x_cm']
            rightear_y = row['rightear_y_cm']
            nose_x = row['nose_x_cm']
            nose_y = row['nose_y_cm']
            angs = []
            for step in range(len(leftear_x)):
                ang = np.arctan2(np.mean([leftear_y[step],rightear_y[step]])-nose_y[step],np.mean([leftear_x[step],rightear_x[step]])-nose_x[step])
                angs.append(ang)
            self.data.at[ind, 'head_angle'] = np.array(angs).astype(object)

    def approaching_target(self, dist):
        if np.nanmean(dist[:10]) > np.nanmean(dist[:-10]):
            return True
        else:
            return False

    def add_tracking(self, name, path):
        tc = Topcam(self.generic_camconfig, name, path, self.camname)
        tc.gather_files()
        tc.pack_position_data()
        tc.filter_likelihood()
        self.positions = tc.xrpts
        self.trial_name = name
        self.trial_path = path

    def pillar_avoidance(self):
        self.make_task_df()

        # label odd/even trials (i.e. moving leftwards or moving rightwards?)
        self.data['odd'] = np.nan
        for i, ind in enumerate(self.data.index.values):
            if ind%2 == 0: # odd values
                self.data.at[ind, 'odd'] = True
            elif ind%2 == 1:
                self.data.at[ind, 'odd'] = False

        dist_to_posts = np.median(self.data['arenaTR_x'].iloc[0],0) - np.median(self.data['arenaTL_x'].iloc[0],0)
        self.pxls2cm = dist_to_posts/self.dist_across_arena
        self.convert_pxls_to_dist()

        self.get_head_angle()

        pdf = PdfPages(os.path.join(self.trial_path, (self.data['animal'].iloc[0]+'_'+str(self.data['date'].iloc[0])+'_'+str(self.data['task'].iloc[0])+'_figs.pdf')))

        # get obstacle position for all times
        for ind, row in self.data.iterrows():
            for x in ['b','w']:
                xvals = np.stack([row['obstacle'+x+'TL_x_cm'], row['obstacle'+x+'TR_x_cm'], row['obstacle'+x+'BL_x_cm'], row['obstacle'+x+'BR_x_cm']]).astype(float)
                xvals_cm = np.stack([row['obstacle'+x+'TL_x_cm'], row['obstacle'+x+'TR_x_cm'], row['obstacle'+x+'BL_x_cm'], row['obstacle'+x+'BR_x_cm']]).astype(float)
                self.data.at[ind, x+'obstacle_x'] = np.nanmean(xvals)
                self.data.at[ind, x+'obstacle_x_cm'] = np.nanmean(xvals_cm)
                self.data.at[ind, x+'obstacle_x_std'] = np.mean(np.nanstd(xvals, axis=1))
                yvals = np.stack([row['obstaclewTL_y_cm'], row['obstaclewTR_y_cm'], row['obstaclewBL_y_cm'], row['obstaclewBR_y_cm']]).astype(float)
                yvals_cm = np.stack([row['obstaclewTL_y_cm'], row['obstaclewTR_y_cm'], row['obstaclewBL_y_cm'], row['obstaclewBR_y_cm']]).astype(float)
                self.data.at[ind, x+'obstacle_y'] = np.nanmean(yvals)
                self.data.at[ind, x+'obstacle_y_cm'] = np.nanmean(yvals_cm)
                self.data.at[ind, x+'obstacle_y_std'] = np.mean(np.nanstd(yvals, axis=1))

        # animal speed
        for ind, row in self.data.iterrows():
            temp_time = np.diff(row['trial_timestamps'])
            x = np.diff(row['nose_x_cm']); y = np.diff(row['nose_y_cm'])
            if len(x) == len(temp_time):
                xspeed = list((x/temp_time)**2)
            elif len(x) > len(temp_time):
                xspeed = list((x[:len(temp_time)]/temp_time)**2)
            elif len(x) < len(temp_time):
                xspeed = list((x/temp_time[:len(x)])**2)
            if len(y) == len(temp_time):
                yspeed = list((y/temp_time)**2)
            elif len(y) > len(temp_time):
                yspeed = list((y[:len(temp_time)]/temp_time)**2)
            elif len(y) < len(temp_time):
                yspeed = list((y/temp_time[:len(y)])**2)
            self.data.at[ind, 'speed'] = np.sqrt(xspeed + yspeed).astype(object)

        # cluster obstacle positions
        kmeans_input = np.stack([self.data['wobstacle_x_cm'].map(np.nanmean), self.data['wobstacle_y_cm'].map(np.nanmean)])
        kmeans_mask = np.any(np.isnan(kmeans_input), axis=0)
        kmeans_input = kmeans_input[:,~kmeans_mask]
        labels = KMeans(n_clusters=self.num_clusters_to_use, random_state=0).fit(kmeans_input.T).labels_

        # transit time histogram
        plt.figure()
        plt.xlabel('transit time')
        plt.hist(self.data['len'], bins=25)
        pdf.savefig(); plt.close()

        # spatial positions of obstacle clusters
        plt.figure()
        for i in range(len(labels)):
            obstacles = kmeans_input[:,i]
            plt.plot(obstacles[0], obstacles[1], '*', color=list(mcolors.TABLEAU_COLORS)[labels[i]])
            plt.ylim([20.03, 0]); plt.xlim([0, 33.30])
        pdf.savefig(); plt.close()

        # add kmeans labels to trial data
        self.data = self.data[~pd.isnull(self.data['wobstacle_x_cm'])][~pd.isnull(self.data['wobstacle_y_cm'])]
        self.data['obstacle_cluster'] = labels

        # drop any transits that were really slow (only drop slowest 10% of transits)
        time_thresh = self.data['len'].quantile(0.9)
        self.data = self.data[self.data['len']<time_thresh]

        # plot all the traces
        odd = self.data[self.data.index%2==0]
        even = self.data[self.data.index%2==1]
        direction_count = 0
        for direction_df in [odd, even]:
            if direction_count == 0:
                leftcolor='g'; rightcolor='b'
            else:
                leftcolor='b'; rightcolor='g'
            plt.subplots(3,3, figsize=(9,6))
            for c in range(9):
                this_cluster = direction_df[direction_df['obstacle_cluster']==c].copy().reset_index()
                plt.subplot(3,3,c+1)
                plt.gca().set_aspect('equal', adjustable='box')
                colors = plt.cm.magma(np.linspace(0,1,len(this_cluster)))
                for ind, row in this_cluster.iterrows():
                    plt.plot([np.median(row['obstaclewTL_x_cm'],0),
                            np.median(row['obstaclewTR_x_cm'],0),
                            np.median(row['obstaclewBR_x_cm'],0),
                            np.median(row['obstaclewBL_x_cm'],0),
                            np.median(row['obstaclewTL_x_cm'],0)],
                            [np.median(row['obstaclewTL_y_cm'],0),
                            np.median(row['obstaclewTR_y_cm'],0),
                            np.median(row['obstaclewBR_y_cm'],0),
                            np.median(row['obstaclewBL_y_cm'],0),
                            np.median(row['obstaclewTL_y_cm'],0)],'k-')
                    plt.plot([np.median(row['arenaTL_x_cm'],0),
                            np.median(row['arenaTR_x_cm'],0),
                            np.median(row['arenaBR_x_cm'],0),
                            np.median(row['arenaBL_x_cm'],0),
                            np.median(row['arenaTL_x_cm'],0)],
                            [np.median(row['arenaTL_y_cm'],0),
                            np.median(row['arenaTR_y_cm'],0),
                            np.median(row['arenaBR_y_cm'],0),
                            np.median(row['arenaBL_y_cm'],0),
                            np.median(row['arenaTL_y_cm'],0)],'k-')
                    plt.plot(row['nose_x_cm'], row['nose_y_cm'], '-', color=colors[ind])
                    plt.plot(row['leftportT_x_cm'], row['leftportT_y_cm'],'.',color=leftcolor)
                    plt.plot(row['rightportT_x_cm'], row['rightportT_y_cm'],'.',color=rightcolor)
                plt.ylim([20.03,0]); plt.xlim([0,33.30])
                if len(this_cluster) > 0:
                    plt.plot(self.get_median_trace(this_cluster)['median_x_cm'].iloc[0], self.get_median_trace(this_cluster)['median_y_cm'].iloc[0], 'b-')
            direction_count += 1
            plt.tight_layout()
            pdf.savefig(); plt.close()

        # overlay of cluster medians
        direction_count = 0
        plt.subplots(1,2, figsize=(8,5))
        for direction_df in [odd, even]:
            if direction_count == 0:
                leftcolor='g'; rightcolor='b'
            else:
                leftcolor='b'; rightcolor='g'
            plt.subplot(1,2,direction_count+1)
            for c in range(9):
                cluster_color = list(mcolors.TABLEAU_COLORS)[c]
                this_cluster = direction_df[direction_df['obstacle_cluster']==c].copy().reset_index()
                for ind, row in this_cluster.iterrows():
                    plt.plot([np.median(row['obstaclewTL_x_cm'],0),
                            np.median(row['obstaclewTR_x_cm'],0),
                            np.median(row['obstaclewBR_x_cm'],0),
                            np.median(row['obstaclewBL_x_cm'],0),
                            np.median(row['obstaclewTL_x_cm'],0)],
                            [np.median(row['obstaclewTL_y_cm'],0),
                            np.median(row['obstaclewTR_y_cm'],0),
                            np.median(row['obstaclewBR_y_cm'],0),
                            np.median(row['obstaclewBL_y_cm'],0),
                            np.median(row['obstaclewTL_y_cm'],0)],'-', color=cluster_color)
                    plt.plot(row['leftportT_x_cm'], row['leftportT_y_cm'], '.', color=leftcolor)
                    plt.plot(row['rightportT_x_cm'], row['rightportT_y_cm'], '.', color=rightcolor)
                plt.ylim([20.03,0]); plt.xlim([0,33.30])
                if len(this_cluster) > 0:
                    plt.plot(self.get_median_trace(this_cluster)['median_x_cm'].iloc[0], self.get_median_trace(this_cluster)['median_y_cm'].iloc[0], '-', color=cluster_color)
            direction_count += 1
            plt.tight_layout()
            pdf.savefig(); plt.close()

        # properties over time (e.g. time active, speed, etc.)
        plt.subplots(2,3, figsize=(15,8))
        median_speed = []; max_speed = []; time_active = []
        for ind, row in self.data.iterrows():
            median_speed.append(np.median(row['speed']))
            max_speed.append(np.max(row['speed']))
            time_active.append(np.sum(row['speed']>5)/60)
        plt.subplot(2,3,1)
        slow = self.data['speed'].iloc[np.nanargmin(median_speed)]
        slowT = np.linspace(0,1,len(slow))
        med = self.data['speed'].iloc[np.argsort(median_speed)[len(median_speed)//2]]
        medT = np.linspace(0,1,len(med))
        fast = self.data['speed'].iloc[np.nanargmax(median_speed)]
        fastT = np.linspace(0,1,len(fast))
        fake_time = np.linspace(0,1,100)
        plt.plot(interp1d(slowT, slow, bounds_error=False)(fake_time))
        plt.plot(interp1d(medT, med, bounds_error=False)(fake_time))
        plt.plot(interp1d(fastT, fast, bounds_error=False)(fake_time))
        plt.subplot(2,3,2)
        plt.hist(median_speed); plt.xlabel('median speed (cm/sec)')
        plt.subplot(2,3,3)
        plt.hist(max_speed); plt.xlabel('max speed (cm/sec)')
        plt.subplot(2,3,4)
        plt.hist(time_active); plt.xlabel('time active (sec)')
        plt.subplot(2,3,5)
        plt.plot([i for i in max_speed if ~np.isnan(i)]); plt.xlabel('trial'); plt.ylabel('max speed (cm/sec)')
        plt.subplot(2,3,6)
        plt.plot([i for i in time_active if ~np.isnan(i)]); plt.xlabel('trial'); plt.ylabel('time active (sec)')
        plt.tight_layout()
        pdf.savefig()
        plt.close()

        plt.subplots(4,4,figsize=(20,18))
        fake_time = np.linspace(0,1,100)
        plt.subplot(4,4,1)
        for ind, data in self.data['speed'].iteritems():
            plt.plot(interp1d(np.linspace(0,1,len(data)), data, bounds_error=False)(fake_time), '-', alpha=0.2); plt.ylabel('speed'); plt.xlabel('time')
        plt.subplot(4,4,2)
        for ind, data in self.data['head_angle'].iteritems():
            plt.plot(interp1d(np.linspace(0,1,len(data)), data, bounds_error=False)(fake_time), '-', alpha=0.2); plt.ylabel('angle to horizontal'); plt.xlabel('time')
        odd = self.data[self.data.index%2==0]
        even = self.data[self.data.index%2==1]
        for direction_num in range(2):
            plt.subplot(4,4,3+direction_num)
            direction_df = [odd, even][direction_num]
            for ind, row in direction_df.iterrows():
                dist = self.distance_from_nose(row, 'leftportT')
                plt.plot(interp1d(np.linspace(0,1,len(dist)), dist, bounds_error=False)(fake_time), alpha=0.2)
                if self.approaching_target(dist) is True:
                    plt.title('distance to target port (moving left)')
                    direction_df.at[ind, 'approach'] = True
                else:
                    plt.title('distance to previous port (moving left)')
                    direction_df.at[ind, 'approach'] = False
                plt.plot(0,dist[0],'.',color='g')
                plt.plot(100,dist[-1],'.',color='b')
        for direction_num in range(2):
            plt.subplot(4,4,5+direction_num)
            direction_df = [odd, even][direction_num]
            for ind, row in direction_df.iterrows():
                dist = self.distance_from_nose(row, 'rightportT')
                plt.plot(interp1d(np.linspace(0,1,len(dist)), dist, bounds_error=False)(fake_time), alpha=0.2)
                if row['approach']:
                    plt.title('distance to previous port (moving right)')
                else:
                    plt.title('distance to target port (moving right)')
                plt.plot(0,dist[0],'.',color='g')
                plt.plot(100,dist[-1],'.',color='b')
        for direction_num in range(2):
            plt.subplot(4,4,7+direction_num)
            direction_df = [odd, even][direction_num]
            for ind, row in direction_df.iterrows():
                ang = self.angle_from_nose(row, 'leftportT')
                plt.plot(interp1d(np.linspace(0,1,len(ang)), ang, bounds_error=False)(fake_time), alpha=0.2)
                if row['approach']:
                    plt.title('angle to target port (headed left')
                else:
                    plt.title('angle to previous port (headed left)')
                plt.plot(0,ang[0],'.',color='g')
                plt.plot(100,ang[-1],'.',color='b')
        for direction_num in range(2):
            plt.subplot(4,4,9+direction_num)
            direction_df = [odd, even][direction_num]
            for ind, row in direction_df.iterrows():
                ang = self.angle_from_nose(row, 'rightportT')
                plt.plot(interp1d(np.linspace(0,1,len(ang)), ang, bounds_error=False)(fake_time), alpha=0.2)
                if row['approach']:
                    plt.title('angle to target port')
                else:
                    plt.title('angle to previous port')
                plt.plot(0,ang[0],'.',color='g')
                plt.plot(100,ang[-1],'.',color='b')
        for direction_num in range(2):
            plt.subplot(4,4,11+direction_num)
            direction_df = [odd, even][direction_num]
            for ind, row in direction_df.iterrows():
                dist = self.distance_from_nose(row, 'wobstacle')
                plt.plot(interp1d(np.linspace(0,1,len(dist)), dist, bounds_error=False)(fake_time), alpha=0.2)
                plt.title('distance to obstacle')
        for direction_num in range(2):
            plt.subplot(4,4,13+direction_num)
            direction_df = [odd, even][direction_num]
            for ind, row in direction_df.iterrows():
                ang = self.angle_from_nose(row, 'wobstacle')
                plt.plot(interp1d(np.linspace(0,1,len(ang)), ang, bounds_error=False)(fake_time), alpha=0.2)
                plt.title('angle to obstacle')
        plt.tight_layout()
        pdf.savefig()
        plt.close()

        pdf.close()
        self.data.to_hdf(os.path.join(self.trial_path, (self.data['animal'].iloc[0]+'_'+str(self.data['date'].iloc[0])+'_'+str(self.data['task'].iloc[0])+'.h5')), 'w')

    def format_frames_oa(self, vid_path):
        # open the .avi file
        vidread = cv2.VideoCapture(vid_path)
        # empty array that is the target shape
        # should be number of frames x downsampled height x downsampled width
        all_frames = np.empty([int(vidread.get(cv2.CAP_PROP_FRAME_COUNT)),
                            int(vidread.get(cv2.CAP_PROP_FRAME_HEIGHT)),
                            int(vidread.get(cv2.CAP_PROP_FRAME_WIDTH))], dtype=np.uint8)
        # iterate through each frame
        for frame_num in tqdm(range(0,int(vidread.get(cv2.CAP_PROP_FRAME_COUNT)))):
            # read the frame in and make sure it is read in correctly
            ret, frame = vidread.read()
            if not ret:
                break
            # convert to grayyscale
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            # add the downsampled frame to all_frames as int8
            all_frames[frame_num,:,:] = frame.astype(np.int8)
        return all_frames

    def get_row_for_timestamp(self, df, seek_timestamp):
        for ind, row in df.iterrows():
            if seek_timestamp in row['trial_timestamps']:
                return row

    def plot_frame(self, vid_arr, timestamps, df, seek_frame, return_as_array=False):
        seek_timestamp = timestamps[seek_frame]
        row = self.get_row_for_timestamp(df, seek_timestamp)
        if row is None:
            if return_as_array:
                return np.zeros(np.shape(vid_arr[0]))
            elif not return_as_array:
                plt.figure()
                plt.imshow(np.zeros(np.shape(vid_arr[0])), cmap='gray')
                plt.show()
        row_time_index = np.where(row['trial_timestamps']==seek_timestamp)
        current_ang = row['head_angle'][row_time_index][0]
        x1 = row['nose_x'][row_time_index]
        y1 = row['nose_y'][row_time_index]
        x2 = x1+60 * np.cos(current_ang)
        y2 = y1+60 * np.sin(current_ang)
        frame = vid_arr[seek_frame,:,:]
        fig = plt.figure()
        plt.imshow(frame, cmap='gray')
        plt.plot((x1,x2), (y1,y2), '-')
        row_time_index = row_time_index[0][0]
        plt.plot(row['nose_x'][:row_time_index], row['nose_y'][:row_time_index],'r.')
        plt.plot(row['leftear_x'][:row_time_index], row['leftear_y'][:row_time_index], 'g.')
        plt.plot(row['rightear_x'][:row_time_index], row['rightear_y'][:row_time_index], 'g.')
        if not return_as_array:
            plt.show()
        elif return_as_array:
            fig.canvas.draw()
            frame_as_array = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
            frame_as_array = frame_as_array.reshape(fig.canvas.get_width_height()[::-1] + (3,))
            plt.close()
            return frame_as_array

    def plot_all_trials(self, vid_arr, timestamps, df, vid_savepath):
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        out_vid = cv2.VideoWriter(vid_savepath, fourcc, 60.0, (640, 480))
        start = 3600*5
        for seek_frame in tqdm(range(start,start+3600)):
            frame = self.plot_frame(vid_arr, timestamps, df, seek_frame, return_as_array=True)
            out_vid.write(frame.astype('uint8'))
        out_vid.release()

    def make_videos(self):
        vid_savepath = os.path.join(self.trial_path, (self.data['animal'].iloc[0]+'_'+str(self.data['date'].iloc[0])+'_'+str(self.data['task'].iloc[0])+'plot.avi'))
        vid_path = find('*'+str(self.s['date'])+'*'+self.s['animal']+'*'+str(self.s['task'])+'*.avi', self.trial_path)[0]
        self.timestamp_path = find('*'+str(self.s['date'])+'*'+self.s['animal']+'*'+str(self.s['task'])+'*_top1_BonsaiTS.csv', self.trial_path)[0]
        print('formating video frames as array')
        vid_arr = self.format_frames_oa(vid_path)
        print('plotting video of traces')
        self.plot_all_trials(vid_arr, self.read_timestamp_file(), self.data, vid_savepath)

    def gap_detection(self):

        self.make_task_df()

        dist_to_posts = np.median(self.data['arenaTR_x'].iloc[0],0) - np.median(self.data['arenaTL_x'].iloc[0],0)
        self.pxls2cm = dist_to_posts/self.dist_across_arena
        self.convert_pxls_to_dist()

        self.get_head_angle()

        pdf = PdfPages(os.path.join(self.trial_path, (self.data['animal'].iloc[0]+'_'+str(self.data['date'].iloc[0])+'_'+str(self.data['task'].iloc[0])+'_figs.pdf')))

        self.timestamp_path = find('*'+str(self.s['date'])+'*'+self.s['animal']+'*'+str(self.s['task'])+'*_top1_BonsaiTS.csv', self.trial_path)[0]

        # get gap positions
        sq_names = ['sBR','sBL','sTR','sTL','s2BR','s2BL','s2TR','s2TL']
        tr_names = ['tBL','tTL','tM','t2BL','t2TL','t2M']
        for ind, row in self.data.iterrows():
            # first gap
            sq_x_positions = np.zeros(8)
            sq_y_positions = np.zeros(8)
            sq_std = np.zeros(8)
            for i, sq_name in enumerate(sq_names):
                x_positions = row['obstacle'+sq_name+'_x_cm'].astype(float)
                y_positions = row['obstacle'+sq_name+'_y_cm'].astype(float)
                if ~np.isnan(x_positions).all() > 0 and ~np.isnan(y_positions).all() > 0:
                    sq_x_positions[i] = np.nanmedian(x_positions)
                    sq_y_positions[i] = np.nanmedian(y_positions)
                    sq_std[i] = np.mean([np.nanstd(x_positions),np.nanstd(y_positions)])
            self.data.at[ind, 'obs_sq_x_pos'] = sq_x_positions.astype(object)
            self.data.at[ind, 'obs_sq_y_pos'] = sq_y_positions.astype(object)
            self.data.at[ind, 'obs_sq_std'] = sq_std.astype(object)

            tr_x_positions = np.zeros(6)
            tr_y_positions = np.zeros(6)
            tr_std = np.zeros(6)
            for i, tr_name in enumerate(tr_names):
                x_positions = row['obstacle'+tr_name+'_x_cm'].astype(float)
                y_positions = row['obstacle'+tr_name+'_y_cm'].astype(float)
                if ~np.isnan(x_positions).all() > 0 and ~np.isnan(y_positions).all() > 0:
                    tr_x_positions[i] = np.nanmedian(x_positions)
                    tr_y_positions[i] = np.nanmedian(y_positions)
                    tr_std[i] = np.mean([np.nanstd(x_positions),np.nanstd(y_positions)])
            self.data.at[ind, 'obs_tr_x_pos'] = tr_x_positions.astype(object)
            self.data.at[ind, 'obs_tr_y_pos'] = tr_y_positions.astype(object)
            self.data.at[ind, 'obs_tr_std'] = tr_std.astype(object)
            self.data.at[ind, 'might_be_bad'] = (True if (np.sum(~np.isnan(tr_x_positions[:4]))<3 or np.sum(~np.isnan(sq_x_positions[:5]))<4) else False)
            has_both_doors = (True if (np.sum(~np.isnan(tr_x_positions))>3 and np.sum(~np.isnan(sq_x_positions))>4) else False)
            door1 = np.array([np.nanmean(sq_x_positions[2:4]), np.nanmean(sq_y_positions[2:4]), np.nanmean(tr_x_positions[:2]), np.nanmean(tr_y_positions[:2])])
            self.data.at[ind, 'door1'] = door1.astype(object)
            self.data.at[ind, 'door1_cent'] = np.array([np.nanmean([door1[0], door1[2]]), np.nanmean([door1[1], door1[3]])]).astype(object)
            if has_both_doors:
                self.data.at[ind, 'has_door2'] = has_both_doors
                door2 = np.array([np.nanmean(sq_x_positions[6:]), np.nanmean(sq_y_positions[6:]), np.nanmean(tr_x_positions[3:5]), np.nanmean(tr_y_positions[3:5])])
                self.data.at[ind, 'door2'] = door2.astype(object)
                self.data.at[ind, 'door2_cent'] = np.array([np.nanmean([door2[0], door2[2]]), np.nanmean([door2[1], door2[3]])]).astype(object)

        # animal speed
        for ind, row in self.data.iterrows():
            temp_time = np.diff(row['trial_timestamps'])
            x = np.diff(row['nose_x_cm']); y = np.diff(row['nose_y_cm'])
            if len(x) == len(temp_time):
                xspeed = list((x/temp_time)**2)
            elif len(x) > len(temp_time):
                xspeed = list((x[:len(temp_time)]/temp_time)**2)
            elif len(x) < len(temp_time):
                xspeed = list((x/temp_time[:len(x)])**2)
            if len(y) == len(temp_time):
                yspeed = list((y/temp_time)**2)
            elif len(y) > len(temp_time):
                yspeed = list((y[:len(temp_time)]/temp_time)**2)
            elif len(y) < len(temp_time):
                yspeed = list((y/temp_time[:len(y)])**2)
            self.data.at[ind, 'speed'] = np.sqrt(xspeed + yspeed).astype(object)

        kmeans_input = np.concatenate([flatten_series(self.data['obs_sq_x_pos']),
                                       flatten_series(self.data['obs_sq_y_pos']),
                                       flatten_series(self.data['obs_tr_x_pos']),
                                       flatten_series(self.data['obs_tr_y_pos'])], axis=1)

        kmeans_input[np.isnan(kmeans_input)] = 0
        labels = KMeans(n_clusters=self.num_clusters_to_use).fit(kmeans_input).labels_

        plt.subplots(2,1)
        plt.subplot(2,1,1)
        plt.hist(labels)
        plt.subplot(2,1,2)
        plt.plot(labels)
        plt.tight_layout()
        pdf.savefig()
        plt.close()

        plt.subplots(10,int(np.ceil(len(labels)/10)),figsize=(30,30))
        sq_idx = [0,1,3,2,4]
        for c, l in enumerate(labels):
            plt.subplot(10,int(np.ceil(len(labels)/10)),c+1)
            obstacles = kmeans_input[c,:]
            obstacles[obstacles==0] = np.nan
            plt.plot(np.append(obstacles[:4], obstacles[0])[sq_idx], np.append(obstacles[8:12], obstacles[8])[sq_idx], 'k-')
            plt.plot(np.append(obstacles[4:8], obstacles[4])[sq_idx], np.append(obstacles[12:16], obstacles[12])[sq_idx], 'k-')
            plt.plot(np.append(obstacles[16:19], obstacles[16]), np.append(obstacles[22:25], obstacles[22]), 'k-')
            plt.plot(np.append(obstacles[19:22], obstacles[19]), np.append(obstacles[25:], obstacles[25]), 'k-')
            door1 = self.data['door1'].iloc[c]
            plt.plot([door1[0], door1[2]], [door1[1], door1[3]], 'g-', linewidth=6)
            if self.data['has_door2'].iloc[c]:
                door2 = self.data['door2'].iloc[c]
                plt.plot([door2[0], door2[2]], [door2[1], door2[3]], 'g-', linewidth=6)
            plt.title(l); plt.ylim([20.03,0]); plt.xlim([0,33.30])
        plt.tight_layout()
        pdf.savefig()
        plt.close()

        self.data['obstacle_cluster'] = labels

        plt.subplots(2,int(np.ceil(self.num_clusters_to_use/2)),figsize=(20,5))
        for c in range(self.num_clusters_to_use):
            plt.subplot(2,int(np.ceil(self.num_clusters_to_use/2)),c+1)
            this_cluster = self.data[self.data['obstacle_cluster']==c].copy().reset_index()
            colors = plt.cm.magma(np.linspace(0,1,len(this_cluster)))
            mediantrace = self.get_median_trace(this_cluster)
            for ind, row in this_cluster.iterrows():
                plt.plot(mediantrace['nose_x_cm'].iloc[ind], mediantrace['nose_y_cm'].iloc[ind], '-', color=colors[ind])
            if len(this_cluster['door1']) > 1:
                door1 = np.nanmedian(flatten_series(this_cluster['door1']),0)
                plt.plot([door1[0], door1[2]], [door1[1], door1[3]], 'g-', linewidth=6)
            if len(this_cluster['door1']) > 1 and this_cluster['has_door2'].iloc[0]:
                door2 = np.nanmean(flatten_series(this_cluster['door2']),0)
                plt.plot([door2[0], door2[2]], [door2[1], door2[3]], 'g-', linewidth=6)
            plt.title(c)
        plt.tight_layout()
        pdf.savefig()
        plt.close()

        # overlay of cluster medians
        odd = self.data[self.data.index%2==0]
        even = self.data[self.data.index%2==1]
        direction_count = 0
        plt.subplots(1,2, figsize=(8,5))
        for direction_df in [odd, even]:
            if direction_count == 0:
                leftcolor='g'; rightcolor='b'
            else:
                leftcolor='b'; rightcolor='g'
            plt.subplot(1,2,direction_count+1)
            for c in range(self.num_clusters_to_use):
                cluster_color = plt.cm.magma(np.linspace(0,1,self.num_clusters_to_use))[c]
                this_cluster = direction_df[direction_df['obstacle_cluster']==c].copy().reset_index()
                for ind, row in this_cluster.iterrows():
                    door1 = np.nanmedian(flatten_series(this_cluster['door1']),0)
                    plt.plot([door1[0], door1[2]], [door1[1], door1[3]], '-', linewidth=6, color=cluster_color)
                    if this_cluster['has_door2'].iloc[0]:
                        door2 = np.nanmean(flatten_series(this_cluster['door2']),0)
                        plt.plot([door2[0], door2[2]], [door2[1], door2[3]], '-', linewidth=6, color=cluster_color)
                    plt.plot(row['leftportT_x_cm'], row['leftportT_y_cm'], '.', color=leftcolor)
                    plt.plot(row['rightportT_x_cm'], row['rightportT_y_cm'], '.', color=rightcolor)
                plt.ylim([20.03,0]); plt.xlim([0,33.30])
                if len(this_cluster) > 0:
                    plt.plot(self.get_median_trace(this_cluster)['median_x_cm'].iloc[0], self.get_median_trace(this_cluster)['median_y_cm'].iloc[0], '-', color=cluster_color)
            direction_count += 1
            plt.tight_layout()
            pdf.savefig(); plt.close()

        # properties over time (e.g. time active, speed, etc.)
        plt.subplots(2,3, figsize=(15,8))
        median_speed = []; max_speed = []; time_active = []
        for ind, row in self.data.iterrows():
            median_speed.append(np.median(row['speed']))
            max_speed.append(np.max(row['speed']))
            time_active.append(np.sum(row['speed']>5)/60)
            self.data.at[ind, 'median_speed'] = np.median(row['speed'])
            self.data.at[ind, 'max_speed'] = np.max(row['speed'])
            self.data.at[ind, 'time_active'] = np.sum(row['speed']>5)/60
        plt.subplot(2,3,1)
        slow = self.data['speed'].iloc[np.nanargmin(median_speed)]
        slowT = np.linspace(0,1,len(slow))
        med = self.data['speed'].iloc[np.argsort(median_speed)[len(median_speed)//2]]
        medT = np.linspace(0,1,len(med))
        fast = self.data['speed'].iloc[np.nanargmax(median_speed)]
        fastT = np.linspace(0,1,len(fast))
        fake_time = np.linspace(0,1,100)
        plt.plot(interp1d(slowT, slow, bounds_error=False)(fake_time))
        plt.plot(interp1d(medT, med, bounds_error=False)(fake_time))
        plt.plot(interp1d(fastT, fast, bounds_error=False)(fake_time))
        plt.subplot(2,3,2)
        plt.hist(median_speed); plt.xlabel('median speed (cm/sec)')
        plt.subplot(2,3,3)
        plt.hist(max_speed); plt.xlabel('max speed (cm/sec)')
        plt.subplot(2,3,4)
        plt.hist(time_active); plt.xlabel('time active (sec)')
        plt.subplot(2,3,5)
        plt.plot([i for i in max_speed if ~np.isnan(i)]); plt.xlabel('trial'); plt.ylabel('max speed (cm/sec)')
        plt.subplot(2,3,6)
        plt.plot([i for i in time_active if ~np.isnan(i)]); plt.xlabel('trial'); plt.ylabel('time active (sec)')
        plt.tight_layout()
        pdf.savefig()
        plt.close()

        plt.subplots(4,4,figsize=(20,18))
        fake_time = np.linspace(0,1,100)
        plt.subplot(4,4,1)
        for ind, data in self.data['speed'].iteritems():
            plt.plot(interp1d(np.linspace(0,1,len(data)), data, bounds_error=False)(fake_time), '-', alpha=0.2); plt.ylabel('speed'); plt.xlabel('time')
        plt.subplot(4,4,2)
        for ind, data in self.data['head_angle'].iteritems():
            plt.plot(interp1d(np.linspace(0,1,len(data)), data, bounds_error=False)(fake_time), '-', alpha=0.2); plt.ylabel('angle to horizontal'); plt.xlabel('time')
        odd = self.data[self.data.index%2==0]
        even = self.data[self.data.index%2==1]
        for direction_num in range(2):
            plt.subplot(4,4,3+direction_num)
            direction_df = [odd, even][direction_num]
            for ind, row in direction_df.iterrows():
                dist = np.array(self.distance_from_nose(row, 'leftportT'))
                plt.plot(interp1d(np.linspace(0,1,len(dist)), dist, bounds_error=False)(fake_time), alpha=0.2)
                if self.approaching_target(dist) is True:
                    plt.title('distance to target port (moving left)')
                    direction_df.at[ind, 'approach'] = True
                else:
                    plt.title('distance to previous port (moving left)')
                    direction_df.at[ind, 'approach'] = False
                plt.plot(0,dist[0],'.',color='g')
                plt.plot(100,dist[-1],'.',color='b')
                if direction_num==0 and self.approaching_target(dist):
                    savename = 'odd_dist_to_target_port'
                elif direction_num==0 and not self.approaching_target(dist):
                    savename = 'odd_dist_to_previous_port'
                elif direction_num==1 and self.approaching_target(dist):
                    savename = 'even_dist_to_target_port'
                elif direction_num==1 and not self.approaching_target(dist):
                    savename = 'even_dist_to_previous_port'
                self.data.at[ind, savename] = dist.astype(object)
        for direction_num in range(2):
            plt.subplot(4,4,5+direction_num)
            direction_df = [odd, even][direction_num]
            for ind, row in direction_df.iterrows():
                dist = np.array(self.distance_from_nose(row, 'rightportT'))
                plt.plot(interp1d(np.linspace(0,1,len(dist)), dist, bounds_error=False)(fake_time), alpha=0.2)
                if self.approaching_target(dist):
                    plt.title('distance to previous port (moving right)')
                else:
                    plt.title('distance to target port (moving right)')
                plt.plot(0,dist[0],'.',color='g')
                plt.plot(100,dist[-1],'.',color='b')
                if direction_num==0 and self.approaching_target(dist):
                    savename = 'odd_dist_to_target_port'
                elif direction_num==0 and not self.approaching_target(dist):
                    savename = 'odd_dist_to_previous_port'
                elif direction_num==1 and self.approaching_target(dist):
                    savename = 'even_dist_to_target_port'
                elif direction_num==1 and not self.approaching_target(dist):
                    savename = 'even_dist_to_previous_port'
                self.data.at[ind, savename] = dist.astype(object)
        for direction_num in range(2):
            plt.subplot(4,4,7+direction_num)
            direction_df = [odd, even][direction_num]
            for ind, row in direction_df.iterrows():
                ang = np.array(self.angle_from_nose(row, 'leftportT'))
                plt.plot(interp1d(np.linspace(0,1,len(ang)), ang, bounds_error=False)(fake_time), alpha=0.2)
                if row['approach']:
                    plt.title('angle to target port (headed left')
                else:
                    plt.title('angle to previous port (headed left)')
                plt.plot(0,ang[0],'.',color='g')
                plt.plot(100,ang[-1],'.',color='b')
                if direction_num==0 and row['approach']:
                    savename = 'odd_ang_to_target_port'
                elif direction_num==0 and not row['approach']:
                    savename = 'odd_ang_to_previous_port'
                elif direction_num==1 and row['approach']:
                    savename = 'even_ang_to_target_port'
                elif direction_num==1 and not row['approach']:
                    savename = 'even_ang_to_previous_port'
                self.data.at[ind, savename] = ang.astype(object)
        for direction_num in range(2):
            plt.subplot(4,4,9+direction_num)
            direction_df = [odd, even][direction_num]
            for ind, row in direction_df.iterrows():
                ang = np.array(self.angle_from_nose(row, 'rightportT'))
                plt.plot(interp1d(np.linspace(0,1,len(ang)), ang, bounds_error=False)(fake_time), alpha=0.2)
                if row['approach']:
                    plt.title('angle to target port')
                else:
                    plt.title('angle to previous port')
                plt.plot(0,ang[0],'.',color='g')
                plt.plot(100,ang[-1],'.',color='b')
                if direction_num==0 and row['approach']:
                    savename = 'odd_ang_to_target_port'
                elif direction_num==0 and not row['approach']:
                    savename = 'odd_ang_to_previous_port'
                elif direction_num==1 and row['approach']:
                    savename = 'even_ang_to_target_port'
                elif direction_num==1 and not row['approach']:
                    savename = 'even_ang_to_previous_port'
                self.data.at[ind, savename] = ang.astype(object)
        for direction_num in range(2):
            plt.subplot(4,4,11+direction_num)
            direction_df = [odd, even][direction_num]
            for ind, row in direction_df.iterrows():
                dist = np.array(self.distance_from_nose(row, 'door1_cent', xy_together=True))
                plt.plot(interp1d(np.linspace(0,1,len(dist)), dist, bounds_error=False)(fake_time), alpha=0.2)
                plt.title('distance to door1')
                if direction_num==0:
                    savename = 'odd_dist_to_door1'
                else:
                    savename = 'even_dist_to_door1'
                self.data.at[ind, savename] = dist.astype(object)
        for direction_num in range(2):
            plt.subplot(4,4,13+direction_num)
            direction_df = [odd, even][direction_num]
            for ind, row in direction_df.iterrows():
                if row['has_door2']:
                    ang = np.array(self.angle_from_nose(row, 'door1_cent', xy_together=True))
                    plt.plot(interp1d(np.linspace(0,1,len(ang)), ang, bounds_error=False)(fake_time), alpha=0.2)
                    plt.title('angle to door1')
                    if direction_num==0:
                        savename = 'odd_ang_to_door1'
                    else:
                        savename = 'even_ang_to_door1'
                    self.data.at[ind, savename] = ang.astype(object)
        plt.tight_layout()
        pdf.savefig()
        plt.close()

        pdf.close()
        self.data.to_hdf(os.path.join(self.trial_path, (self.data['animal'].iloc[0]+'_'+str(self.data['date'].iloc[0])+'_'+str(self.data['task'].iloc[0])+'.h5')), 'w')

class AvoidanceSession(BaseInput):
    def __init__(self, metadata_path, task='oa'):
        with open(metadata_path) as f:
            self.metadata = json.load(f)

        self.path = self.metadata['path']
        self.dates_list = [i for i in list(self.metadata.keys()) if i != 'path']

        if task=='oa':
            self.is_pillar_avoidance = True
            self.is_gap_detection = False
        elif task=='gd':
            self.is_pillar_avoidance = False
            self.is_gap_detection = True

        if self.is_pillar_avoidance:
            self.dlc_project = {'light':'/home/niell_lab/Documents/deeplabcut_projects/object_avoidance-Mike-2021-08-31/config.yaml',
                                'dark':None}
        elif self.is_gap_detection:
            self.dlc_project = {'light':'/home/niell_lab/Documents/deeplabcut_projects/gap_determination-Kana-2021-10-19/config.yaml',
                                'dark':'/home/niell_lab/Documents/deeplabcut_projects/dark_gap_determination-Kana-2021-11-08/config.yaml'}
        self.camname = 'top1'
        self.generic_camconfig = {
            'paths': {
                'dlc_projects': {
                    self.camname: self.dlc_project
                },
            },
            'internals': {
                'follow_strict_naming': False,
                'crop_for_dlc': False,
                'filter_dlc_predictions': False,
                'multianimal_top_project': False,
                'likelihood_threshold': 0.99
            }
        }

    def change_dlc_project(self, project_paths):
        """
        project_paths should be a dict like... {'light':/path/to/config.yaml, 'dark':/path/to/config.yaml}
        """
        # update object
        self.dlc_project = project_paths
        # also update config
        self.generic_camconfig['paths']['dlc_projects'][self.camname] = self.dlc_project

    def preprocess(self):
        for date in self.dates_list:
            date_dir = os.path.join(self.path, date)
            for animal in [i for i in list(self.metadata[date].keys())]:
                animal_dir = os.path.join(date_dir, animal)
                camconfig = self.generic_camconfig
                camconfig['animal_directory'] = animal_dir
                isdark = ('dark' in animal_dir)
                if not isdark:
                    print('using light network')
                    camconfig['paths']['dlc_projects'][self.camname] = self.dlc_project['light']
                elif isdark:
                    print('using dark network')
                    camconfig['paths']['dlc_projects'][self.camname] = self.dlc_project['dark']
                for recording_name in [k for k,v in self.metadata[date][animal].items()]:
                    recording_dir = os.path.join(animal_dir, recording_name)
                    name = '_'.join(os.path.splitext(os.path.split([i for i in find('*.avi', recording_dir) if all(bad not in i for bad in ['plot','IR','rep11','betafpv','side_gaze','._'])][0])[1])[0].split('_')[:-2])
                    tc = Topcam(config=camconfig, recording_name=name, recording_path=recording_dir, camname=self.camname)
                    tc.pose_estimation()
                    tc.gather_camera_files()
                    tc.pack_position_data()
                    tc.pack_video_frames()
                    tc.pt_names = list(tc.xrpts['point_loc'].values)
                    tc.filter_likelihood()
                    tc.get_head_body_yaw()
                    tc.save_params()

    def gather_all_sessions(self):
        data_dict = {'date': [],
                    'animal': [],
                    'task': [],
                    'poke1_ts':[],
                    'poke2_ts': [],
                    'top1_ts': [],
                    'poke1_t0':[],
                    'poke2_t0': [],
                    'top1_t0': []}
        # list of dates for analysis
        data_path = Path(self.path).expanduser()
        # populate dict with metadata and timestamps
        for date in self.dates_list:
            use_animals = [k for k,v in self.metadata[date].items()]
            for ani in use_animals:
                for task in os.listdir(data_path / date / ani):
                    data_paths = [str(i) for i in list((data_path / date / ani/ task).rglob('*.csv'))]
                    data_paths = [i for i in data_paths if 'spout' not in i]
                    if data_paths != []:
                        _, name = os.path.split(data_paths[1])
                        split_name = name.split('_')
                        data_dict['date'].append(split_name[0])
                        data_dict['animal'].append(split_name[1])
                        data_dict['task'].append(split_name[4])
                    for ind, csv in enumerate(data_paths):
                        self.timestamp_path = csv
                        time = self.read_timestamp_file()
                        _, name = os.path.split(data_paths[ind])
                        split_name = name.split('_')
                        data_dict[split_name[5] +'_ts'].append(time)
                        data_dict[split_name[5] +'_t0'].append(time[0])
        self.all_sessions = pd.DataFrame.from_dict(data_dict)

    def process(self, videos=False):
        self.gather_all_sessions()
        for trial_ind, trial_row in tqdm(self.all_sessions.iterrows()):
            try:
                # analyze each trial
                trial = AvoidanceTrial(trial_row, self.path, self.metadata)
                dlc_h5 = find('*'+str(trial_row['date'])+'*'+trial_row['animal']+'*'+str(trial_row['task'])+'*.h5', self.path)
                if dlc_h5 == []:
                    continue
                trial_path, _ = os.path.split(dlc_h5[0])
                trial_name = '_'.join(os.path.splitext(os.path.split([i for i in find('*.avi', trial_path) if all(bad not in i for bad in ['plot','IR','rep11','betafpv','side_gaze','._'])][0])[1])[0].split('_')[:-1])
                trial.add_tracking(trial_name, trial_path)
                if self.is_pillar_avoidance:
                    trial.pillar_avoidance()
                elif self.is_gap_detection:
                    trial.gap_detection()
                # make short diagnostic video
                if videos:
                    trial.make_videos()
            except:
                pass


"""CRF training utilities.

Two feature sets are provided:

   1. Default feature set by sklearn_crfsuite
   2. Liu et al. (2015) features used in de-identification shared task.

Features are encoded in python-crfsuite format:
https://python-crfsuite.readthedocs.io/en/latest/pycrfsuite.html#pycrfsuite.ItemSequence
"""

import re
import string
from typing import Callable, Dict, List, Tuple

import sklearn_crfsuite
from tqdm import tqdm
from unidecode import unidecode

from deidentify.methods.tagging_utils import Token

NEWLINE_REGEX = re.compile(r'\n')
SPACE_REGEX = re.compile(r'\s')

def sent2features(sent: List[Token],
                  feature_extractor: Callable[[List[Token], int], Dict]) -> List[Dict]:
    """Convert a sentence to features in python-crfsuite format.

    python-crfsuite can't handle feature values that contain whitespace or newline characters. These
    characters are replaced with a special #SPACE and #NEWLINE token.

    See issues:
    https://github.com/scrapinghub/python-crfsuite/issues/14
    https://github.com/scrapinghub/python-crfsuite/issues/71

    Parameters
    ----------
    sent : List[Token]
        A sentence constituded of a list of tokens.
    feature_extractor : Callable[[List[Token], int], Dict]
        Callable that represents a token at position `i: int` as a feature dict.

    Returns
    -------
    sent_features : List[Dict]
        List of feature dicts per token. `len(sent_features) == len(sent)`

    """
    sent_features = []

    for i in range(len(sent)):
        token_features = feature_extractor(sent, i)

        for feature_name, value in token_features.items():
            if not isinstance(value, str):
                continue

            value = NEWLINE_REGEX.sub('#NEWLINE', value)
            value = SPACE_REGEX.sub('#SPACE', value)
            token_features[feature_name] = value

        sent_features.append(token_features)

    return sent_features

def sent2labels(sent):
    return [token.label for token in sent]

def sents_to_features_and_labels(sents, feature_extractor):
    X = [sent2features(s, feature_extractor) for s in sents]
    y = [sent2labels(s) for s in sents]
    return X, y

class SentenceFilterCRF(sklearn_crfsuite.CRF):
    """Custom CRF implementation that allows to ignore entire sentences during training/prediction
    time. A default label will be assigned to all tokens within that sentence.

    The confidence of predictions within ignored sentences is set to 1.
    """

    # pylint: disable=too-many-arguments
    # pylint: disable=R0914
    # scikit-learn estimators explicitly have to mention keyword arguments (no *args, **kwargs)
    def __init__(self,
                 ignore_sentence,
                 ignored_label,
                 algorithm=None,
                 min_freq=None,
                 all_possible_states=None,
                 all_possible_transitions=None,
                 c1=None,
                 c2=None,
                 max_iterations=None,
                 num_memories=None,
                 epsilon=None,
                 period=None,
                 delta=None,
                 linesearch=None,
                 max_linesearch=None,
                 calibration_eta=None,
                 calibration_rate=None,
                 calibration_samples=None,
                 calibration_candidates=None,
                 calibration_max_trials=None,
                 pa_type=None,
                 c=None,
                 error_sensitive=None,
                 averaging=None,
                 variance=None,
                 gamma=None,
                 verbose=False,
                 model_filename=None,
                 keep_tempfiles=False,
                 trainer_cls=None):
        self.ignore_sentence = ignore_sentence
        self.ignored_label = ignored_label

        super(SentenceFilterCRF, self).__init__(
            algorithm=algorithm,
            min_freq=min_freq,
            all_possible_states=all_possible_states,
            all_possible_transitions=all_possible_transitions,
            c1=c1,
            c2=c2,
            max_iterations=max_iterations,
            num_memories=num_memories,
            epsilon=epsilon,
            period=period,
            delta=delta,
            linesearch=linesearch,
            max_linesearch=max_linesearch,
            calibration_eta=calibration_eta,
            calibration_rate=calibration_rate,
            calibration_samples=calibration_samples,
            calibration_candidates=calibration_candidates,
            calibration_max_trials=calibration_max_trials,
            pa_type=pa_type,
            c=c,
            error_sensitive=error_sensitive,
            averaging=averaging,
            variance=variance,
            gamma=gamma,
            verbose=verbose,
            model_filename=model_filename,
            keep_tempfiles=keep_tempfiles,
            trainer_cls=trainer_cls)

    def _filter_xy(self, X, y):
        return zip(*filter(lambda xy: not self.ignore_sentence(xy[0]), zip(X, y)))

    def fit(self, X, y, X_dev=None, y_dev=None):
        X, y = self._filter_xy(X, y)

        if X_dev is not None:
            X_dev, y_dev = self._filter_xy(X_dev, y_dev)

        return super(SentenceFilterCRF, self).fit(X, y, X_dev=X_dev, y_dev=y_dev)

    def predict(self, X, verbose=False):
        X = tqdm(X, disable=not verbose, desc='Tag sentences')
        return super().predict(X)

    def predict_single(self, xseq):
        if self.ignore_sentence(xseq):
            return [self.ignored_label] * len(xseq)

        return super().predict_single(xseq)

    def predict_marginals_single(self, xseq):
        if self.ignore_sentence(xseq):
            ignored_marginals = {c: 0 for c in self.classes_}
            ignored_marginals[self.ignored_label] = 1
            return [ignored_marginals] * len(xseq)

        return super().predict_marginals_single(xseq)

def sklearn_crfsuite_feature_extractor(sent, i):
    """
    Taken from:
    https://sklearn-crfsuite.readthedocs.io/en/latest/tutorial.html
    """
    word = sent[i].text
    pos_tag = sent[i].pos_tag

    features = {
        'bias': 1.0,
        'word.lower()': word.lower(),
        'word[-3:]': word[-3:],
        'word[-2:]': word[-2:],
        'word.isupper()': word.isupper(),
        'word.istitle()': word.istitle(),
        'word.isdigit()': word.isdigit(),
        'pos_tag': pos_tag,
        'pos_tag[:2]': pos_tag[:2],
    }
    if i > 0:
        word1 = sent[i - 1].text
        pos_tag1 = sent[i - 1].pos_tag
        features.update({
            '-1:word.lower()': word1.lower(),
            '-1:word.istitle()': word1.istitle(),
            '-1:word.isupper()': word1.isupper(),
            '-1:pos_tag': pos_tag1,
            '-1:pos_tag[:2]': pos_tag1[:2],
        })
    else:
        features['BOS'] = True

    if i < len(sent) - 1:
        word1 = sent[i + 1].text
        pos_tag1 = sent[i + 1].pos_tag
        features.update({
            '+1:word.lower()': word1.lower(),
            '+1:word.istitle()': word1.istitle(),
            '+1:word.isupper()': word1.isupper(),
            '+1:pos_tag': pos_tag1,
            '+1:pos_tag[:2]': pos_tag1[:2],
        })
    else:
        features['EOS'] = True

    return features

def liu_feature_extractor(sent, i):
    """Reproduces the features used by Liu et al. (2015).

    Does not include word representation (word2vec, brown clusters) and gazetteer features.

    Reference:
    Liu, Z., et al. (2015). Automatic de-identification of electronic medical records using
    token-level and character-level conditional random fields. Journal of Biomedical Informatics,
    58, S47–S52. https://doi.org/10.1016/J.JBI.2015.06.009
    """
    token = sent[i]

    null_token = Token(text='<PAD>', pos_tag='<PAD>', label='', ner_tag=None)
    sent_window = list_window(sent, center=i, window=(2, 2), oob_item=null_token)
    token_window = [t.text.lower() for t in sent_window]
    pos_window = [t.pos_tag for t in sent_window]
    text_lower = token.text.lower()

    features = {}
    features.update(_ngram_feature_group(token_window, N=1, group_name='bow[-2:2].uni'))
    features.update(_ngram_feature_group(token_window, N=2, group_name='bow[-2:2].bi'))
    features.update(_ngram_feature_group(token_window, N=3, group_name='bow[-2:2].tri'))

    features.update(_ngram_feature_group(pos_window, N=1, group_name='pos[-2:2].uni'))
    features.update(_ngram_feature_group(pos_window, N=2, group_name='pos[-2:2].bi'))
    features.update(_ngram_feature_group(pos_window, N=3, group_name='pos[-2:2].tri'))

    sent_window = list_window(sent, center=i, window=(1, 1), oob_item=null_token)
    pos_window = [t.pos_tag for t in sent_window]
    sep = join_features
    features['bowpos.w0p-1'] = sep((text_lower, pos_window[0]))
    features['bowpos.w0p0'] = sep((text_lower, pos_window[1]))
    features['bowpos.w0p1'] = sep((text_lower, pos_window[2]))
    features['bowpos.w0p-1p0'] = sep((text_lower, pos_window[0], pos_window[1]))
    features['bowpos.w0p0p1'] = sep((text_lower, pos_window[1], pos_window[2]))
    features['bowpos.w0p-1p1'] = sep((text_lower, pos_window[0], pos_window[2]))
    features['bowpos.w0p-1p0p1'] = sep((text_lower, pos_window[0], pos_window[1], pos_window[2]))

    features['sent.len(sent)'] = len(sent)
    features['sent.end_mark'] = sent[-1].text.strip() in ['!', '?', '.']
    features['sent.has_unmatched_bracket'] = has_unmatched_bracket(sent)

    for j in range(1, 6):
        features['suffix[-{}:]'.format(j)] = text_lower[-j:]
        features['prefix[:{}]'.format(j)] = text_lower[:j]

    features['word.isupper()'] = token.text.isupper()
    features['word.istitle()'] = token.text.istitle()
    features['word.isdigit()'] = token.text.isdigit()
    features['word.contains_digit'] = any(c.isdigit() for c in token.text)
    features['word.has_upper_inside'] = any(c.isupper() for c in token.text[1:])
    features['word.has_punct_inside'] = any(c in string.punctuation for c in token.text[1:])
    features['word.has_digit_inside'] = any(c.isdigit() for c in token.text[1:])
    features['word.is_ascii'] = all(ord(c) < 128 for c in token.text)
    features['word.ner_tag'] = token.ner_tag
    features['word.pos_tag'] = token.pos_tag

    shape = word_shape(token.text)
    features['shape.long'] = shape
    features['shape.short'] = collapse_word_shape(shape)

    return features

def join_features(feature_list):
    return '|'.join(feature_list)

def ngrams(tokens, N):
    return [tuple(tokens[i:i + N]) for i in range(len(tokens) - N + 1)]

def list_window(sent: List, center: int, window: Tuple[int, int], oob_item=None) -> List:
    """Get a window of tokens within a sentence.

    Parameters
    ----------
    sent : List
        A list of tokens.
    center : int
        The index acting as center of the window.
    window : Tuple[int, int]
        The window width. `window[0]` is elements before center, `window[1]` is elements after
        center. Interval is closed.
    oob_item : type
        The item to return if window indexes are out of bounds of `sent`.

    Returns
    -------
    tokens : List
        The tokens within the given window.

    """
    tokens = []
    for i in range(center - window[0], center + window[1] + 1):
        if i < 0:
            tokens.append(oob_item)
        elif i >= len(sent):
            tokens.append(oob_item)
        else:
            tokens.append(sent[i])
    return tokens

def _ngram_feature_group(tokens, N, group_name, sep=join_features):
    features = {}
    token_ngrams = ngrams(tokens, N)
    for j, item in enumerate(token_ngrams):
        features['{}.{}'.format(group_name, j)] = sep(item)
    return features

def has_unmatched_bracket(sent):
    n_open = 0

    for token in sent:
        if token.text == '(':
            n_open += 1
        elif token.text == ')':
            n_open -= 1

    return n_open > 0

def word_shape(token):
    shape = ''
    for c in unidecode(token):
        if c in string.ascii_lowercase:
            shape += 'a'
        elif c in string.ascii_uppercase:
            shape += 'A'
        elif c in string.digits:
            shape += '#'
        else:
            shape += '-'
    return shape

def collapse_word_shape(shape):
    collapsed = ''
    current = None
    for c in shape:
        if c == current:
            continue
        collapsed += c
        current = c
    return collapsed


from urllib.request import urlopen as uReq
from bs4 import BeautifulSoup as soup

filename = "nike_cr7.csv"
f = open(filename, "w")

for i in range(1,6):
	url = 'https://www.amazon.com/s/ref=sr_pg_'+ str(i) +'?rh=i%3Aaps%2Ck%3Anike+cr7&page=3&keywords=nike+cr7&ie=UTF8&qid=1535481489'

	client = uReq(url)
	webpage_html = client.read()
	client.close()

	html_page = soup(webpage_html, "html.parser")
	headers = "Name, Price\n"

	f.write(headers)

	div_containers = html_page.findAll("div", {"class" : "a-fixed-left-grid-inner"})

	for div_container in div_containers:
		name_temp = div_container.findAll("div", {"class:", "a-col-right"})[0]
		name = name_temp.div.div.a["title"]
		print(name)
		price = name_temp.find("span", {"class:", "a-offscreen"})
		if price != None:
			print(price.get_text())
			f.write(str(name).replace(",", "|") + "," + str(price.get_text()) + "\n")

f.close()


from precalc_feats import *
from resnet_model import *

from kaggle import *
from setup import *
from utils import save_array

def train_ensemble():
    nb_models = 5 # train 5 ensemble models
    models = []
    model_paths = []

    for run in range(nb_models):
        print("====== Ensemble model: %d ======" % run)
        m = ResnetModel('data/')
        model_prefix = "resnet_ft_da%d" % run
        m.model_name = model_prefix + m.model_name

        # finetune for last identity block
        m.fine_tune()

        print("====== training model ======")
        train_batches = create_batches('data/train/', shuffle=True, use_da=True)
        val_batches = create_batches('data/valid/', shuffle=True, use_da=True)
        # m.train(nb_epoch=20, use_da=True)
        m.train_on_batches(train_batches, val_batches, nb_epoch=20)

        # append model
        models = models + [m]
        model_paths = model_paths + [m.model_path]

    return models, model_paths

def test_ensemble(models):
    nb_test_samples = 1000
    nb_classes = 8
    nb_augmentations = 5
    preds = np.zeros((nb_test_samples, nb_classes))

    for test_run in range(nb_augmentations):
        # make test batch randomly with data aug
        print("====== data-aug test batch: %d ======" % test_run)
        test_batches = create_batches('data/test/', shuffle=False, use_da=True)
        preds_aug = np.zeros((nb_test_samples, nb_classes))

        for ind, m in enumerate(models):
            print("====== running test model: %d ======" % ind)
            _preds = m.test_on_batch(test_batches)
            preds_aug = preds_aug + _preds

        preds_aug /= len(models)
        preds = preds + preds_aug

    preds /= nb_augmentations
    save_array('submits/resnet_ft_ens_preds.gz', preds)
    return preds

def submit_ensemble(preds):
    # get test batch from any model
    test_batches = create_batches('data/test/', shuffle=False, use_da=False)

    print("======= making submission ========")
    submits_path = 'submits/resnet_ft_ens_preds.gz'
    submit(preds, test_batches, submits_path)

    print("======= pushing to kaggle ========")
    push_to_kaggle(submits_path)

def load_models():
    models = []
    nb_models = 5
    for run in range(nb_models):
        print("====== Loading ensemble model: %d ======" % run)
        m = ResnetModel('data/')
        model_prefix = "resnet_ft_da%d" % run
        m.model_name = model_prefix + m.model_name
        print("model path: %s" % m.model_path)
        m.load_model()

        models = models + [m]
    return models

if __name__ == "__main__":
    models, model_paths = train_ensemble()
    preds = test_ensemble(models)
    submit_ensemble(preds)


from __future__ import print_function
from __future__ import division

import matplotlib.pyplot as pl
import numpy as np
import sys
import os

def convert(this_one, dest, name):
    array = np.load(this_one)
    new_array = 1 - array
    green = 255*new_array
    red = blue = np.zeros_like(green)
    full_image = np.dstack((red, green, blue))
    save_path = os.path.join(dest, '{}.png'.format(name))
    pl.imsave(save_path, full_image)
    pass

if __name__ == '__main__':
    folder_path = sys.argv[1]
    files = [f for f in os.listdir(folder_path)]
    files.sort(key=lambda f: int(filter(str.isdigit, f)))
    print(files)
    count = 0
    for np_arr in files:
        convert(os.path.join(folder_path, np_arr), dest=folder_path, name=count)
        count += 1


from __future__ import absolute_import, division, print_function
__metaclass__ = type

ANSIBLE_METADATA = {'metadata_version': '1.1',
                    'status': ['preview'],
                    'supported_by': 'network'}

DOCUMENTATION = """
---
module: iosxr_facts
version_added: "2.2"
author: "Ricardo Carrillo Cruz (@rcarrillocruz)"
short_description: Collect facts from remote devices running IOS XR
description:
  - Collects a base set of device facts from a remote device that
    is running IOS XR.  This module prepends all of the
    base network fact keys with C(ansible_net_<fact>).  The facts
    module will always collect a base set of facts from the device
    and can enable or disable collection of additional facts.
extends_documentation_fragment: iosxr
notes:
  - Tested against IOS XRv 6.1.2
  - This module does not support netconf connection
options:
  gather_subset:
    description:
      - When supplied, this argument will restrict the facts collected
        to a given subset.  Possible values for this argument include
        all, hardware, config, and interfaces.  Can specify a list of
        values to include a larger subset.  Values can also be used
        with an initial C(M(!)) to specify that a specific subset should
        not be collected.
    required: false
    default: '!config'
"""

EXAMPLES = """
# Collect all facts from the device
- iosxr_facts:
    gather_subset: all

# Collect only the config and default facts
- iosxr_facts:
    gather_subset:
      - config

# Do not collect hardware facts
- iosxr_facts:
    gather_subset:
      - "!hardware"
"""

RETURN = """
ansible_net_gather_subset:
  description: The list of fact subsets collected from the device
  returned: always
  type: list

# default
ansible_net_version:
  description: The operating system version running on the remote device
  returned: always
  type: str
ansible_net_hostname:
  description: The configured hostname of the device
  returned: always
  type: string
ansible_net_image:
  description: The image file the device is running
  returned: always
  type: string

# hardware
ansible_net_filesystems:
  description: All file system names available on the device
  returned: when hardware is configured
  type: list
ansible_net_memfree_mb:
  description: The available free memory on the remote device in Mb
  returned: when hardware is configured
  type: int
ansible_net_memtotal_mb:
  description: The total memory on the remote device in Mb
  returned: when hardware is configured
  type: int

# config
ansible_net_config:
  description: The current active config from the device
  returned: when config is configured
  type: str

# interfaces
ansible_net_all_ipv4_addresses:
  description: All IPv4 addresses configured on the device
  returned: when interfaces is configured
  type: list
ansible_net_all_ipv6_addresses:
  description: All IPv6 addresses configured on the device
  returned: when interfaces is configured
  type: list
ansible_net_interfaces:
  description: A hash of all interfaces running on the system
  returned: when interfaces is configured
  type: dict
ansible_net_neighbors:
  description: The list of LLDP neighbors from the remote device
  returned: when interfaces is configured
  type: dict
"""
import re

from ansible.module_utils.basic import AnsibleModule
from ansible.module_utils.network.iosxr.iosxr import iosxr_argument_spec, run_command
from ansible.module_utils.six import iteritems
from ansible.module_utils.six.moves import zip

class FactsBase(object):

    def __init__(self):
        self.facts = dict()
        self.commands()

    def commands(self):
        raise NotImplementedError

class Default(FactsBase):

    def commands(self):
        return(['show version | utility head -n 20'])

    def populate(self, results):
        self.facts['version'] = self.parse_version(results['show version | utility head -n 20'])
        self.facts['image'] = self.parse_image(results['show version | utility head -n 20'])
        self.facts['hostname'] = self.parse_hostname(results['show version | utility head -n 20'])

    def parse_version(self, data):
        match = re.search(r'Version (\S+)$', data, re.M)
        if match:
            return match.group(1)

    def parse_hostname(self, data):
        match = re.search(r'^(.+) uptime', data, re.M)
        if match:
            return match.group(1)

    def parse_image(self, data):
        match = re.search(r'image file is "(.+)"', data)
        if match:
            return match.group(1)

class Hardware(FactsBase):

    def commands(self):
        return(['dir /all', 'show memory summary'])

    def populate(self, results):
        self.facts['filesystems'] = self.parse_filesystems(
            results['dir /all'])

        match = re.search(r'Physical Memory: (\d+)M total \((\d+)',
                          results['show memory summary'])
        if match:
            self.facts['memtotal_mb'] = match.group(1)
            self.facts['memfree_mb'] = match.group(2)

    def parse_filesystems(self, data):
        return re.findall(r'^Directory of (\S+)', data, re.M)

class Config(FactsBase):

    def commands(self):
        return(['show running-config'])

    def populate(self, results):
        self.facts['config'] = results['show running-config']

class Interfaces(FactsBase):

    def commands(self):
        return(['show interfaces', 'show ipv6 interface',
                'show lldp', 'show lldp neighbors detail'])

    def populate(self, results):
        self.facts['all_ipv4_addresses'] = list()
        self.facts['all_ipv6_addresses'] = list()

        interfaces = self.parse_interfaces(results['show interfaces'])
        self.facts['interfaces'] = self.populate_interfaces(interfaces)

        data = results['show ipv6 interface']
        if len(data) > 0:
            data = self.parse_interfaces(data)
            self.populate_ipv6_interfaces(data)

        if 'LLDP is not enabled' not in results['show lldp']:
            neighbors = results['show lldp neighbors detail']
            self.facts['neighbors'] = self.parse_neighbors(neighbors)

    def populate_interfaces(self, interfaces):
        facts = dict()
        for key, value in iteritems(interfaces):
            intf = dict()
            intf['description'] = self.parse_description(value)
            intf['macaddress'] = self.parse_macaddress(value)

            ipv4 = self.parse_ipv4(value)
            intf['ipv4'] = self.parse_ipv4(value)
            if ipv4:
                self.add_ip_address(ipv4['address'], 'ipv4')

            intf['mtu'] = self.parse_mtu(value)
            intf['bandwidth'] = self.parse_bandwidth(value)
            intf['duplex'] = self.parse_duplex(value)
            intf['lineprotocol'] = self.parse_lineprotocol(value)
            intf['operstatus'] = self.parse_operstatus(value)
            intf['type'] = self.parse_type(value)

            facts[key] = intf
        return facts

    def populate_ipv6_interfaces(self, data):
        for key, value in iteritems(data):
            if key in ['No', 'RPF'] or key.startswith('IP'):
                continue
            self.facts['interfaces'][key]['ipv6'] = list()
            addresses = re.findall(r'\s+(.+), subnet', value, re.M)
            subnets = re.findall(r', subnet is (.+)$', value, re.M)
            for addr, subnet in zip(addresses, subnets):
                ipv6 = dict(address=addr.strip(), subnet=subnet.strip())
                self.add_ip_address(addr.strip(), 'ipv6')
                self.facts['interfaces'][key]['ipv6'].append(ipv6)

    def add_ip_address(self, address, family):
        if family == 'ipv4':
            self.facts['all_ipv4_addresses'].append(address)
        else:
            self.facts['all_ipv6_addresses'].append(address)

    def parse_neighbors(self, neighbors):
        facts = dict()
        nbors = neighbors.split('------------------------------------------------')
        for entry in nbors[1:]:
            if entry == '':
                continue
            intf = self.parse_lldp_intf(entry)
            if intf not in facts:
                facts[intf] = list()
            fact = dict()
            fact['host'] = self.parse_lldp_host(entry)
            fact['port'] = self.parse_lldp_port(entry)
            facts[intf].append(fact)
        return facts

    def parse_interfaces(self, data):
        parsed = dict()
        key = ''
        for line in data.split('\n'):
            if len(line) == 0:
                continue
            elif line[0] == ' ':
                parsed[key] += '\n%s' % line
            else:
                match = re.match(r'^(\S+)', line)
                if match:
                    key = match.group(1)
                    parsed[key] = line
        return parsed

    def parse_description(self, data):
        match = re.search(r'Description: (.+)$', data, re.M)
        if match:
            return match.group(1)

    def parse_macaddress(self, data):
        match = re.search(r'address is (\S+)', data)
        if match:
            return match.group(1)

    def parse_ipv4(self, data):
        match = re.search(r'Internet address is (\S+)/(\d+)', data)
        if match:
            addr = match.group(1)
            masklen = int(match.group(2))
            return dict(address=addr, masklen=masklen)

    def parse_mtu(self, data):
        match = re.search(r'MTU (\d+)', data)
        if match:
            return int(match.group(1))

    def parse_bandwidth(self, data):
        match = re.search(r'BW (\d+)', data)
        if match:
            return int(match.group(1))

    def parse_duplex(self, data):
        match = re.search(r'(\w+)(?: D|-d)uplex', data, re.M)
        if match:
            return match.group(1)

    def parse_type(self, data):
        match = re.search(r'Hardware is (.+),', data, re.M)
        if match:
            return match.group(1)

    def parse_lineprotocol(self, data):
        match = re.search(r'line protocol is (.+)\s+?$', data, re.M)
        if match:
            return match.group(1)

    def parse_operstatus(self, data):
        match = re.search(r'^(?:.+) is (.+),', data, re.M)
        if match:
            return match.group(1)

    def parse_lldp_intf(self, data):
        match = re.search(r'^Local Interface: (.+)$', data, re.M)
        if match:
            return match.group(1)

    def parse_lldp_host(self, data):
        match = re.search(r'System Name: (.+)$', data, re.M)
        if match:
            return match.group(1)

    def parse_lldp_port(self, data):
        match = re.search(r'Port id: (.+)$', data, re.M)
        if match:
            return match.group(1)

FACT_SUBSETS = dict(
    default=Default,
    hardware=Hardware,
    interfaces=Interfaces,
    config=Config,
)

VALID_SUBSETS = frozenset(FACT_SUBSETS.keys())

def main():
    spec = dict(
        gather_subset=dict(default=['!config'], type='list')
    )

    spec.update(iosxr_argument_spec)

    module = AnsibleModule(argument_spec=spec,
                           supports_check_mode=True)

    warnings = list()

    gather_subset = module.params['gather_subset']

    runable_subsets = set()
    exclude_subsets = set()

    for subset in gather_subset:
        if subset == 'all':
            runable_subsets.update(VALID_SUBSETS)
            continue

        if subset.startswith('!'):
            subset = subset[1:]
            if subset == 'all':
                exclude_subsets.update(VALID_SUBSETS)
                continue
            exclude = True
        else:
            exclude = False

        if subset not in VALID_SUBSETS:
            module.fail_json(msg='Bad subset')

        if exclude:
            exclude_subsets.add(subset)
        else:
            runable_subsets.add(subset)

    if not runable_subsets:
        runable_subsets.update(VALID_SUBSETS)

    runable_subsets.difference_update(exclude_subsets)
    runable_subsets.add('default')

    facts = dict()
    facts['gather_subset'] = list(runable_subsets)

    instances = list()
    for key in runable_subsets:
        instances.append(FACT_SUBSETS[key]())

    try:
        for inst in instances:
            commands = inst.commands()
            responses = run_command(module, commands)
            results = dict(zip(commands, responses))
            inst.populate(results)
            facts.update(inst.facts)
    except Exception:
        module.exit_json(out=module.from_json(results))

    ansible_facts = dict()
    for key, value in iteritems(facts):
        key = 'ansible_net_%s' % key
        ansible_facts[key] = value

    module.exit_json(ansible_facts=ansible_facts, warnings=warnings)

if __name__ == '__main__':
    main()