"""Collect rasters of responses for all LGN opto experiments as a function of stimulus type
(movies, gratings), run state, and opto state. Then, calculate various measures and save
results to DataFrames. Some variables (DataFrames and mseu string lists) can be optionally
saved to and loaded from .pickle files to avoid long and unnecessary recalculation times.

Run this script at the IPython/Jupyter command line; the `-i` runs scripts in IPython's
namespace instead of an empty one:

`run -i main.py`

`pvmvis` (default), `ntsrmvis` or `negntsrmvis` can be provided as an `--exptype=` argument
to this script, to load and later plot PV-Cre (most of the paper), Ntsr1-Cre (Fig1S4),
or Ntsr negative (Fig1S5) data, respectively.
"""

import os
import sys
import argparse

from numpy import pi
from numpy.random import choice
import scipy.stats
from scipy.stats import (linregress, ttest_rel, ttest_ind, ttest_1samp, wilcoxon,
                         ks_2samp)
from scipy.stats.mstats import gmean
from scipy.optimize import least_squares

import matplotlib.pyplot as plt
from matplotlib.offsetbox import AnchoredText
from matplotlib.ticker import ScalarFormatter, FormatStrFormatter
import seaborn as sns
import pandas as pd

from util import (load, load_all, save, export2csv,
                  desat, axes_disable_scientific,
                  ms2msstr, msu2msustr, mseustr2msustr, mseustrs2msustrs, mseustr2mstr,
                  mseustrs2mstrs, mseustr2msestr, mseustrs2msestrs,
                  findmse, fitmodel, residual_rsquared, linear_loss, get_max_snr,
                  intround, split_tranges, wrap_raster,
                  cf, saveall, wintitle, simpletraster,
                  raster2psth, raster2freqcomp, sparseness, reliability, snr_baden2016,
                  get_psth_peaks_gac, vector_OSI, percentile_ci, sum_of_gaussians)

## Define constants:

EXPTYPE = 'pvmvis' # pvmvis, ntsrmvis, negntsrmvis
if __name__ == '__main__':
    # parse command-line arguments:
    parser = argparse.ArgumentParser(description='Main script for natfeedback paper')
    parser.add_argument("--exptype", help="pvmvis, ntsrmvis, negntsrmvis")
    args = parser.parse_args()
    EXPTYPE = args.exptype or EXPTYPE
print('*** EXPTYPE: %r' % EXPTYPE)

RATETHRESH = 0.01 # mean firing rate threshold for unit inclusion, across all trial types, Hz
SNRTHRESH = 0.015
SCATTERPTHRESH = 0.05
SGNF2ALPHA = {True: 1, False: 0.4}
# visual drive unit inclusion criteria for gratings, ignores st8 and opto:
VDNCRIT = 1 # min num points in tuning curve that are at least VDZCRIT sems away from spont rate
VDZCRIT = 2.5

PSTHBASELINEMEDIANX = 2 # PSTH median multiplier to estimate baseline
PSTHPEAKMINTHRESH = 3 # PSTH peak detection threshold relative to estimated baseline, Hz

STIMTYPES = 'mvi', 'grt'
stimtype2axislabel = {'grt':'grating', 'mvi':'movie'}
MVIKINDS = 'nat',
ALLSTIMTYPES = MVIKINDS + ('grt',)
STIMTYPESPLUSALL = list(STIMTYPES) + ['mvi+grt'] # strings
STIMTYPESPLUSALLI = list(STIMTYPES) + [slice(None)] # for use as indices

# iterate over run states in this order:
ST8S = ['run', 'sit'] # excludes 'none'
ALLST8S = ['none', 'run', 'sit']
ST8COMBOS = [['none'], ['run', 'sit'], ['run', 'sit']]
OPTOCOMBOS = [[False, True], [False], [True]]
st82crit = {'run':'(run_speed > 1).mean() > 0.5',
            'sit':'(run_speed < 0.25).mean() > 0.5',
            'none':'none'}
rungreen = '#' + ''.join([ format(val, '02x') for val in [66, 168, 117] ])
sitorange = '#' + ''.join([ format(val, '02x') for val in [235, 115, 9] ])
st82clr = {'run':rungreen, 'sit':sitorange, 'none':'black'}
st82alpha = {None:1, 'run':1, 'sit':0.5} # sitting is like suppression, so desaturate colour
st82clrmap = {'run':plt.cm.Greens, 'sit':plt.cm.Oranges, 'none':plt.cm.Greys}

# define example mseus and their colors:
green = '#007f00' # deep green
violet = '#9f3fff' # pure violet (7f00ff) is a little too dark
darkblue = '#000064' # control condition
optoblue = '#2a7fff' # light shade of blue for opto condition
black = '#000000'
asterisk = 5, 2, 0 # 5 sided, asterisk style, 0 deg rotation, has rounded corners though
DEFSZ = 50 # default scatter plot point size (open points)
exmpli2clr = {1:violet, 2:green, 3:optoblue} # example neurons 1, 2, 3
# use different marker shapes to help distinguish example points:
exmpli2mrk = {1:'x', 2:'+', 3:asterisk}
exmpli2sz = {1:60, 2:70, 3:60} # marker size
exmpli2lw = {1:2, 2:2, 3:2} # (marker) line width
fig1exmpli = 1 # fig1 uses example neuron 1
fig5exmpli = 1 # fig5 uses example neuron 1

# designate movie example units:
pvmvimseu2exmpli = {'PVCre_2018_0003_s03_e03_u51':1, # example neuron 1
                    'PVCre_2017_0008_s12_e06_u56':2} # example neuron 2
ntsrmvimseu2exmpli = {'Ntsr1Cre_2019_0008_s03_e04_u18':1}
negntsrmvimseu2exmpli = {}
exptype2mviexmplis = {'pvmvis':pvmvimseu2exmpli, 'ntsrmvis':ntsrmvimseu2exmpli,
                      'negntsrmvis':negntsrmvimseu2exmpli}
mvimseu2exmpli = exptype2mviexmplis[EXPTYPE]

# designate grating example units:
pvgrtmseu2exmpli = {'PVCre_2018_0003_s03_e02_u51':1,  # example neuron 1
                    'PVCre_2018_0003_s05_e02_u164':3} # example neuron 3
ntsrgrtmseu2exmpli = {'Ntsr1Cre_2019_0008_s06_e02_u46':1}
ntsrgrtmseu2exmpli = {}
exptype2grtexmplis = {'pvmvis':pvgrtmseu2exmpli, 'ntsrmvis':ntsrgrtmseu2exmpli,
                      'negntsrmvis':ntsrgrtmseu2exmpli}
grtmseu2exmpli = exptype2grtexmplis[EXPTYPE]

mvimsu2exmpli = { mseustr2msustr(k):v for k, v in mvimseu2exmpli.items() }
grtmsu2exmpli = { mseustr2msustr(k):v for k, v in grtmseu2exmpli.items() }
msu2exmpli = {**mvimsu2exmpli, **grtmsu2exmpli} # merge the two

burstclr = 'red'

OPTOS = False, True
opto2tuni = {False:0, True:1} # for indexing into tun_params field in Tuning table
opto2fb = {None:'', False:'feedback', True:'suppression'}
opto2clr = {None:darkblue, False:darkblue, True:optoblue}
opto2alpha = {None:1, False:1, True:0.5} # with opto, feedback is removed, so desaturate colour

modmeasures = ['meanrate', 'meanrate02', 'meanrate35', 'blankmeanrate', 'blankcondmeanrate',
               'meanburstratio', 'blankmeanburstratio', 'blankcondmeanburstratio',
               'spars', 'rel', 'meanpkw', 'snr']
modmeasuresnoblankcond = modmeasures.copy()
modmeasuresnoblankcond.remove('blankcondmeanrate')
modmeasuresnoblankcond.remove('blankcondmeanburstratio')
modmeasuresnoblank = modmeasuresnoblankcond.copy()
modmeasuresnoblank.remove('blankmeanrate')
modmeasuresnoblank.remove('blankmeanburstratio')
measure2axislabel = {'meanrate':'FR', 'meanrate02':'FR', 'meanrate35':'FR',
                     'blankmeanrate':'FR', 'blankcondmeanrate':'FR',
                     'meanburstratio':'burst ratio', 'blankmeanburstratio':'burst ratio',
                     'blankcondmeanburstratio':'burst ratio',
                     'spars':'sparseness', 'rel':'reliability',
                     'meanpkw':'mean peak width', 'snr':'SNR'}
measure2axisunits = {'meanrate':' (spk/s)', 'meanrate02':' (spk/s)', 'meanrate35':' (spk/s)',
                     'meanpkw':' (s)'}
short2longaxislabel = {'FR':'Firing rate'}

fmodes = ['', 'nb', 'b', 'nr'] # (narrow) firing modes
fmode2txt = {'':'all', 'nb':'nonburst', 'b':'burst', 'nr':'nonrand'}

# Tuning table keys used for gratings:
ivs_order = 'grat_orientation, opto1'
tun_model = 'sum_of_gaussians'

model = 'threshlin' # linear or threshlin
relaverage = 'mean' # reliability averaging method: mean or median

DEFAULTFIGURESIZE = 2.8, 2.8 # normal
RASTERSCALEX = 0.75
OFFSETS = -0.5, 0.75 # wide raster time offsets, sec
OFFSETS02 = 0, -3 # limits to spikes from t=0 to 2 s, assuming 5 sec movies
OFFSETS35 = 3, 0 # limits to spikes from t=3 to 5 s
NTRIALSMVIGRT = 120 # limit movies to first 120 trials for more direct comparison with gratings

PSTHHEIGHT = 3
binw, tres, kernel = 0.02, 0.001, 'gauss' # for calculating PSTHs, sec
psthplottres = 0.010 # sec
ssx = intround(psthplottres / tres) # subsample factor to get the desired plot tres

MINNTRIALSAMPLETHRESH = 10

# ON/OFF/transient classification:
PDT = 0.050 # propagation delay time for movie onset (s)
WINDT = 0.20 # duration separating start of pre and PDT, and PDT and end of post (s)
EXCLTR = 0, 2*PDT # time range around PDT to treat as transient response
assert EXCLTR[0] <= PDT <= EXCLTR[1]

# load data from pickles:
print("NOTE: for a given EXPTYPE, only some pickle files exist and are relevant,\n"
      "      expect some 'file not found errors':")
subfolder = EXPTYPE if EXPTYPE != 'pvmvis' else ''
name2val = load_all(subfolder=subfolder)
locals().update(name2val)