# Allen Institute Software License - This software license is the 2-clause BSD
# license plus a third clause that prohibits redistribution for commercial
# purposes without further permission.
#
# Copyright 2017. Allen Institute. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Redistributions for commercial purposes are not permitted without the
# Allen Institute's written permission.
# For purposes of this license, commercial purposes is the incorporation of the
# Allen Institute's software into anything for which you will charge fees or
# other compensation. Contact terms@alleninstitute.org for commercial licensing
# opportunities.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
import numpy as np
import pandas as pd
import os
import logging


def plot_drifting_grating_traces(dg, save_dir):
    '''saves figures with a Ori X TF grid of mean resposes'''
    logging.info("Plotting Ori and TF mean response for all cells")

    blank = dg.sweep_response[dg.stim_table.temporal_frequency == 0]
    for nc in range(dg.numbercells):
        if np.mod(nc, 20) == 0:
            logging.info("Cell #%s", str(nc))
        xtime = np.arange(-1 * dg.interlength / dg.acquisition_rate, (dg.sweeplength +
                                                                      dg.interlength) / dg.acquisition_rate, 1 / dg.acquisition_rate)
        plt.figure(nc, figsize=(20, 16))
        vmax = 0
        vmin = 0
        try:
            blank_p = blank[str(nc)].mean() + \
                (blank[str(nc)].std() / len(blank[str(nc)]))
            blank_n = blank[str(nc)].mean() - \
                (blank[str(nc)].std() / len(blank[str(nc)]))
        except:
            blank_p = blank.iloc[:, nc].apply(
                np.mean) + (blank.iloc[:, nc].apply(np.std) / blank.iloc[:, nc].apply(len))
            blank_n = blank.iloc[:, nc].apply(
                np.mean) - (blank.iloc[:, nc].apply(np.std) / blank.iloc[:, nc].apply(len))
        for ori in dg.orivals:
            ori_pt = np.where(dg.orivals == ori)[0][0]
            for tf in dg.tfvals[1:]:
                tf_pt = np.where(dg.tfvals == tf)[0][0]
                sp_pt = (5 * ori_pt) + tf_pt
                subset_response = dg.sweep_response[
                    (dg.stim_table.temporal_frequency == tf) & (dg.stim_table.orientation == ori)]
                try:
                    subset_response_p = subset_response[str(nc)].mean(
                    ) + (subset_response[str(nc)][:-1].std() / len(subset_response[str(nc)]))
                    subset_response_n = subset_response[str(nc)].mean(
                    ) - (subset_response[str(nc)][:-1].std() / len(subset_response[str(nc)]))
                except:
                    subset_response_p = subset_response.iloc[:, nc].apply(
                        np.mean) + (subset_response.iloc[:, nc].apply(np.std) / subset_response.iloc[:, nc].apply(len))
                    subset_response_n = subset_response.iloc[:, nc].apply(
                        np.mean) - (subset_response.iloc[:, nc].apply(np.std) / subset_response.iloc[:, nc].apply(len))
                ax = plt.subplot(8, 5, sp_pt)
                while len(xtime) > len(subset_response[str(nc)].mean()):
                    xtime = np.delete(xtime, -1)
                try:
                    ax.fill_between(xtime, subset_response_p,
                                    subset_response_n, color='b', alpha=0.5)
                except:
                    pass
                try:
                    ax.fill_between(xtime, blank_p, blank_n,
                                    color='k', alpha=0.5)
                except:
                    pass
                try:
                    ax.plot(xtime, subset_response[
                            str(nc)].mean(), color='b', lw=2)
                except:
                    pass
                ax.plot(xtime, subset_response[
                        str(nc)].mean(), color='b', lw=2)
                # TODO: remove the [:119]  and [:-1] and the try/except
                ax.plot(xtime, blank[str(nc)].mean(), color='k', lw=2)
                ax.axvspan(0, dg.sweeplength / dg.acquisition_rate,
                           ymin=0, ymax=1, facecolor='gray', alpha=0.3)
                ax.set_xlim(-1, 3)
                ax.set_xticks(range(-1, 4))
                ax.yaxis.set_major_locator(MaxNLocator(4))
                vmax = np.where(np.amax(subset_response_p) >
                                vmax, np.amax(subset_response_p), vmax)
                vmin = np.where(np.amin(subset_response_n) <
                                vmin, np.amin(subset_response_n), vmin)

                if ori_pt < 7:
                    ax.set_xticks([])
                else:
                    ax.set_xlabel("Time (s)", fontsize=20)
                if tf_pt > 1:
                    ax.set_yticks([])
                else:
                    ax.set_ylabel(str(dg.orivals[ori_pt]), fontsize=24)
                if ori_pt == 0:
                    ax.set_title(str(dg.tfvals[tf_pt]), fontsize=24)

        for i in range(1, sp_pt + 1):
            ax = plt.subplot(8, 5, i)
            ax.set_ylim(vmin, vmax)
        plt.tick_params(labelsize=16)
        plt.tight_layout()
        plt.suptitle("Cell " + str(nc + 1), fontsize=20)
        plt.subplots_adjust(top=0.9)
        filename = 'Traces DG Cell_' + str(nc + 1) + '.png'
        fullfilename = os.path.join(save_dir, filename)
        plt.savefig(fullfilename)
        plt.close()


def plot_ns_traces(nsa, save_dir):
    logging.info("Plotting Natural Scene traces for each cell")
    xtime = np.arange(-1 * nsa.interlength / nsa.acquisition_rate, (nsa.sweeplength +
                                                                    nsa.interlength) / nsa.acquisition_rate, 1 / nsa.acquisition_rate)
    blank = nsa.sweep_response[nsa.stim_table.frame == -1]
    for nc in range(nsa.numbercells):
        if np.mod(nc, 20) == 0:
            logging.info("Cell #%s", str(nc))
        vmax = 0
        vmin = 0
        blank_p = blank[str(nc)].mean() + \
            (blank[str(nc)].std() / len(blank[str(nc)]))
        blank_n = blank[str(nc)].mean() - \
            (blank[str(nc)].std() / len(blank[str(nc)]))
        plt.figure(nc, figsize=(30, 25))
        for ns in range(nsa.number_scenes - 1):
            subset_response = nsa.sweep_response[nsa.stim_table.frame == ns]
            subset_response_p = subset_response[str(nc)].mean(
            ) + (subset_response[str(nc)][:].std() / len(subset_response[str(nc)]))
            subset_response_n = subset_response[str(nc)].mean(
            ) - (subset_response[str(nc)][:].std() / len(subset_response[str(nc)]))
            ax = plt.subplot(10, 12, ns + 1)
            try:
                ax.fill_between(xtime, subset_response_p,
                                subset_response_n, color='b', alpha=0.5)
            except:
                xtime = xtime[:-1]
                ax.fill_between(xtime, subset_response_p,
                                subset_response_n, color='b', alpha=0.5)
            ax.fill_between(xtime, blank_p, blank_n, color='k', alpha=0.5)
            ax.plot(xtime, subset_response[str(nc)].mean(), color='b', lw=2)
            ax.plot(xtime, blank[str(nc)].mean(), color='k', lw=2)
            ax.axvspan(0, nsa.sweeplength / nsa.acquisition_rate,
                       ymin=0, ymax=1, facecolor='gray', alpha=0.3)
            ax.yaxis.set_major_locator(MaxNLocator(4))
            vmax = np.where(np.amax(subset_response_p) > vmax,
                            np.amax(subset_response_p), vmax)
            vmin = np.where(np.amin(subset_response_n) < vmin,
                            np.amin(subset_response_n), vmin)
            if ns < 108:
                ax.set_xticks([])
            if np.mod(ns, 12):
                ax.set_yticks([])
        for i in range(1, nsa.number_scenes):
            ax = plt.subplot(10, 12, i)
            ax.set_ylim(vmin, vmax)
        plt.tight_layout()
        plt.suptitle("Cell " + str(nc + 1), fontsize=20)
        plt.subplots_adjust(top=0.9)
        filename = 'NS Traces Cell_' + str(nc + 1) + '.png'
        fullfilename = os.path.join(save_dir, filename)
        plt.savefig(fullfilename)
        plt.close()


def plot_sg_traces(sg, save_dir):
    logging.info("Plotting Static Grating traces for each cell")
    xtime = np.arange(-1 * sg.interlength / sg.acquisition_rate, (sg.sweeplength +
                                                                  sg.interlength) / sg.acquisition_rate, 1 / sg.acquisition_rate)
    blank = sg.sweep_response[sg.stim_table.spatial_frequency == 0]
    for nc in range(sg.numbercells):
        if np.mod(nc, 20) == 0:
            logging.info("Cell #%s", str(nc))
        vmax = 0
        vmin = 0
        blank_p = blank[str(nc)].mean() + \
            (blank[str(nc)].std() / len(blank[str(nc)]))
        blank_n = blank[str(nc)].mean() - \
            (blank[str(nc)].std() / len(blank[str(nc)]))
        while len(xtime) > len(blank_p):
            xtime = np.delete(xtime, -1)
        plt.figure(nc, figsize=(30, 30))
        ph_dict = {0: 0, 0.25: 6, 0.5: 77, 0.75: 83}
        for ori in sg.orivals:
            ori_pt = np.where(sg.orivals == ori)[0][0]
            for sf in sg.sfvals[1:]:
                sf_pt = np.where(sg.sfvals == sf)[0][0]
                for phase in sg.phasevals:
                    ph_pt = ph_dict[phase]
                    subplotnum = sf_pt + (ori_pt * 11) + ph_pt
                    subset_response = sg.sweep_response[(sg.stim_table.spatial_frequency == sf) & (
                        sg.stim_table.orientation == ori) & (sg.stim_table.phase == phase)]
                    subset_response_p = subset_response[str(nc)].mean(
                    ) + (subset_response[str(nc)][:].std() / len(subset_response[str(nc)]))
                    subset_response_n = subset_response[str(nc)].mean(
                    ) - (subset_response[str(nc)][:].std() / len(subset_response[str(nc)]))
                    ax = plt.subplot(13, 11, subplotnum)
                    ax.fill_between(xtime, subset_response_p,
                                    subset_response_n, color='b', alpha=0.5)
                    ax.fill_between(xtime, blank_p, blank_n,
                                    color='k', alpha=0.5)
                    ax.plot(xtime, subset_response[
                            str(nc)].mean(), color='b', lw=2)
                    ax.plot(xtime, blank[str(nc)].mean(), color='k', lw=2)
                    ax.axvspan(0, sg.sweeplength / sg.acquisition_rate,
                               ymin=0, ymax=1, facecolor='gray', alpha=0.3)
                    ax.yaxis.set_major_locator(MaxNLocator(4))
                    vmax = np.where(np.amax(subset_response_p)
                                    > vmax, np.amax(subset_response_p), vmax)
                    vmin = np.where(np.amin(subset_response_n)
                                    < vmin, np.amin(subset_response_n), vmin)
                    if np.mod(subplotnum, 11) != 1:
                        ax.set_yticks([])
                    else:
                        ax.set_ylabel(ori, fontsize=20)
                    if subplotnum < 133:
                        ax.set_xticks([])
                    if subplotnum < 12:
                        ax.set_title(sf, fontsize=20)
                    if subplotnum == 3:
                        ax.set_title("Phase 0.0", fontsize=20)
                    if subplotnum == 9:
                        ax.set_title("Phase 0.25", fontsize=20)
                    if subplotnum == 80:
                        ax.set_title("Phase 0.5", fontsize=20)
                    if subplotnum == 86:
                        ax.set_title("Phase 0.75", fontsize=20)
        for i in range(1, 144):
            ax = plt.subplot(13, 11, i)
            ax.set_ylim(vmin, vmax)
        plt.tight_layout()
        plt.suptitle("Cell " + str(nc + 1), fontsize=20)
        plt.subplots_adjust(top=0.9)
        filename = 'SG Traces Cell_' + str(nc + 1) + '.png'
        fullfilename = os.path.join(save_dir, filename)
        plt.savefig(fullfilename)
        plt.close()


def plot_lsn_traces(lsn, save_dir, suffix=''):
    logging.info("Plotting LSN traces for all cells")
    xtime = np.arange(-lsn.interlength / lsn.acquisition_rate,
                      (lsn.interlength + lsn.sweeplength) / lsn.acquisition_rate,
                      1.0 / lsn.acquisition_rate)

    for nc in range(lsn.numbercells):
        if np.mod(nc, 20) == 0:
            logging.info("Cell #%s", str(nc))

        plt.figure(nc, figsize=(24, 20))
        vmax = 0
        vmin = 0
        one_cell = lsn.sweep_response[str(nc)]

        for yp in range(16):
            for xp in range(28):
                sp_pt = (yp * 28) + xp + 1
                on_frame = np.where(lsn.LSN[:, yp, xp] == 255)[0]
                off_frame = np.where(lsn.LSN[:, yp, xp] == 0)[0]
                subset_on = one_cell[lsn.stim_table.frame.isin(on_frame)]
                subset_off = one_cell[lsn.stim_table.frame.isin(off_frame)]

                subset_on_mean = subset_on.mean()
                subset_off_mean = subset_off.mean()

                ax = plt.subplot(16, 28, sp_pt)
                ax.plot(xtime, subset_on_mean, color='r', lw=2)
                ax.plot(xtime, subset_off_mean, color='b', lw=2)
                ax.axvspan(0, lsn.sweeplength / lsn.acquisition_rate,
                           ymin=0, ymax=1, facecolor='gray', alpha=0.3)
                vmax = np.where(np.amax(subset_on_mean) > vmax,
                                np.amax(subset_on_mean), vmax)
                vmax = np.where(np.amax(subset_off_mean) > vmax,
                                np.amax(subset_off_mean), vmax)
                vmin = np.where(np.amin(subset_on_mean) < vmin,
                                np.amin(subset_on_mean), vmin)
                vmin = np.where(np.amin(subset_off_mean) < vmin,
                                np.amin(subset_off_mean), vmin)
                ax.set_xticks([])
                ax.set_yticks([])

        for i in range(1, sp_pt + 1):
            ax = plt.subplot(16, 28, i)
            ax.set_ylim(vmin, vmax)

        plt.tight_layout()
        plt.suptitle("Cell " + str(nc + 1), fontsize=20)
        plt.subplots_adjust(top=0.9)
        filename = 'Traces LSN Cell_' + str(nc + 1) + suffix + '.png'
        fullfilename = os.path.join(save_dir, filename)
        plt.savefig(fullfilename)
        plt.close()


def _plot_3sa(dg, nm1, nm3, save_dir):
    logging.info("Plotting for all cell")
    for nc in range(dg.numbercells):
        if np.mod(nc, 20) == 0:
            logging.info("Cell #%s", str(nc))
        plt.figure(nc, figsize=(20, 20))
        ax1 = plt.subplot2grid((6, 6), (0, 0), colspan=4)  # full trace
        ax2 = plt.subplot2grid((6, 6), (0, 4))  # histogram of F
        ax3 = plt.subplot2grid((6, 6), (1, 0), colspan=4)  # running speed
        ax4 = plt.subplot2grid((6, 6), (1, 4))
        ax5 = plt.subplot2grid((6, 6), (1, 5), sharex=ax4, sharey=ax4)
        ax6 = plt.subplot2grid((6, 6), (2, 0), colspan=2)
        ax7 = plt.subplot2grid((6, 6), (2, 2), colspan=2)
        ax8 = plt.subplot2grid((6, 6), (2, 4), colspan=2)
        ax9 = plt.subplot2grid((6, 6), (3, 0))
        ax10 = plt.subplot2grid((6, 6), (4, 0), colspan=3)
        ax11 = plt.subplot2grid((6, 6), (5, 0), colspan=3)
        ax12 = plt.subplot2grid((6, 6), (4, 3), colspan=3)
        ax13 = plt.subplot2grid((6, 6), (5, 3), colspan=3)

        xtime = np.arange(0, np.size(dg.celltraces, 1), 1.)
        xtime /= dg.acquisition_rate
        dif = np.ediff1d(dg.stim_table.start.values,
                         to_begin=8000, to_end=8000)
        test = np.argwhere(dif > 5000)
        ax1.plot(xtime, dg.celltraces[nc, :])
        for i in range(len(test) - 1):
            ax1.axvspan(xmin=(dg.stim_table.start.iloc[test[i]].values / dg.acquisition_rate), xmax=(
                dg.stim_table.end.iloc[test[i + 1] - 1].values / dg.acquisition_rate), color='gray', alpha=0.3)
        ax1.axvspan(xmin=nm1.stim_table.start.min() / nm1.acquisition_rate, xmax=(
            (nm1.stim_table.start.max() + nm1.sweeplength) / nm1.acquisition_rate), color='red', alpha=0.3)
        dif = np.ediff1d(nm3.stim_table.start.values,
                         to_begin=8000, to_end=8000)
        test = np.argwhere(dif > 5000)
        for i in range(len(test) - 1):
            ax1.axvspan(xmin=(nm3.stim_table.start.iloc[test[i]].values / nm3.acquisition_rate), xmax=(
                (nm3.stim_table.end.iloc[test[i + 1] - 1].values + nm3.sweeplength) / nm3.acquisition_rate), color='blue', alpha=0.3)
        ax1.set_xlabel("Time (s)", fontsize=20)
        ax1.set_ylabel("Fluorescence", fontsize=20)

        ax2.hist(dg.celltraces[nc, :], bins=70)
        ax2.set_yscale('log')
        ax2.set_xlabel("Fluorescence", fontsize=20)
        ax2.set_ylabel("Count", fontsize=20)

        xtime = np.arange(0, np.size(dg.dxcm), 1.)
        xtime /= dg.acquisition_rate
        ax3.plot(xtime, dg.dxcm, color='k')
        ax3.set_xlabel("Time (s)", fontsize=20)
        ax3.set_xlabel("Speed (cm/s)", fontsize=20)

        smax = nm1.binned_cells_sp[nc, np.argmax(nm1.binned_cells_sp[
                                                 nc, :, 0]), 0] + nm1.binned_cells_sp[nc, np.argmax(nm1.binned_cells_sp[nc, :, 0]), 1]
        vmax = nm1.binned_cells_vis[nc, np.argmax(nm1.binned_cells_vis[
                                                  nc, :, 0]), 0] + nm1.binned_cells_vis[nc, np.argmax(nm1.binned_cells_vis[nc, :, 0]), 1]
        rmax = np.where(smax > vmax, smax, vmax)
        smin = nm1.binned_cells_sp[nc, np.argmin(nm1.binned_cells_sp[
                                                 nc, :, 0]), 0] - nm1.binned_cells_sp[nc, np.argmin(nm1.binned_cells_sp[nc, :, 0]), 1]
        vmin = nm1.binned_cells_vis[nc, np.argmin(nm1.binned_cells_vis[
                                                  nc, :, 0]), 0] - nm1.binned_cells_vis[nc, np.argmin(nm1.binned_cells_vis[nc, :, 0]), 1]
        rmin = np.where(smin < vmin, smin, vmin)

        ax4.errorbar(nm1.binned_dx_sp[:, 0], nm1.binned_cells_sp[
                     nc, :, 0], yerr=nm1.binned_cells_sp[nc, :, 1], fmt='.', color='k')
        ax4.set_ylim(rmin, rmax)
        ax4.set_xlabel("Speed (cm/s)", fontsize=20)
        ax4.set_ylabel("DF/F", fontsize=20)
        ax4.set_title("Spontaneous", fontsize=20)

        ax5.errorbar(nm1.binned_dx_vis[:, 0], nm1.binned_cells_vis[
                     nc, :, 0], yerr=nm1.binned_cells_vis[nc, :, 1], fmt='.')
        ax5.set_ylim(rmin, rmax)
        ax5.set_xlabel("Speed (cm/s)", fontsize=20)
        ax5.set_ylabel("DF/F", fontsize=20)
        ax5.set_title("Visual Stimuli", fontsize=20)

        peakori = dg.peak.ori_dg[nc]
        peaktf = dg.peak.tf_dg[nc]
        ax6.errorbar(dg.orivals, dg.response[:, peaktf, nc, 0], yerr=dg.response[
                     :, peaktf, nc, 1], fmt='bo-', lw=2)
        ax6.fill_between(dg.orivals, np.repeat(dg.response[0, 0, nc, 0] + dg.response[0, 0, nc, 1], dg.number_ori), np.repeat(
            dg.response[0, 0, nc, 0] - dg.response[0, 0, nc, 1], dg.number_ori), color='gray', alpha=0.5)
        ax6.axhline(y=dg.response[0, 0, nc, 0], ls='--', color='k')
        ax6.annotate(str(dg.tfvals[peaktf]) + " Hz",
                     xy=(0, 0.9), xycoords='axes fraction', fontsize=14)
        ax6.set_xticks(dg.orivals)
        ax7.set_xlim(-10, 325)
        ax6.set_xlabel("Direction (d)", fontsize=20)
        ax6.set_ylabel("Mean DF/F (%)", fontsize=20)
        ax6.yaxis.set_major_locator(MaxNLocator(6))

        ax7.errorbar(range(5), dg.response[peakori, 1:, nc, 0], yerr=dg.response[
                     peakori, 1:, nc, 1], fmt='bo-', lw=2)
        ax7.fill_between(range(5), np.repeat(dg.response[0, 0, nc, 0] + dg.response[0, 0, nc, 1], 5), np.repeat(
            dg.response[0, 0, nc, 0] - dg.response[0, 0, nc, 1], 5), color='gray', alpha=0.5)
        ax7.axhline(y=dg.response[0, 0, nc, 0], ls='--', color='k', lw=2)
        ax7.annotate(str(dg.orivals[peakori]) + " Deg",
                     xy=(0, 0.9), xycoords='axes fraction', fontsize=14)
        ax7.set_xlim(-0.2, 4.2)
        ax7.set_xticks(range(5))
        ax7.set_xticklabels(dg.tfvals[1:])
        ax7.set_xlabel("Temporal frequency (Hz)", fontsize=20)

        subset = dg.sweep_response[(dg.stim_table.orientation == dg.orivals[peakori]) & (
            dg.stim_table.temporal_frequency == dg.tfvals[peaktf])]
        xtime = np.arange(-1 * dg.interlength / dg.acquisition_rate, (dg.sweeplength +
                                                                      dg.interlength) / dg.acquisition_rate, 1 / dg.acquisition_rate)
        while len(xtime) > len(subset[str(nc)].mean()):
            xtime = np.delete(xtime, -1)
        for index, row in subset.iterrows():
            ax8.plot(xtime, subset[str(nc)][index], lw=2)
        ax8.set_xlim(-1, 3)
        ax8.annotate(str(dg.orivals[peakori]) + " Deg / " + str(
            dg.tfvals[peaktf]) + " Hz", xy=(0, 0.9), xycoords='axes fraction', fontsize=14)
        ax8.set_xlabel("Time (s)", fontsize=20)
        ax8.set_ylabel("DF/F (%)", fontsize=20)
        ax8.axvspan(0, dg.sweeplength / dg.acquisition_rate,
                    ymin=0, ymax=1, facecolor='gray', alpha=0.3)
        ax8.yaxis.set_major_locator(MaxNLocator(6))
        ax8.set_title("Trial responses to prefered ori/tf", fontsize=20)

        im = ax9.imshow(dg.response[:, 1:, nc, 0],
                        cmap='gray', interpolation='none')
        ax9.set_ylabel("Direction (d)", fontsize=20)
        ax9.set_yticks(range(8))
        ax9.set_yticklabels(dg.orivals)
        ax9.set_xlabel("Temporal frequency (Hz)", fontsize=20)
        ax9.set_xticks(range(5))
        ax9.set_xticklabels(dg.tfvals[1:])
        cbar = plt.colorbar(im, ax=ax9)
        cbar.ax.set_ylabel('DF/F (%)', fontsize=8)
        for t in cbar.ax.get_yticklabels():
            t.set_fontsize(8)

        xtime = np.arange(0, nm1.sweeplength /
                          nm1.acquisition_rate, 1 / nm1.acquisition_rate)
        while len(xtime) > len(nm1.sweep_response[str(nc)].mean()):
            xtime = np.delete(xtime, -1)
        for index, row in nm1.sweep_response.iterrows():
            ax10.plot(xtime, nm1.sweep_response[str(nc)][index], lw=2)
        ax10.set_xlabel("Time (s)", fontsize=20)
        ax10.set_ylabel("DF/F", fontsize=20)
        ax10.set_title("Natural Movie 1", fontsize=20, color='red')

        temp = np.empty((len(nm1.stim_table), nm1.sweeplength))
        for i in range(len(nm1.stim_table)):
            temp[i, :] = nm1.sweep_response[str(nc)].iloc[i]
        ax11.imshow(temp, cmap='gray', interpolation='none', aspect=40)
        ax11.set_ylabel("Trials", fontsize=20)
        ax11.set_xticks([])

        xtime = np.arange(0, nm3.sweeplength /
                          nm3.acquisition_rate, 1 / nm3.acquisition_rate)
        while len(xtime) > len(nm3.sweep_response[str(nc)].mean()):
            xtime = np.delete(xtime, -1)
        for index, row in nm3.sweep_response.iterrows():
            ax12.plot(xtime, nm3.sweep_response[str(nc)][index], lw=2)
        ax12.set_xlabel("Time (s)", fontsize=20)
        ax12.set_ylabel("DF/F", fontsize=20)
        ax12.set_title("Natural Movie Long", fontsize=20, color='blue')

        temp = np.empty((len(nm3.stim_table), nm3.sweeplength))
        for i in range(len(nm3.stim_table)):
            temp[i, :] = nm3.sweep_response[str(nc)].iloc[i]
        ax13.imshow(temp, cmap='gray', interpolation='none', aspect=100)
        ax13.set_ylabel("Trials", fontsize=20)
        ax13.set_xticks([])

        plt.tick_params(labelsize=16)
        plt.tight_layout()
        filename = 'Cell_' + str(nc + 1) + '_3SA.png'
        fullfilename = os.path.join(save_dir, filename)
        plt.savefig(fullfilename)
        plt.close()


def _plot_3sc(lsn, nm1, nm2, save_dir, suffix=''):
    logging.info("Plotting for all cells")
    for nc in range(lsn.numbercells):
        if np.mod(nc, 20) == 0:
            logging.info("Cell #%s", str(nc))

        plt.figure(nc, figsize=(20, 20))
        ax1 = plt.subplot2grid((6, 6), (0, 0), colspan=4)  # full trace
        ax2 = plt.subplot2grid((6, 6), (0, 4))  # histogram of F
        ax11 = plt.subplot2grid((6, 6), (1, 0), colspan=4)
        ax12 = plt.subplot2grid((6, 6), (1, 4))
        ax13 = plt.subplot2grid((6, 6), (1, 5), sharex=ax12, sharey=ax12)
        ax3 = plt.subplot2grid((6, 6), (2, 0), colspan=3)  # movie 1
        ax4 = plt.subplot2grid((6, 6), (3, 0), colspan=3)
        ax5 = plt.subplot2grid((6, 6), (2, 3), colspan=3)  # movie 2
        ax6 = plt.subplot2grid((6, 6), (3, 3), colspan=3)
        ax7 = plt.subplot2grid((6, 6), (4, 0), colspan=3)  # receptive fields
        ax8 = plt.subplot2grid((6, 6), (4, 3), colspan=3)
        ax9 = plt.subplot2grid((6, 6), (5, 0), colspan=3)
        ax10 = plt.subplot2grid((6, 6), (5, 3), colspan=3)

        xtime = np.arange(0, np.size(lsn.celltraces, 1), 1.)
        xtime /= lsn.acquisition_rate
        dif = np.ediff1d(lsn.stim_table.start.values,
                         to_begin=8000, to_end=8000)
        test = np.argwhere(dif > 5000)
        ax1.plot(xtime, lsn.celltraces[nc, :])
        for i in range(len(test) - 1):
            ax1.axvspan(xmin=(lsn.stim_table.start.iloc[test[i]].values / lsn.acquisition_rate), xmax=(
                lsn.stim_table.end.iloc[test[i + 1] - 1].values / lsn.acquisition_rate), color='gray', alpha=0.3)
        ax1.axvspan(nm1.stim_table.start.min() / nm1.acquisition_rate, nm1.stim_table.end.max() /
                    nm1.acquisition_rate, ymin=0, ymax=1, color='red', alpha=0.3)
        ax1.axvspan(nm2.stim_table.start.min() / nm2.acquisition_rate, nm2.stim_table.end.max() /
                    nm2.acquisition_rate, ymin=0, ymax=1, color='green', alpha=0.3)
        ax1.set_xlabel("Time (s)", fontsize=20)
        ax1.set_ylabel("Fluorescence", fontsize=20)

        ax2.hist(lsn.celltraces[nc, :], bins=70)
        ax2.set_yscale('log')
        ax2.set_xlabel("Fluorescence", fontsize=20)
        ax2.set_ylabel("Count", fontsize=20)

        xtime = np.arange(0, np.size(lsn.dxcm), 1.)
        xtime /= lsn.acquisition_rate
        ax11.plot(xtime, lsn.dxcm, color='k')
        ax11.set_xlabel("Time (s)", fontsize=20)
        ax11.set_xlabel("Speed (cm/s)", fontsize=20)

        smax = nm1.binned_cells_sp[nc, np.argmax(nm1.binned_cells_sp[
                                                 nc, :, 0]), 0] + nm1.binned_cells_sp[nc, np.argmax(nm1.binned_cells_sp[nc, :, 0]), 1]
        vmax = nm1.binned_cells_vis[nc, np.argmax(nm1.binned_cells_vis[
                                                  nc, :, 0]), 0] + nm1.binned_cells_vis[nc, np.argmax(nm1.binned_cells_vis[nc, :, 0]), 1]
        rmax = np.where(smax > vmax, smax, vmax)
        smin = nm1.binned_cells_sp[nc, np.argmin(nm1.binned_cells_sp[
                                                 nc, :, 0]), 0] - nm1.binned_cells_sp[nc, np.argmin(nm1.binned_cells_sp[nc, :, 0]), 1]
        vmin = nm1.binned_cells_vis[nc, np.argmin(nm1.binned_cells_vis[
                                                  nc, :, 0]), 0] - nm1.binned_cells_vis[nc, np.argmin(nm1.binned_cells_vis[nc, :, 0]), 1]
        rmin = np.where(smin < vmin, smin, vmin)

        ax12.errorbar(nm1.binned_dx_sp[:, 0], nm1.binned_cells_sp[
                      nc, :, 0], yerr=nm1.binned_cells_sp[nc, :, 1], fmt='.', color='k')
        ax12.set_ylim(rmin, rmax)
        ax12.set_xlabel("Speed (cm/s)", fontsize=20)
        ax12.set_ylabel("DF/F", fontsize=20)
        ax12.set_title("Spontaneous", fontsize=20)

        ax13.errorbar(nm1.binned_dx_vis[:, 0], nm1.binned_cells_vis[
                      nc, :, 0], yerr=nm1.binned_cells_vis[nc, :, 1], fmt='.')
        ax13.set_ylim(rmin, rmax)
        ax13.set_xlabel("Speed (cm/s)", fontsize=20)
        ax13.set_ylabel("DF/F", fontsize=20)
        ax13.set_title("Visual Stimuli", fontsize=20)

        xtime = np.arange(0, nm1.sweeplength /
                          nm1.acquisition_rate, 1 / nm1.acquisition_rate)
        for index, row in nm1.sweep_response.iterrows():
            ax3.plot(xtime, nm1.sweep_response[str(nc)][index], lw=2)
        ax3.set_xlabel("Time (s)", fontsize=20)
        ax3.set_ylabel("DF/F", fontsize=20)
        ax3.set_title("Natural Movie 1", fontsize=20, color='red')

        temp = np.empty((len(nm1.stim_table), nm1.sweeplength))
        for i in range(len(nm1.stim_table)):
            temp[i, :] = nm1.sweep_response[str(nc)].iloc[i]
        ax4.imshow(temp, cmap='gray', interpolation='none', aspect=40)
        ax4.set_ylabel("Trials", fontsize=20)
        ax4.set_xticks([])

        xtime = np.arange(0, nm2.sweeplength /
                          nm2.acquisition_rate, 1 / nm2.acquisition_rate)
        for index, row in nm2.sweep_response.iterrows():
            ax5.plot(xtime, nm2.sweep_response[str(nc)][index], lw=2)
        ax5.set_xlabel("Time (s)", fontsize=20)
        ax5.set_ylabel("DF/F", fontsize=20)
        ax5.set_title("Natural Movie 2", fontsize=20, color='green')

        temp = np.empty((len(nm2.stim_table), nm2.sweeplength))
        for i in range(len(nm2.stim_table)):
            temp[i, :] = nm2.sweep_response[str(nc)].iloc[i]
        ax6.imshow(temp, cmap='gray', interpolation='none', aspect=40)
        ax6.set_ylabel("Trials", fontsize=20)
        ax6.set_xticks([])

        vMax = np.where(np.amax(lsn.receptive_field[:, :, nc, 0]) > np.amax(lsn.receptive_field[
                        :, :, nc, 1]), np.amax(lsn.receptive_field[:, :, nc, 0]), np.amax(lsn.receptive_field[:, :, nc, 1]))
        vMin = np.where(np.amin(lsn.receptive_field[:, :, nc, 0]) < np.amin(lsn.receptive_field[
                        :, :, nc, 1]), np.amin(lsn.receptive_field[:, :, nc, 0]), np.amin(lsn.receptive_field[:, :, nc, 1]))

        imon = ax7.imshow(lsn.receptive_field[
                          :, :, nc, 0], cmap='gray', interpolation='None', vmin=vMin, vmax=vMax)
        ax7.set_title("ON", fontsize=20)
        ax7.set_xticks([])
        ax7.set_yticks([])
        cbar = plt.colorbar(imon, ax=ax7, fraction=0.046, pad=0.04)
        cbar.ax.set_ylabel('DF/F (%)', fontsize=10)
        for t in cbar.ax.get_yticklabels():
            t.set_fontsize(8)

        imoff = ax8.imshow(lsn.receptive_field[
                           :, :, nc, 1], cmap='gray', interpolation='None', vmin=vMin, vmax=vMax)
        ax8.set_title("OFF", fontsize=20)
        ax8.set_xticks([])
        ax8.set_yticks([])
        cbar = plt.colorbar(imoff, ax=ax8, fraction=0.046, pad=0.04)
        cbar.ax.set_ylabel('DF/F (%)', fontsize=10)
        for t in cbar.ax.get_yticklabels():
            t.set_fontsize(8)

        zon = (lsn.receptive_field[:, :, nc, 0] - np.mean(lsn.receptive_field[
               :, :, nc, 0])) / np.std(lsn.receptive_field[:, :, nc, 0])
        zon = np.where(abs(zon) > 2, zon, 0)
        Vmax_on = np.where(abs(np.amax(zon)) > abs(
            np.amin(zon)), np.amax(zon), -1 * np.amin(zon))
        zoff = (lsn.receptive_field[:, :, nc, 1] - np.mean(lsn.receptive_field[
                :, :, nc, 1])) / np.std(lsn.receptive_field[:, :, nc, 1])
        zoff = np.where(abs(zoff) > 2, zoff, 0)
        Vmax_off = np.where(abs(np.amax(zoff)) > abs(
            np.amin(zoff)), np.amax(zoff), -1 * np.amin(zoff))
        Vmax = np.where(Vmax_on > Vmax_off, Vmax_on, Vmax_off)
        imzon = ax9.imshow(zon, cmap='RdBu_r',
                           interpolation='none', vmin=-1 * Vmax, vmax=Vmax)
        ax9.set_title("On Z-score", fontsize=20)
        ax9.set_xticks([])
        ax9.set_yticks([])
        cbar = plt.colorbar(imzon, ax=ax9, fraction=0.046, pad=0.04)

        zoff = (lsn.receptive_field[:, :, nc, 1] - np.mean(lsn.receptive_field[
                :, :, nc, 1])) / np.std(lsn.receptive_field[:, :, nc, 1])
        zoff = np.where(abs(zoff) > 2, zoff, 0)
        imzoff = ax10.imshow(
            zoff, cmap='RdBu', interpolation='none', vmin=-1 * Vmax, vmax=Vmax)
        ax10.set_title("Off Z-score", fontsize=20)
        ax10.set_xticks([])
        ax10.set_yticks([])
        cbar = plt.colorbar(imzoff, ax=ax10, fraction=0.046, pad=0.04)

        plt.tick_params(labelsize=16)
        plt.tight_layout()
        filename = 'Cell_' + str(nc + 1) + '_3SC' + suffix + '.png'
        fullfilename = os.path.join(save_dir, filename)
        plt.savefig(fullfilename)
        plt.close()


def _plot_3sb(sg, nm1, ns, save_dir):
    logging.info("Plotting for all cells")
    for nc in range(sg.numbercells):
        if np.mod(nc, 20) == 0:
            logging.info("Cell #%s", str(nc))
        plt.figure(nc, figsize=(20, 24))
        ax1 = plt.subplot2grid((6, 6), (0, 0), colspan=4)  # full trace
        ax2 = plt.subplot2grid((6, 6), (0, 4))  # histogram of F
        ax14 = plt.subplot2grid((6, 6), (1, 4))  # speed tuning
        ax16 = plt.subplot2grid((6, 6), (1, 5), sharex=ax14, sharey=ax14)
        ax15 = plt.subplot2grid((6, 6), (1, 0), colspan=4)
        ax3 = plt.subplot2grid((6, 6), (2, 0), colspan=2)  # Ori tuning
        ax4 = plt.subplot2grid((6, 6), (2, 2), colspan=2)  # sf tuning
        ax13 = plt.subplot2grid((6, 6), (2, 4), colspan=2)  # response at peak
        ax5 = plt.subplot2grid((6, 6), (3, 0))
        ax6 = plt.subplot2grid((6, 6), (3, 1))
        ax7 = plt.subplot2grid((6, 6), (3, 2))
        ax8 = plt.subplot2grid((6, 6), (3, 3))
        ax9 = plt.subplot2grid((6, 6), (4, 0), colspan=3)  # movie
        ax10 = plt.subplot2grid((6, 6), (5, 0), colspan=3)
        ax11 = plt.subplot2grid((6, 6), (4, 3), colspan=2)  # natural scenes
        ax12 = plt.subplot2grid((6, 6), (5, 3), colspan=2)

        xtime = np.arange(0, np.size(sg.celltraces, 1), 1.)
        xtime /= sg.acquisition_rate
        dif = np.ediff1d(sg.stim_table.start.values,
                         to_begin=8000, to_end=8000)
        test = np.argwhere(dif > 5000)
        ax1.plot(xtime, sg.celltraces[nc, :])
        for i in range(len(test) - 1):
            ax1.axvspan(xmin=(sg.stim_table.start.iloc[test[i]].values / sg.acquisition_rate), xmax=(
                sg.stim_table.end.iloc[test[i + 1] - 1].values / sg.acquisition_rate), color='gray', alpha=0.3)
        ax1.axvspan(nm1.stim_table.start.min() / nm1.acquisition_rate, nm1.stim_table.end.max() /
                    nm1.acquisition_rate, ymin=0, ymax=1, color='red', alpha=0.3)
        dif = np.ediff1d(ns.stim_table.start.values,
                         to_begin=8000, to_end=8000)
        test = np.argwhere(dif > 5000)
        for i in range(len(test) - 1):
            ax1.axvspan(xmin=(ns.stim_table.start.iloc[test[i]].values / ns.acquisition_rate), xmax=(
                ns.stim_table.end.iloc[test[i + 1] - 1].values / ns.acquisition_rate), color='blue', alpha=0.3)
        ax1.set_xlabel("Time (s)", fontsize=20)
        ax1.set_ylabel("Fluorescence", fontsize=20)

        ax2.hist(sg.celltraces[nc, :], bins=70)
        ax2.set_yscale('log')
        ax2.set_xlabel("Fluorescence", fontsize=20)
        ax2.set_ylabel("Count", fontsize=20)

        xtime = np.arange(0, np.size(sg.dxcm), 1.)
        xtime /= sg.acquisition_rate
        ax15.plot(xtime, sg.dxcm, color='k')
        ax15.set_xlabel("Time (s)", fontsize=20)
        ax15.set_xlabel("Speed (cm/s)", fontsize=20)

        peakori = sg.peak.ori_sg[nc]
        peaksf = sg.peak.sf_sg[nc]
        ax3.errorbar(sg.orivals, sg.response[:, peaksf, 0, nc, 0], yerr=sg.response[
                     :, peaksf, 0, nc, 1], color='blue', fmt='o-', lw=2)
        ax3.errorbar(sg.orivals, sg.response[:, peaksf, 1, nc, 0], yerr=sg.response[
                     :, peaksf, 1, nc, 1], color='cornflowerblue', fmt='o-', lw=2)
        ax3.errorbar(sg.orivals, sg.response[:, peaksf, 2, nc, 0], yerr=sg.response[
                     :, peaksf, 2, nc, 1], color='steelblue', fmt='o-', lw=2)
        ax3.errorbar(sg.orivals, sg.response[:, peaksf, 3, nc, 0], yerr=sg.response[
                     :, peaksf, 3, nc, 1], color='lightskyblue', fmt='o-', lw=2)
        ax3.fill_between(sg.orivals, np.repeat(sg.response[0, 0, 0, nc, 0] + sg.response[0, 0, 0, nc, 1], sg.number_ori), np.repeat(
            sg.response[0, 0, 0, nc, 0] - sg.response[0, 0, 0, nc, 1], sg.number_ori), color='gray', alpha=0.5)
        ax3.axhline(y=sg.response[0, 0, 0, nc, 0], ls='--', color='k', lw=2)
        ax3.set_xlim(-10, 160)
        ax3.set_xticks(sg.orivals)
        ax3.set_xlabel("Orientation (d)", fontsize=20)
        ax3.set_ylabel("DF/F (%)", fontsize=20)

        ax4.errorbar(range(5), sg.response[peakori, 1:, 0, nc, 0], yerr=sg.response[
                     peakori, 1:, 0, nc, 1], color='blue', fmt='o-', lw=2)
        ax4.errorbar(range(5), sg.response[peakori, 1:, 1, nc, 0], yerr=sg.response[
                     peakori, 1:, 1, nc, 1], color='cornflowerblue', fmt='o-', lw=2)
        ax4.errorbar(range(5), sg.response[peakori, 1:, 2, nc, 0], yerr=sg.response[
                     peakori, 1:, 2, nc, 1], color='steelblue', fmt='o-', lw=2)
        ax4.errorbar(range(5), sg.response[peakori, 1:, 3, nc, 0], yerr=sg.response[
                     peakori, 1:, 3, nc, 1], color='lightskyblue', fmt='o-', lw=2)
        ax4.fill_between(range(5), np.repeat(sg.response[0, 0, 0, nc, 0] + sg.response[0, 0, 0, nc, 1], 5), np.repeat(
            sg.response[0, 0, 0, nc, 0] - sg.response[0, 0, 0, nc, 1], 5), color='gray', alpha=0.5)
        ax4.axhline(y=sg.response[0, 0, 0, nc, 0], ls='--', color='k', lw=2)
        ax4.set_xlim(-0.2, 4.2)
        ax4.set_xticks(range(5))
        ax4.set_xticklabels(sg.sfvals[1:])
        ax4.set_xlabel("Spatial frequency (cpd)", fontsize=20)

        xtime = np.arange(-1 * sg.interlength / sg.acquisition_rate, (sg.sweeplength +
                                                                      sg.interlength) / sg.acquisition_rate, 1 / sg.acquisition_rate)
        peakori = sg.peak.ori_sg[nc]
        peaksf = sg.peak.sf_sg[nc]
        peakphase = sg.peak.phase_sg[nc]
        subset = sg.sweep_response[(sg.stim_table.orientation == sg.orivals[peakori]) & (
            sg.stim_table.spatial_frequency == sg.sfvals[peaksf]) & (sg.stim_table.phase == sg.phasevals[peakphase])]
        subset_p = subset[str(nc)].mean(
        ) + (subset[str(nc)].std() / np.sqrt(len(subset[str(nc)])))
        subset_n = subset[str(nc)].mean(
        ) - (subset[str(nc)].std() / np.sqrt(len(subset[str(nc)])))
        try:
            ax13.fill_between(xtime, subset_p, subset_n, color='b', alpha=0.5)
        except:
            xtime = xtime[:-1]
            ax13.fill_between(xtime, subset_p, subset_n, color='b', alpha=0.5)
        blank = sg.sweep_response[(sg.stim_table.orientation == 0) & (
            sg.stim_table.spatial_frequency == 0) & (sg.stim_table.phase == 0)]
        blank_p = blank[str(nc)].mean() + \
            (blank[str(nc)].std() / np.sqrt(len(blank[str(nc)])))
        blank_n = blank[str(nc)].mean() - \
            (blank[str(nc)].std() / np.sqrt(len(blank[str(nc)])))
        ax13.fill_between(xtime, blank_p, blank_n, color='gray', alpha=0.5)
        ax13.plot(xtime, subset[str(nc)].mean(), color='b', lw=2)
        ax13.plot(xtime, blank[str(nc)].mean(), color='k', lw=2)
        ax13.axvspan(0, sg.sweeplength / sg.acquisition_rate,
                     ymin=0, ymax=1, facecolor='gray', alpha=0.3)
        ax13.yaxis.set_major_locator(MaxNLocator(4))
        ax13.set_xlabel("Time (s)", fontsize=20)
        ax13.set_ylabel("DF/F (%)", fontsize=20)

        Vmax = sg.response[:, 1:, :, nc, 0].max()
        ax5.imshow(sg.response[:, 1:, 0, nc, 0], cmap='gray',
                   interpolation='none', vmin=0, vmax=Vmax)
        ax5.set_ylabel("Orientation (d)", fontsize=20)
        ax5.set_yticks(range(6))
        ax5.set_yticklabels(sg.orivals)
        ax5.set_xlabel("Spatial frequency (cpd)", fontsize=20)
        ax5.set_xticks(range(5))
        ax5.set_xticklabels(sg.sfvals[1:])
        ax5.set_title("Phase 0.0", color='blue', fontsize=20)

        ax6.imshow(sg.response[:, 1:, 1, nc, 0], cmap='gray',
                   interpolation='none', vmin=0, vmax=Vmax)
        ax6.set_xlabel("Spatial frequency (cpd)", fontsize=20)
        ax6.set_xticks(range(5))
        ax6.set_xticklabels(sg.sfvals[1:])
        ax6.set_yticks(range(6))
        ax6.set_yticklabels(sg.orivals)
        ax6.set_title("Phase 0.25", color='cornflowerblue', fontsize=20)

        ax7.imshow(sg.response[:, 1:, 2, nc, 0], cmap='gray',
                   interpolation='none', vmin=0, vmax=Vmax)
        ax7.set_xlabel("Spatial frequency (cpd)", fontsize=20)
        ax7.set_xticks(range(5))
        ax7.set_xticklabels(sg.sfvals[1:])
        ax7.set_yticks(range(6))
        ax7.set_yticklabels(sg.orivals)
        ax7.set_title("Phase 0.5", color='steelblue', fontsize=20)

        ax8.imshow(sg.response[:, 1:, 3, nc, 0], cmap='gray',
                   interpolation='none', vmin=0, vmax=Vmax)
        ax8.set_xlabel("Spatial frequency (cpd)", fontsize=20)
        ax8.set_xticks(range(5))
        ax8.set_xticklabels(sg.sfvals[1:])
        ax8.set_yticks(range(6))
        ax8.set_yticklabels(sg.orivals)
        ax8.set_title("Phase 0.75", color='lightskyblue', fontsize=20)

        xtime = np.arange(0, nm1.sweeplength /
                          nm1.acquisition_rate, 1 / nm1.acquisition_rate)
        while len(xtime) > nm1.sweeplength:
            xtime = np.delete(xtime, -1)
        for index, row in nm1.sweep_response.iterrows():
            ax9.plot(xtime, nm1.sweep_response[str(nc)][index], lw=2)
        ax9.set_xlabel("Time (s)", fontsize=20)
        ax9.set_ylabel("DF/F", fontsize=20)
        ax9.set_title("Natural Movie 1", fontsize=20, color='red')

        temp = np.empty((len(nm1.stim_table), nm1.sweeplength))
        for i in range(len(nm1.stim_table)):
            temp[i, :] = nm1.sweep_response[str(nc)].iloc[i]
        ax10.imshow(temp, cmap='gray', interpolation='none', aspect=40)
        ax10.set_ylabel("Trials", fontsize=20)
        ax10.set_xticks([])

        temp = np.copy(ns.response[1:, nc, :2])
        scene_response = pd.DataFrame(temp, columns=('response', 'error'))
        scene_response = scene_response.sort(
            columns='response', ascending=False)
        ax11.errorbar(range(ns.number_scenes - 1), scene_response.response,
                      yerr=scene_response.error, fmt='o', color='k')
        ax11.fill_between(range(ns.number_scenes - 1), np.repeat(ns.response[0, nc, 0] + ns.response[0, nc, 1], ns.number_scenes - 1), np.repeat(
            ns.response[0, nc, 0] - ns.response[0, nc, 1], ns.number_scenes - 1), color='gray', alpha=0.3)
        ax11.axhline(y=ns.response[0, nc, 0], ls='--', lw=2, color='k')
        ax11.set_xlim(-2, 120)
        ax11.set_title("Natural Scenes", fontsize=20, color='blue')
        ax11.set_xlabel("Scene", fontsize=20)
        ax11.set_ylabel("DF/F (%)", fontsize=20)

        xtime = np.arange(-1 * ns.interlength / ns.acquisition_rate, (ns.sweeplength +
                                                                      ns.interlength) / ns.acquisition_rate, 1 / ns.acquisition_rate)
        nsp = np.argmax(ns.response[1:, nc, 0])
        subset_response = ns.sweep_response[ns.stim_table.frame == nsp]
        subset_response_p = subset_response[str(nc)].mean(
        ) + (subset_response[str(nc)][:].std() / np.sqrt(len(subset_response[str(nc)])))
        subset_response_n = subset_response[str(nc)].mean(
        ) - (subset_response[str(nc)][:].std() / np.sqrt(len(subset_response[str(nc)])))
        try:
            ax12.fill_between(xtime, subset_response_p,
                              subset_response_n, color='b', alpha=0.5)
        except:
            xtime = xtime[:-1]
            ax12.fill_between(xtime, subset_response_p,
                              subset_response_n, color='b', alpha=0.5)
        blank = ns.sweep_response[ns.stim_table.frame == -1]
        blank_p = blank[str(nc)].mean() + \
            (blank[str(nc)].std() / np.sqrt(len(blank[str(nc)])))
        blank_n = blank[str(nc)].mean() - \
            (blank[str(nc)].std() / np.sqrt(len(blank[str(nc)])))
        ax12.fill_between(xtime, blank_p, blank_n, color='gray', alpha=0.5)
        ax12.plot(xtime, subset_response[str(nc)].mean(), color='b', lw=2)
        ax12.plot(xtime, blank[str(nc)].mean(), color='k', lw=2)
        ax12.axvspan(0, ns.sweeplength / ns.acquisition_rate,
                     ymin=0, ymax=1, facecolor='gray', alpha=0.3)
        ax12.yaxis.set_major_locator(MaxNLocator(4))
        ax12.set_xlabel("Time (s)", fontsize=20)
        ax12.set_ylabel("DF/F (%)", fontsize=20)

        plt.tick_params(labelsize=16)
        plt.tight_layout()
        filename = 'Cell_' + str(nc + 1) + '_3SB.png'
        fullfilename = os.path.join(save_dir, filename)
        plt.savefig(fullfilename)
        plt.close()


def plot_running_a(dg, nm1, nm3, save_dir):
    logging.info("Plotting running data summary")
    nc = -1
    plt.figure(1, figsize=(10, 8))
    ax1 = plt.subplot2grid((4, 4), (0, 0), colspan=3)
    ax2 = plt.subplot2grid((4, 4), (0, 3))
    ax3 = plt.subplot2grid((4, 4), (1, 0))
    ax4 = plt.subplot2grid((4, 4), (1, 1))
    ax5 = plt.subplot2grid((4, 4), (1, 2))
    ax6 = plt.subplot2grid((4, 4), (2, 0), colspan=2)
    ax7 = plt.subplot2grid((4, 4), (3, 0), colspan=2)
    ax8 = plt.subplot2grid((4, 4), (2, 2), colspan=2)
    ax9 = plt.subplot2grid((4, 4), (3, 2), colspan=2)

    xtime = np.arange(0, np.size(dg.dxcm), 1.)
    xtime /= dg.acquisition_rate
    dif = np.ediff1d(dg.stim_table.start.values, to_begin=8000, to_end=8000)
    test = np.argwhere(dif > 5000)
    ax1.plot(xtime, dg.dxcm, color='k')
    for i in range(len(test) - 1):
        ax1.axvspan(xmin=(dg.stim_table.start.iloc[test[i]].values / dg.acquisition_rate), xmax=(
            dg.stim_table.end.iloc[test[i + 1] - 1].values / dg.acquisition_rate), color='gray', alpha=0.3)
    ax1.axvspan(xmin=nm1.stim_table.start.min() / nm1.acquisition_rate, xmax=(
        (nm1.stim_table.start.max() + nm1.sweeplength) / nm1.acquisition_rate), color='red', alpha=0.3)
    dif = np.ediff1d(nm3.stim_table.start.values, to_begin=8000, to_end=8000)
    test = np.argwhere(dif > 5000)
    for i in range(len(test) - 1):
        ax1.axvspan(xmin=(nm3.stim_table.start.iloc[test[i]].values / nm3.acquisition_rate), xmax=(
            (nm3.stim_table.end.iloc[test[i + 1] - 1].values + nm3.sweeplength) / nm3.acquisition_rate), color='blue', alpha=0.3)
    ax1.set_xlabel("Time (s)", fontsize=20)
    ax1.set_ylabel("Speed (cm/s)", fontsize=20)

    dx = dg.dxcm[np.logical_not(np.isnan(dg.dxcm))]
    ax2.hist(dx, bins=80, range=(-20, 100), color='gray')
    ax2.set_xlabel("Speed (cm/s)", fontsize=20)

    run_peak = np.where(dg.response[:, 1:, nc, 0]
                        == np.nanmax(dg.response[:, 1:, nc, 0]))
    peakori = run_peak[0][0]
    peaktf = run_peak[1][0] + 1

    ax3.errorbar(dg.orivals, dg.response[:, peaktf, nc, 0], yerr=dg.response[
                 :, peaktf, nc, 1], fmt='b.-', lw=2)
    ax3.fill_between(dg.orivals, np.repeat(dg.response[0, 0, nc, 0] + dg.response[0, 0, nc, 1], dg.number_ori), np.repeat(
        dg.response[0, 0, nc, 0] - dg.response[0, 0, nc, 1], dg.number_ori), color='gray', alpha=0.5)
    ax3.axhline(y=dg.response[0, 0, nc, 0], ls='--', color='k')
    ax3.annotate(str(dg.tfvals[peaktf]) + " Hz",
                 xy=(0, 0.9), xycoords='axes fraction', fontsize=14)
    ax3.set_xtick = (dg.orivals)
    ax3.set_xlabel("Direction (deg)", fontsize=20)
    ax3.set_ylabel("Speed (cm/s)", fontsize=20)
    ax3.yaxis.set_major_locator(MaxNLocator(6))

    ax4.errorbar(dg.tfvals[1:], dg.response[peakori, 1:, nc, 0], yerr=dg.response[
                 peakori, 1:, nc, 1], fmt='b.-', lw=2)
    ax4.fill_between(dg.tfvals[1:], np.repeat(dg.response[0, 0, nc, 0] + dg.response[0, 0, nc, 1], dg.number_tf - 1),
                     np.repeat(dg.response[0, 0, nc, 0] - dg.response[0, 0, nc, 1], dg.number_tf - 1), color='gray', alpha=0.5)
    ax4.axhline(y=dg.response[0, 0, nc, 0], ls='--', color='k')
    ax4.annotate(str(dg.orivals[peakori]) + " Deg",
                 xy=(0, 0.9), xycoords='axes fraction', fontsize=14)
    ax4.set_xticks = (dg.tfvals[1:])
    ax4.set_xlabel("Temporal frequency (Hz)", fontsize=20)
    ax4.yaxis.set_major_locator(MaxNLocator(6))

    im = ax5.imshow(dg.response[:, 1:, nc, 0],
                    cmap='gray', interpolation='none')
    ax5.set_ylabel("Direction", fontsize=16)
    ax5.set_xlabel("TF", fontsize=16)
    ax5.set_yticks(range(dg.number_ori))
    ax5.set_yticklabels(list(dg.orivals.astype(int).astype(str)))
    ax5.set_xticks(range(dg.number_tf - 1))
    ax5.set_xticklabels(list(dg.tfvals[1:].astype(int).astype(str)))
    cbar = plt.colorbar(im, ax=ax5)
    cbar.ax.set_ylabel('Speed (cm/s)', fontsize=8)
    for t in cbar.ax.get_yticklabels():
        t.set_fontsize(8)

    xtime = np.arange(0, nm1.sweeplength /
                      nm1.acquisition_rate, 1 / nm1.acquisition_rate)
    while len(xtime) > len(nm1.sweep_response['dx'].mean()):
        xtime = np.delete(xtime, -1)
    for index, row in nm1.sweep_response.iterrows():
        ax6.plot(xtime, nm1.sweep_response['dx'][index], lw=2)
    ax6.set_xlabel("Time (s)", fontsize=20)
    ax6.set_ylabel("DF/F", fontsize=20)
    ax6.set_title("Natural Movie 1", fontsize=20, color='red')

    temp = np.empty((len(nm1.stim_table), nm1.sweeplength))
    for i in range(len(nm1.stim_table)):
        temp[i, :] = nm1.sweep_response['dx'].iloc[i][:, 0]
    ax7.imshow(temp, cmap='gray', interpolation='none', aspect=40)
    ax7.set_ylabel("Trials", fontsize=20)
    ax7.set_xticks([])

    xtime = np.arange(0, nm3.sweeplength /
                      nm3.acquisition_rate, 1 / nm3.acquisition_rate)
    while len(xtime) > len(nm3.sweep_response['dx'].mean()):
        xtime = np.delete(xtime, -1)
    for index, row in nm3.sweep_response.iterrows():
        ax8.plot(xtime, nm3.sweep_response['dx'][index], lw=2)
    ax8.set_xlabel("Time (s)", fontsize=20)
    ax8.set_ylabel("DF/F", fontsize=20)
    ax8.set_title("Natural Movie Long", fontsize=20, color='blue')

    temp = np.empty((len(nm3.stim_table), nm3.sweeplength))
    for i in range(len(nm3.stim_table)):
        temp[i, :] = nm3.sweep_response['dx'].iloc[i][:, 0]
    ax9.imshow(temp, cmap='gray', interpolation='none', aspect=100)
    ax9.set_ylabel("Trials", fontsize=20)
    ax9.set_xticks([])

    plt.tick_params(labelsize=16)
    plt.tight_layout()
    plt.suptitle("Running Summary", fontsize=20)
    plt.subplots_adjust(top=0.9)
    filename = 'Running Summary.png'
    fullfilename = os.path.join(save_dir, filename)
    plt.savefig(fullfilename)
    plt.close()
