"""
Set of functions to deal with dlc data
"""
import logging
import pandas as pd
import warnings

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import scipy.interpolate as interpolate
from scipy.stats import zscore

from neurodsp.smooth import smooth_interpolate_savgol
from brainbox.processing import bincount2D
import brainbox.behavior.wheel as bbox_wheel

logger = logging.getLogger('ibllib')

SAMPLING = {'left': 60,
            'right': 150,
            'body': 30}
RESOLUTION = {'left': 2,
              'right': 1,
              'body': 1}

T_BIN = 0.02  # sec
WINDOW_LEN = 2  # sec
WINDOW_LAG = -0.5  # sec


# For plotting we use a window around the event the data is aligned to WINDOW_LAG before and WINDOW_LEN after the event
def plt_window(x):
    return x + WINDOW_LAG, x + WINDOW_LEN


def insert_idx(array, values):
    idx = np.searchsorted(array, values, side="left")
    # Choose lower index if insertion would be after last index or if lower index is closer
    idx[idx == len(array)] -= 1
    idx[np.where(abs(values - array[idx - 1]) < abs(values - array[idx]))] -= 1
    # If 0 index was reduced, revert
    idx[idx == -1] = 0
    if np.all(idx == 0):
        raise ValueError('Something is wrong, all values to insert are outside of the array.')
    return idx


def likelihood_threshold(dlc, threshold=0.9):
    """
    Set dlc points with likelihood less than threshold to nan
    :param dlc: dlc pqt object
    :param threshold: likelihood threshold
    :return:
    """
    features = np.unique(['_'.join(x.split('_')[:-1]) for x in dlc.keys()])
    for feat in features:
        nan_fill = dlc[f'{feat}_likelihood'] < threshold
        dlc[f'{feat}_x'][nan_fill] = np.nan
        dlc[f'{feat}_y'][nan_fill] = np.nan

    return dlc


def get_speed(dlc, dlc_t, camera, feature='paw_r'):
    """

    :param dlc: dlc pqt table
    :param dlc_t: dlc time points
    :param camera: camera type e.g 'left', 'right', 'body'
    :param feature: dlc feature to compute speed over
    :return:
    """
    x = dlc[f'{feature}_x'] / RESOLUTION[camera]
    y = dlc[f'{feature}_y'] / RESOLUTION[camera]

    # get speed in px/sec [half res]
    s = ((np.diff(x) ** 2 + np.diff(y) ** 2) ** .5) * SAMPLING[camera]

    dt = np.diff(dlc_t)
    tv = dlc_t[:-1] + dt / 2

    # interpolate over original time scale
    if tv.size > 1:
        ifcn = interpolate.interp1d(tv, s, fill_value="extrapolate")
        return ifcn(dlc_t)


def get_speed_for_features(dlc, dlc_t, camera, features=['paw_r', 'paw_l', 'nose_tip']):
    """
    Wrapper to compute speed for a number of dlc features and add them to dlc table
    :param dlc: dlc pqt table
    :param dlc_t: dlc time points
    :param camera: camera type e.g 'left', 'right', 'body'
    :param features: dlc features to compute speed for
    :return:
    """
    for feat in features:
        dlc[f'{feat}_speed'] = get_speed(dlc, dlc_t, camera, feat)

    return dlc


def get_feature_event_times(dlc, dlc_t, features):
    """
    Detect events from the dlc traces. Based on the standard deviation between frames
    :param dlc: dlc pqt table
    :param dlc_t: dlc times
    :param features: features to consider
    :return:
    """

    for i, feat in enumerate(features):
        f = dlc[feat]
        threshold = np.nanstd(np.diff(f)) / 4
        if i == 0:
            events = np.where(np.abs(np.diff(f)) > threshold)[0]
        else:
            events = np.r_[events, np.where(np.abs(np.diff(f)) > threshold)[0]]

    return dlc_t[np.unique(events)]


def get_licks(dlc, dlc_t):
    """
    Compute lick times from the tongue dlc points
    :param dlc: dlc pqt table
    :param dlc_t: dlc times
    :return:
    """
    lick_times = get_feature_event_times(dlc, dlc_t, ['tongue_end_l_x', 'tongue_end_l_y',
                                                      'tongue_end_r_x', 'tongue_end_r_y'])
    return lick_times


def get_sniffs(dlc, dlc_t):
    """
    Compute sniff times from the nose tip
    :param dlc: dlc pqt table
    :param dlc_t: dlc times
    :return:
    """

    sniff_times = get_feature_event_times(dlc, dlc_t, ['nose_tip_y'])
    return sniff_times


def get_dlc_everything(dlc_cam, camera):
    """
    Get out features of interest for dlc
    :param dlc_cam: dlc object
    :param camera: camera type e.g 'left', 'right'
    :return:
    """

    aligned = True
    if dlc_cam.times.shape[0] != dlc_cam.dlc.shape[0]:
        # logger warning and print out status of the qc, specific serializer django!
        logger.warning('Dimension mismatch between dlc points and timestamps')
        min_samps = min(dlc_cam.times.shape[0], dlc_cam.dlc.shape[0])
        dlc_cam.times = dlc_cam.times[:min_samps]
        dlc_cam.dlc = dlc_cam.dlc[:min_samps]
        aligned = False

    dlc_cam.dlc = likelihood_threshold(dlc_cam.dlc)
    dlc_cam.dlc = get_speed_for_features(dlc_cam.dlc, dlc_cam.times, camera)
    dlc_cam['licks'] = get_licks(dlc_cam.dlc, dlc_cam.times)
    dlc_cam['sniffs'] = get_sniffs(dlc_cam.dlc, dlc_cam.times)
    dlc_cam['aligned'] = aligned

    return dlc_cam


def get_pupil_diameter(dlc):
    """
    Estimates pupil diameter by taking median of different computations.

    The two most straightforward estimates: d1 = top - bottom, d2 = left - right
    In addition, assume the pupil is a circle and estimate diameter from other pairs of points

    :param dlc: dlc pqt table with pupil estimates, should be likelihood thresholded (e.g. at 0.9)
    :return: np.array, pupil diameter estimate for each time point, shape (n_frames,)
    """
    diameters = []
    # Get the x,y coordinates of the four pupil points
    top, bottom, left, right = [np.vstack((dlc[f'pupil_{point}_r_x'], dlc[f'pupil_{point}_r_y']))
                                for point in ['top', 'bottom', 'left', 'right']]
    # First compute direct diameters
    diameters.append(np.linalg.norm(top - bottom, axis=0))
    diameters.append(np.linalg.norm(left - right, axis=0))

    # For non-crossing edges, estimate diameter via circle assumption
    for pair in [(top, left), (top, right), (bottom, left), (bottom, right)]:
        diameters.append(np.linalg.norm(pair[0] - pair[1], axis=0) * 2 ** 0.5)

    # Ignore all nan runtime warning
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)
        return np.nanmedian(diameters, axis=0)


def get_smooth_pupil_diameter(diameter_raw, camera, std_thresh=5, nan_thresh=1):
    """
    :param diameter_raw: np.array, raw pupil diameters, calculated from (thresholded) dlc traces
    :param camera: str ('left', 'right'), which camera to run the smoothing for
    :param std_thresh: threshold (in standard deviations) beyond which a point is labeled as an outlier
    :param nan_thresh: threshold (in seconds) above which we will not interpolate nans, but keep them
                       (for long stretches interpolation may not be appropriate)
    :return:
    """
    # set framerate of camera
    if camera == 'left':
        fr = SAMPLING['left']  # set by hardware
        window = 31  # works well empirically
    elif camera == 'right':
        fr = SAMPLING['right']  # set by hardware
        window = 75  # works well empirically
    else:
        raise NotImplementedError("camera has to be 'left' or 'right")

    # Raise error if too many NaN time points, in this case it doesn't make sense to interpolate
    if np.mean(np.isnan(diameter_raw)) > 0.9:
        raise ValueError(f"Raw pupil diameter for {camera} is too often NaN, cannot smooth.")
    # run savitzy-golay filter on non-nan time points to denoise
    diameter_smoothed = smooth_interpolate_savgol(diameter_raw, window=window, order=3, interp_kind='linear')

    # find outliers and set them to nan
    difference = diameter_raw - diameter_smoothed
    outlier_thresh = std_thresh * np.nanstd(difference)
    without_outliers = np.copy(diameter_raw)
    without_outliers[(difference < -outlier_thresh) | (difference > outlier_thresh)] = np.nan
    # run savitzy-golay filter again on (possibly reduced) non-nan timepoints to denoise
    diameter_smoothed = smooth_interpolate_savgol(without_outliers, window=window, order=3, interp_kind='linear')

    # don't interpolate long strings of nans
    t = np.diff(np.isnan(without_outliers).astype(int))
    begs = np.where(t == 1)[0]
    ends = np.where(t == -1)[0]
    if begs.shape[0] > ends.shape[0]:
        begs = begs[:ends.shape[0]]
    for b, e in zip(begs, ends):
        if (e - b) > (fr * nan_thresh):
            diameter_smoothed[(b + 1):(e + 1)] = np.nan  # offset by 1 due to earlier diff

    return diameter_smoothed


def plot_trace_on_frame(frame, dlc_df, cam):
    """
    Plots dlc traces as scatter plots on a frame of the video.
    For left and right video also plots whisker pad and eye and tongue zoom.

    :param frame: np.array, single video frame to plot on
    :param dlc_df: pd.Dataframe, dlc traces with _x, _y and _likelihood info for each trace
    :param cam: str, which camera to process ('left', 'right', 'body')
    :returns: matplolib.axis
    """
    # Define colors
    colors = {'tail_start': '#636EFA',
              'nose_tip': '#636EFA',
              'paw_l': '#EF553B',
              'paw_r': '#00CC96',
              'pupil_bottom_r': '#AB63FA',
              'pupil_left_r': '#FFA15A',
              'pupil_right_r': '#19D3F3',
              'pupil_top_r': '#FF6692',
              'tongue_end_l': '#B6E880',
              'tongue_end_r': '#FF97FF'}
    # Threshold the dlc traces
    dlc_df = likelihood_threshold(dlc_df)
    # Features without tube
    features = np.unique(['_'.join(x.split('_')[:-1]) for x in dlc_df.keys() if 'tube' not in x])
    # Normalize the number of points across cameras
    dlc_df_norm = pd.DataFrame()
    for feat in features:
        dlc_df_norm[f'{feat}_x'] = dlc_df[f'{feat}_x'][0::int(SAMPLING[cam] / 10)]
        dlc_df_norm[f'{feat}_y'] = dlc_df[f'{feat}_y'][0::int(SAMPLING[cam] / 10)]
        # Scatter
        plt.scatter(dlc_df_norm[f'{feat}_x'], dlc_df_norm[f'{feat}_y'], alpha=0.05, s=2, label=feat, c=colors[feat])

    plt.axis('off')
    plt.imshow(frame, cmap='gray')
    plt.tight_layout()

    ax = plt.gca()
    if cam == 'body':
        plt.title(f'{cam.capitalize()} camera')
        return ax
    # For left and right cam plot whisker pad rectangle
    # heuristic: square with side length half the distance between nose and pupil and anchored on midpoint
    p_nose = np.array(dlc_df[['nose_tip_x', 'nose_tip_y']].mean())
    p_pupil = np.array(dlc_df[['pupil_top_r_x', 'pupil_top_r_y']].mean())
    p_anchor = np.mean([p_nose, p_pupil], axis=0)
    dist = np.linalg.norm(p_nose - p_pupil)
    rect = matplotlib.patches.Rectangle((int(p_anchor[0] - dist / 4), int(p_anchor[1])), int(dist / 2), int(dist / 3),
                                        linewidth=1, edgecolor='lime', facecolor='none')
    ax.add_patch(rect)
    # Plot eye region zoom
    inset_anchor = 0 if cam == 'right' else 0.5
    ax_ins = ax.inset_axes([inset_anchor, -0.5, 0.5, 0.5])
    ax_ins.imshow(frame, cmap='gray', origin="lower")
    for feat in features:
        ax_ins.scatter(dlc_df_norm[f'{feat}_x'], dlc_df_norm[f'{feat}_y'], alpha=1, s=0.001, label=feat, c=colors[feat])
    ax_ins.set_xlim(int(p_pupil[0] - 33 * RESOLUTION[cam] / 2), int(p_pupil[0] + 33 * RESOLUTION[cam] / 2))
    ax_ins.set_ylim(int(p_pupil[1] + 38 * RESOLUTION[cam] / 2), int(p_pupil[1] - 28 * RESOLUTION[cam] / 2))
    ax_ins.axis('off')
    # Plot tongue region zoom
    p1 = np.array(dlc_df[['tube_top_x', 'tube_top_y']].mean())
    p2 = np.array(dlc_df[['tube_bottom_x', 'tube_bottom_y']].mean())
    p_tongue = np.nanmean([p1, p2], axis=0)
    inset_anchor = 0 if cam == 'left' else 0.5
    ax_ins = ax.inset_axes([inset_anchor, -0.5, 0.5, 0.5])
    ax_ins.imshow(frame, cmap='gray', origin="upper")
    for feat in features:
        ax_ins.scatter(dlc_df_norm[f'{feat}_x'], dlc_df_norm[f'{feat}_y'], alpha=1, s=0.001, label=feat, c=colors[feat])
    ax_ins.set_xlim(int(p_tongue[0] - 60 * RESOLUTION[cam] / 2), int(p_tongue[0] + 100 * RESOLUTION[cam] / 2))
    ax_ins.set_ylim(int(p_tongue[1] + 60 * RESOLUTION[cam] / 2), int(p_tongue[1] - 100 * RESOLUTION[cam] / 2))
    ax_ins.axis('off')

    plt.title(f'{cam.capitalize()} camera')
    return ax


def plot_wheel_position(wheel_position, wheel_time, trials_df):
    """
    Plots wheel position across trials, color by which side was chosen

    :param wheel_position: np.array, interpolated wheel position
    :param wheel_time: np.array, interpolated wheel timestamps
    :param trials_df: pd.DataFrame, with column 'stimOn_times' (time of stimulus onset times for each trial)
    :returns: matplotlib.axis
    """
    # Interpolate wheel data
    wheel_position, wheel_time = bbox_wheel.interpolate_position(wheel_time, wheel_position, freq=1 / T_BIN)
    # Create a window around the stimulus onset
    start_window, end_window = plt_window(trials_df['stimOn_times'])
    # Translating the time window into an index window
    start_idx = insert_idx(wheel_time, start_window)
    end_idx = np.array(start_idx + int(WINDOW_LEN / T_BIN), dtype='int64')
    # Getting the wheel position for each window, normalize to first value of each window
    trials_df['wheel_position'] = [wheel_position[start_idx[w]: end_idx[w]] - wheel_position[start_idx[w]]
                                   for w in range(len(start_idx))]
    # Plotting
    times = np.arange(len(trials_df['wheel_position'].iloc[0])) * T_BIN + WINDOW_LAG
    for side, label, color in zip([-1, 1], ['right', 'left'], ['darkred', '#1f77b4']):
        side_df = trials_df[trials_df['choice'] == side]
        for idx in side_df.index:
            plt.plot(times, side_df.loc[idx, 'wheel_position'], c=color, alpha=0.5, linewidth=0.05)
        plt.plot(times, side_df['wheel_position'].mean(), c=color, linewidth=2, label=f'{label} turn')

    plt.axvline(x=0, linestyle='--', c='k', label='stimOn')
    plt.axhline(y=-0.26, linestyle='--', c='g', label='reward')
    plt.axhline(y=0.26, linestyle='--', c='g', label='reward')
    plt.ylim([-0.27, 0.27])
    plt.xlabel('time [sec]')
    plt.ylabel('wheel position diff to first value [rad]')
    plt.legend(loc='center right')
    plt.title('Wheel position trial avg\n(and individual trials)')
    plt.tight_layout()

    return plt.gca()


def _bin_window_licks(lick_times, trials_df):
    """
    Helper function to bin and window the lick times and get them into trials df for plotting

    :param lick_times: np.array, timestamps of lick events
    :param trials_df: pd.DataFrame, with column 'feedback_times' (time of feedback for each trial)
    :returns: pd.DataFrame with binned, windowed lick times for plotting
    """
    # Bin the licks
    lick_bins, bin_times, _ = bincount2D(lick_times, np.ones(len(lick_times)), T_BIN)
    lick_bins = np.squeeze(lick_bins)
    start_window, end_window = plt_window(trials_df['feedback_times'])
    # Translating the time window into an index window
    try:
        start_idx = insert_idx(bin_times, start_window)
    except ValueError:
        logger.error('Lick time stamps are outside of the trials windows')
        raise
    end_idx = np.array(start_idx + int(WINDOW_LEN / T_BIN), dtype='int64')
    # Get the binned licks for each window
    trials_df['lick_bins'] = [lick_bins[start_idx[i]:end_idx[i]] for i in range(len(start_idx))]
    # Remove windows that the exceed bins
    trials_df['end_idx'] = end_idx
    trials_df = trials_df[trials_df['end_idx'] <= len(lick_bins)]
    return trials_df


def plot_lick_hist(lick_times, trials_df):
    """
    Plots histogramm of lick events aligned to feedback time, separate for correct and incorrect trials

    :param lick_times: np.array, timestamps of lick events
    :param trials_df: pd.DataFrame, with column 'feedback_times' (time of feedback for each trial) and
                      'feedbackType' (1 for correct, -1 for incorrect trials)
    :returns: matplotlib axis
    """
    licks_df = _bin_window_licks(lick_times, trials_df)
    # Plot
    times = np.arange(len(licks_df['lick_bins'].iloc[0])) * T_BIN + WINDOW_LAG
    correct = licks_df[licks_df['feedbackType'] == 1]['lick_bins']
    incorrect = licks_df[licks_df['feedbackType'] == -1]['lick_bins']
    plt.plot(times, pd.DataFrame.from_dict(dict(zip(correct.index, correct.values))).mean(axis=1),
             c='k', label='correct trial')
    plt.plot(times, pd.DataFrame.from_dict(dict(zip(correct.index, incorrect.values))).mean(axis=1),
             c='gray', label='incorrect trial')
    plt.axvline(x=0, label='feedback', linestyle='--', c='purple')
    plt.title('Lick events trial avg')
    plt.xlabel('time [sec]')
    plt.ylabel('lick events [a.u.]')
    plt.legend(loc='lower right')
    return plt.gca()


def plot_lick_raster(lick_times, trials_df):
    """
    Plots lick raster for correct trials

    :param lick_times: np.array, timestamps of lick events
    :param trials_df: pd.DataFrame, with column 'feedback_times' (time of feedback for each trial) and
                      feedbackType (1 for correct, -1 for incorrect trials)
    :returns: matplotlib.axis
    """
    licks_df = _bin_window_licks(lick_times, trials_df)
    plt.imshow(list(licks_df[licks_df['feedbackType'] == 1]['lick_bins']), aspect='auto',
               extent=[-0.5, 1.5, len(licks_df['lick_bins'].iloc[0]), 0], cmap='gray_r')
    plt.xticks([-0.5, 0, 0.5, 1, 1.5])
    plt.ylabel('trials')
    plt.xlabel('time [sec]')
    plt.axvline(x=0, label='feedback', linestyle='--', c='purple')
    plt.title('Lick events per correct trial')
    plt.tight_layout()
    return plt.gca()


def plot_motion_energy_hist(camera_dict, trials_df):
    """
    Plots mean motion energy of given cameras, aligned to stimulus onset.

    :param camera_dict: dict, one key for each camera to be plotted (e.g. 'left'), value is another dict with items
                        'motion_energy' (np.array, motion energy calculated from this camera) and
                        'times' (np.array, camera timestamps)
    :param trials_df: pd.DataFrame, with column 'stimOn_times' (time of stimulus onset for each trial)
    :returns: matplotlib.axis
    """
    colors = {'left': '#bd7a98',
              'right': '#2b6f39',
              'body': '#035382'}

    start_window, end_window = plt_window(trials_df['stimOn_times'])
    missing_data = []
    for cam in camera_dict.keys():
        if (camera_dict[cam]['motion_energy'] is not None and len(camera_dict[cam]['motion_energy']) > 0
                and camera_dict[cam]['times'] is not None and len(camera_dict[cam]['times']) > 0):
            try:
                motion_energy = zscore(camera_dict[cam]['motion_energy'], nan_policy='omit')
                try:
                    start_idx = insert_idx(camera_dict[cam]['times'], start_window)
                    end_idx = np.array(start_idx + int(WINDOW_LEN * SAMPLING[cam]), dtype='int64')
                    me_all = [motion_energy[start_idx[i]:end_idx[i]] for i in range(len(start_idx))]
                    me_all = [m for m in me_all if len(m) > 0]
                    times = np.arange(len(me_all[0])) / SAMPLING[cam] + WINDOW_LAG
                    me_mean = np.mean(me_all, axis=0)
                    me_std = np.std(me_all, axis=0) / np.sqrt(len(me_all))
                    plt.plot(times, me_mean, label=f'{cam} cam', color=colors[cam], linewidth=2)
                    plt.fill_between(times, me_mean + me_std, me_mean - me_std, color=colors[cam], alpha=0.2)
                except ValueError:
                    logger.error(f"{cam}Camera camera.times are outside of the trial windows")
                    missing_data.append(cam)
            except AttributeError:
                logger.warning(f"Cannot load motion energy and/or times data for {cam} camera")
                missing_data.append(cam)
        else:
            logger.warning(f"Data missing or empty for motion energy and/or times data for {cam} camera")
            missing_data.append(cam)

    plt.xticks([-0.5, 0, 0.5, 1, 1.5])
    plt.ylabel('z-scored motion energy [a.u.]')
    plt.xlabel('time [sec]')
    plt.axvline(x=0, label='stimOn', linestyle='--', c='k')
    plt.legend(loc='lower right')
    plt.title('Motion Energy trial avg\n(+/- std)')
    if len(missing_data) > 0:
        ax = plt.gca()
        ax.text(.95, .35, f"Data incomplete for\n{' and '.join(missing_data)} camera", color='r', fontsize=10,
                fontweight='bold', horizontalalignment='right', verticalalignment='center', transform=ax.transAxes)
    return plt.gca()


def plot_speed_hist(dlc_df, cam_times, trials_df, feature='paw_r', cam='left', legend=True):
    """
    Plots speed histogram of a given dlc feature, aligned to stimulus onset, separate for correct and incorrect trials

    :param dlc_df: pd.Dataframe, dlc traces with _x, _y and _likelihood info for each trace
    :param cam_times: np.array, camera timestamps
    :param trials_df: pd.DataFrame, with column 'stimOn_times' (time of stimulus onset for each trial)
    :param feature: str, feature with trace in dlc_df for which to plot speed hist, default is 'paw_r'
    :param cam: str, camera to use ('body', 'left', 'right') default is 'left'
    :param legend: bool, whether to add legend to the plot, default is True
    :returns: matplotlib.axis
    """
    # Threshold the dlc traces
    dlc_df = likelihood_threshold(dlc_df)
    # For pre-GPIO sessions, remove the first few timestamps to match the number of frames
    cam_times = cam_times[-len(dlc_df):]
    if len(cam_times) != len(dlc_df):
        raise ValueError("Camera times length and DLC length are inconsistent")
    # Get speeds
    speeds = get_speed(dlc_df, cam_times, camera=cam, feature=feature)
    # Windows aligned to align_to
    start_window, end_window = plt_window(trials_df['stimOn_times'])
    start_idx = insert_idx(cam_times, start_window)
    end_idx = np.array(start_idx + int(WINDOW_LEN * SAMPLING[cam]), dtype='int64')
    # Add speeds to trials_df
    trials_df[f'speed_{feature}'] = [speeds[start_idx[i]:end_idx[i]] for i in range(len(start_idx))]
    # Plot
    times = np.arange(len(trials_df[f'speed_{feature}'].iloc[0])) / SAMPLING[cam] + WINDOW_LAG
    # Need to expand the series of lists into a dataframe first, for the nan skipping to work
    correct = trials_df[trials_df['feedbackType'] == 1][f'speed_{feature}']
    incorrect = trials_df[trials_df['feedbackType'] == -1][f'speed_{feature}']
    plt.plot(times, pd.DataFrame.from_dict(dict(zip(correct.index, correct.values))).mean(axis=1),
             c='k', label='correct trial')
    plt.plot(times, pd.DataFrame.from_dict(dict(zip(incorrect.index, incorrect.values))).mean(axis=1),
             c='gray', label='incorrect trial')
    plt.axvline(x=0, label='stimOn', linestyle='--', c='r')
    plt.title(f'{feature.capitalize()} speed trial avg\n({cam.upper()} cam)')
    plt.xticks([-0.5, 0, 0.5, 1, 1.5])
    plt.xlabel('time [sec]')
    plt.ylabel('speed [px/sec]')
    if legend:
        plt.legend()

    return plt.gca()


def plot_pupil_diameter_hist(pupil_diameter, cam_times, trials_df, cam='left'):
    """
    Plots histogram of pupil diameter aligned to simulus onset and feedback time.

    :param pupil_diameter: np.array, (smoothed) pupil diameter estimate
    :param cam_times: np.array, camera timestamps
    :param trials_df: pd.DataFrame, with column 'stimOn_times' (time of stimulus onset for each trial) and
                      feedback_times (time of feedback for each trial)
    :param cam: str, camera to use ('body', 'left', 'right') default is 'left'
    :returns: matplotlib.axis
    """
    for align_to, color in zip(['stimOn_times', 'feedback_times'], ['red', 'purple']):
        start_window, end_window = plt_window(trials_df[align_to])
        start_idx = insert_idx(cam_times, start_window)
        end_idx = np.array(start_idx + int(WINDOW_LEN * SAMPLING[cam]), dtype='int64')
        # Per trial norm
        pupil_all = [zscore(list(pupil_diameter[start_idx[i]:end_idx[i]])) for i in range(len(start_idx))]
        pupil_all_norm = [trial - trial[0] for trial in pupil_all]

        pupil_mean = np.nanmean(pupil_all_norm, axis=0)
        pupil_std = np.nanstd(pupil_all_norm, axis=0) / np.sqrt(len(pupil_all_norm))
        times = np.arange(len(pupil_all_norm[0])) / SAMPLING[cam] + WINDOW_LAG

        plt.plot(times, pupil_mean, label=align_to.split("_")[0], color=color)
        plt.fill_between(times, pupil_mean + pupil_std, pupil_mean - pupil_std, color=color, alpha=0.5)
    plt.axvline(x=0, linestyle='--', c='k')
    plt.title(f'Pupil diameter trial avg\n({cam.upper()} cam)')
    plt.xlabel('time [sec]')
    plt.xticks([-0.5, 0, 0.5, 1, 1.5])
    plt.ylabel('z-scored smoothed pupil diameter [px]')
    plt.legend(loc='lower right', title='aligned to')
