from dataclasses import asdict, dataclass
from functools import cached_property
from time import sleep
from typing import Any, Dict, List, Optional, Union

from airflow import AirflowException
from airflow.models.taskinstance import Context
from airflow.providers.http.hooks.http import HttpHook
from constants import CRYPTO_COMPARE_HTTP_CONN_ID
from hooks.wrappers.http_stream import HttpStreamHook
from kafka import KafkaProducer
from operators.cryptocurrency.price.base import CryptocurrencyBaseOperator
from utils.exception import raise_airflow_exception
from utils.kafka import kafka_producer_context
from utils.request import get_request_json

@dataclass
class CryptocurrencyMultiPriceApiData:
    fsyms: str
    tsyms: str

class CryptocurrencyPriceSourcingStreamOperator(CryptocurrencyBaseOperator):
    cryptocurrency_http_conn_id: str = CRYPTO_COMPARE_HTTP_CONN_ID

    def __init__(
        self,
        symbol_list: List[str],
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.symbol_list = symbol_list

    def read(
        self,
        endpoint: str,
        data: Optional[Dict[str, Any]] = None,
    ):
        return self.try_to_get_request_json(
            http_hook=self.http_hook,
            endpoint=endpoint,
            data=data,
        )

    @cached_property
    def kafka_topic_name(self) -> str:
        return "cryptocurrency"

    @cached_property
    def standard_currency(self) -> str:
        return "USD"

    @cached_property
    def sleep_second(self) -> float:
        return 1 / len(self.symbol_list)

    @property
    def api_endpoint(self):
        return "pricemulti"

    def try_to_get_request_json(
        self,
        http_hook: Union[HttpHook, HttpStreamHook],
        endpoint: str,
        data: Optional[Dict[str, Any]] = None,
        retry_count: int = 5,
        err_msg: str = "",
    ) -> Dict[str, Any]:
        if retry_count <= 0:
            raise_airflow_exception(
                error_msg=err_msg,
                logger=self.log,
            )
        try:
            response_json = get_request_json(
                http_hook=http_hook,
                endpoint=endpoint,
                data=data,
                headers=self.api_header,
                back_off_cap=self.back_off_cap,
                back_off_base=self.back_off_base,
                proxies=self.proxies,
            )

        except AirflowException as e:
            self.log.info(f"raise AirflowException err_msg: {e}")
            sleep(10)
            return self.try_to_get_request_json(
                http_hook=http_hook,
                endpoint=endpoint,
                data=data,
                retry_count=retry_count - 1,
                err_msg=f"{err_msg} retry_count : {retry_count}\nerr_msg : {e} \n\n",
            )
        response_status = response_json.get("Response")
        if response_status == "Error":
            response_message = response_json.get("Message")
            if (
                response_message
                == "You are over your rate limit please upgrade your account!"
            ):
                self.PROXY_IP_IDX += 1
                self.log.info(
                    f"{response_message}, raise PROXY_IP_IDX to {self.PROXY_IP_IDX}"
                )
            return self.try_to_get_request_json(
                http_hook=http_hook,
                endpoint=endpoint,
                data=data,
                retry_count=retry_count - 1,
                err_msg=f"{err_msg} retry_count : {retry_count}\nerr_msg : {response_message} \n\n",
            )
        return response_json

    def write(
        self,
        json_data: List[Dict[str, Any]],
        kafka_producer: KafkaProducer,
    ) -> None:
        for data in json_data:
            kafka_producer.send(self.kafka_topic_name, value=data)
            kafka_producer.flush()

    @cached_property
    def api_data(
        self,
    ) -> CryptocurrencyMultiPriceApiData:
        return CryptocurrencyMultiPriceApiData(
            fsyms=",".join(self.symbol_list),
            tsyms=self.standard_currency,
        )

    @staticmethod
    def transform(data: Dict[str, Any]) -> List[Dict[str, Any]]:
        return [
            {"symbol": symbol, "close": usd.get("USD")} for symbol, usd in data.items()
        ]

    def execute(self, context: Context) -> None:
        # pass
        with kafka_producer_context() as kafka_producer:
            while 1:
                json_data = self.read(
                    endpoint=self.api_endpoint,
                    data=asdict(self.api_data),
                )
                transformed_data = self.transform(data=json_data)
                self.write(
                    json_data=transformed_data,
                    kafka_producer=kafka_producer,
                )
                sleep(10)


from flask import Blueprint, render_template, url_for

lonely = Blueprint('lonely',
                   __name__,
                   template_folder='./',
                   static_folder='./',
                   static_url_path='/')
lonely.display_name = 'Lonely'
lonely.published = True
lonely.description = "An interactive visualization of original music."

@lonely.route('/')
def _lonely():
    return render_template('lonely.html')


# -*- coding: utf-8 -*-
"""
Created on Wed Dec 11 15:05:36 2019

@author: Ashley
"""

# Manuscript Malezieux, Kees, Mulle submitted to Cell Reports
# Figure S3 - Complex spikes
# Description: changes in complex spikes with theta and LIA, plotted separately

# %% import modules

import os
import numpy as np
from scipy import stats
import matplotlib.pyplot as plt
from matplotlib.colors import rgb2hex
from itertools import compress
import matplotlib as mpl
from mpl_toolkits.axes_grid1 import Divider, Size
from mpl_toolkits.axes_grid1.mpl_axes import Axes

# %% definitions

# bootstrap: one-factor ANOVA-like (for any number of groups):
# is between-group variance bigger than within-group?
def calculate_F(groups_list):
    num_g = len(groups_list)
    box = np.concatenate(groups_list)
    GM = np.nanmedian(box)
    gm = np.zeros(num_g)
    gs = np.zeros(num_g)
    denom = np.zeros(num_g)
    for i in np.arange(num_g):
        gm[i] = np.nanmedian(groups_list[i])
        gs[i] = groups_list[i].size
        denom[i] = np.nansum(np.abs(groups_list[i]-np.nanmedian(groups_list[i])))
    F = (np.sum(gs*np.abs(GM-gm)))/(np.sum(denom))
    return F

# one-way anova for many groups: resampling from big box; only when variances are the same
# returns real_F, p_boot
def boot_anova(groups_list, num_b):
    # compute the real F
    real_F = calculate_F(groups_list)
    faux_F = np.zeros(num_b)
    # compute size and variance of groups
    groups_size = [np.nan]*len(groups_list)
    groups_var = [np.nan]*len(groups_list)
    for g in np.arange(len(groups_list)):
        groups_size[g] = groups_list[g].size
        groups_var[g] = MADAM(groups_list[g], np.nanmedian(groups_list[g]))
    # if the largest variance is more than 2x the smallest, resample within groups
    # demean each group and sample with replacement
    if max(groups_var)/min(groups_var) > 2:
        # subtract the median from each group before resampling
        dm_groups_list = [np.nan] * len(groups_list)
        for g in np.arange(len(groups_list)):
            dm_groups_list[g] = groups_list[g] - np.nanmedian(groups_list[g])
        # shuffle and deal from each group with replacement
        for b in np.arange(num_b):
            # deal into faux groups, each one the same size as in real data
            f_groups_list = [None] * len(groups_list)
            for g in np.arange(len(groups_list)):
                group = dm_groups_list[g]
                resample = group[np.random.randint(0, group.size, size=group.size)]
                f_groups_list[g] = resample
            faux_F[b] = calculate_F(f_groups_list)
        p_boot = np.sum(faux_F > real_F)/num_b
    # if the variances are mostly the same, resample from the big box without replacement
    else:
        box = np.concatenate(groups_list)
        for b in np.arange(num_b):
            np.random.shuffle(box)
            box1 = np.copy(box)
            # deal into fax groups, each one the same size as in real data
            f_groups_list = list()
            for g in np.arange(len(groups_list)):
                f_groups_list.append(box1[0:int(groups_size[g])])
                box1 = box1[int(groups_size[g]):]
            faux_F[b] = calculate_F(f_groups_list)
        p_boot = np.sum(faux_F > real_F)/num_b
    return real_F, p_boot

# definition for self_calculated variance (called MADAM??)
# VERSION: accounts for nans when dividing by number of samples
def MADAM(data_pts, descriptor):
    v = np.nansum(np.abs(data_pts-descriptor))/np.sum(~np.isnan(data_pts))
    return v

def boot_t(t_g0, t_g1, num_b):
    real_d = np.nanmedian(t_g1) - np.nanmedian(t_g0)
    faux_d = np.zeros(num_b)
    box = np.append(t_g0, t_g1)
    for b in np.arange(num_b):
        f_g0 = box[np.random.randint(0, box.size, size=t_g0.size)]
        f_g1 = box[np.random.randint(0, box.size, size=t_g1.size)]
        faux_d[b] = np.nanmedian(f_g1) - np.nanmedian(f_g0)
    p = np.sum(np.abs(faux_d) > np.abs(real_d))/num_b
    return real_d, p

def boot_pair_t(diff, num_b):
    real_d = np.mean(diff)
    faux_d = np.zeros(num_b)
    for b in np.arange(num_b):
        sample = np.random.choice([-1, 1], size = diff.size, replace=True)
        faux_d[b] = np.mean(diff*sample)
    p = np.sum(faux_d<real_d)/num_b
    return real_d, p

# definiton for finding 95% confidence intervals for each bin in histogram
# Version: for a **mean** histogram of **several** histograms
# H_array must be arranged [samples, bins]
def CI_avg_hist(H_array, num_b, CI_perc):
    real_H = np.nanmean(H_array, axis=0)
    faux_H = np.full([H_array.shape[1], num_b], np.nan)
    for b in np.arange(num_b):
        samp = np.random.randint(0, H_array.shape[0], H_array.shape[0])
        faux_H[:, b] = np.nanmean(H_array[samp, :], axis=0)
    CI_low, CI_high = np.nanpercentile(faux_H, [(100-CI_perc)/2, 100-((100-CI_perc)/2)],
                                    axis=1)
    return real_H, CI_high, CI_low

# eta = event triggered averages.  CHANGE: nans instead of removing events
def prepare_eta(signal, ts, event_times, win):
    win_npts = [ts[ts < ts[0] + np.abs(win[0])].size,
                ts[ts < ts[0] + np.abs(win[1])].size]
    et_ts = ts[0:np.sum(win_npts)] - ts[0] + win[0]
    et_signal = np.empty(0)
    if event_times.size > 0:
        if signal.ndim == 1:
            et_signal = np.zeros((et_ts.size, event_times.size))
            for i in np.arange(event_times.size):
                if np.logical_or((event_times[i]+win[0]<ts[0]), (event_times[i]+win[1]>ts[-1])):
                    et_signal[:, i] = np.nan*np.ones(et_ts.size)
                else:
                    # find index of closest timestamp to the event time
                    ind = np.argmin(np.abs(ts-event_times[i]))
                    et_signal[:, i] = signal[(ind - win_npts[0]): (ind + win_npts[1])]
        elif signal.ndim == 2:
            et_signal = np.zeros((signal.shape[0], et_ts.size, event_times.size))
            for i in np.arange(event_times.size):
                if np.logical_or((event_times[i]+win[0]<ts[0]), (event_times[i]+win[1]>ts[-1])):
                    et_signal[:, :, i] = np.nan*np.ones([signal.shape[0], et_ts.size])
                else:
                    # find index of closest timestamp to the event time
                    ind = np.argmin(np.abs(ts-event_times[i]))
                    et_signal[:, :, i] = signal[:, (ind - win_npts[0]):
                                                (ind + win_npts[1])]
    return et_signal, et_ts

# eta = event triggered averages
# this code is for point processes, but times instead of inds
def prepare_eta_times(pt_times, event_times, win):
    et_signal = []
    if (pt_times.size > 0) & (event_times.size > 0):
        # find pt_times that occur within window of each event_time
        for i in np.arange(event_times.size):
            ts_section = pt_times[(pt_times > event_times[i] + win[0]) &
                                  (pt_times < event_times[i] + win[1])]
            ts_section = ts_section - event_times[i]
            et_signal.append(ts_section)
    else:
        et_signal = [np.empty(0) for k in np.arange(event_times.size)]
    return et_signal

# eta = event triggered averages: Version: skip events too close to edge
def prepare_eta_skip(signal, ts, event_times, win):
    win_npts = [ts[ts < ts[0] + np.abs(win[0])].size,
                ts[ts < ts[0] + np.abs(win[1])].size]
    et_ts = ts[0:np.sum(win_npts)] - ts[0] + win[0]
    et_signal = np.empty(0)
    if event_times.size > 0:
        # remove any events that are too close to the beginning or end of recording
        if event_times[0]+win[0] < ts[0]:
            event_times = event_times[1:]
        if event_times[-1]+win[1] > ts[-1]:
            event_times = event_times[:-1]
        if signal.ndim == 1:
            et_signal = np.zeros((et_ts.size, event_times.size))
            for i in np.arange(event_times.size):
                # find index of closest timestamp to the event time
                ind = np.argmin(np.abs(ts-event_times[i]))
                et_signal[:, i] = signal[(ind - win_npts[0]): (ind + win_npts[1])]
        elif signal.ndim == 2:
            et_signal = np.zeros((signal.shape[0], et_ts.size, event_times.size))
            for i in np.arange(event_times.size):
                # find index of closest timestamp to the event time
                ind = np.argmin(np.abs(ts-event_times[i]))
                et_signal[:, :, i] = signal[:, (ind - win_npts[0]):
                                            (ind + win_npts[1])]
    return et_signal, et_ts

# %% load data

dataset_folder = (r'C:\Users\akees\Documents\Ashley\Papers\MIND 1\Cell Reports\Dryad upload\Dataset')

cell_files = os.listdir(dataset_folder)
data = [{} for k in np.arange(len(cell_files))]
for i in np.arange(len(cell_files)):
    full_file = os.path.join(dataset_folder, cell_files[i])
    data[i] = np.load(full_file, allow_pickle=True).item()

states = [{'id':'theta', 'bef':-2.5, 'aft':0.5, 'samp_time':2, 't_win':[-3, 3]},
          {'id':'LIA', 'bef':-4, 'aft':-1, 'samp_time':2, 't_win':[-4, 2]}]
ntl = ['nost', 'theta', 'LIA']

# %% process data - for dVm vs dFR analysis

# for each cell, find start and stop times for unlabeled times
for i in np.arange(len(data)):
    state_start = np.concatenate([data[i]['theta_start'], data[i]['LIA_start']])
    state_start = np.sort(state_start)
    state_stop = np.concatenate([data[i]['theta_stop'], data[i]['LIA_stop']])
    state_stop = np.sort(state_stop)
    data[i]['nost_start'] = np.append(data[i]['Vm_ds_ts'][0], state_stop)
    data[i]['nost_stop'] = np.append(state_start, data[i]['Vm_ds_ts'][-1])

# for each cell, make a new spike_times for specifically non-spikelets
for i in np.arange(len(data)):
    data[i]['spike_times'] = np.delete(data[i]['sp_times'],
                                       data[i]['spikelets_ind'])

# for each cell, calculate the isi (inter-spike-interval)
# for true spikes only
for i in np.arange(len(data)):
    if data[i]['spike_times'].size > 0:
        isi0 = data[i]['spike_times'][0] - data[i]['Vm_ds_ts'][0]
        data[i]['isi'] = np.ediff1d(data[i]['spike_times'], to_begin=isi0)
    else:
        data[i]['isi'] = np.empty(0)

# find the (true) spikes that are within bursts
burst_isi = 0.006  # seconds (Mizuseki 2012 0.006)
for i in np.arange(len(data)):
    if data[i]['spike_times'].size > 0:
        burst_bool = 1*(data[i]['isi'] < burst_isi)
        burst_sp = np.where(data[i]['isi'] < burst_isi)[0]
        burst_sp0 = np.where(np.ediff1d(burst_bool) == 1)[0]
        bursts = [None]*len(burst_sp0)
        if burst_sp0.size > 0:
            for j in np.arange(len(burst_sp0)-1):
                inds = np.append(burst_sp0[j], burst_sp[np.logical_and(burst_sp > burst_sp0[j],
                                 burst_sp < burst_sp0[j+1])])
                bursts[j] = data[i]['spike_times'][inds]
            # special case for the last burst:
            j = len(burst_sp0)-1
            inds = np.append(burst_sp0[j], burst_sp[burst_sp > burst_sp0[j]])
            bursts[j] = data[i]['spike_times'][inds]
        data[i]['bursts'] = bursts
    else:
        data[i]['bursts'] = [None]*0

# add windows triggered by start of some brain states
# collect relative times for (true) spikes, singles, doublets, bursts, and CS
for l in np.arange(len(states)):
    for i in np.arange(len(data)):
        t_Vm, t_ts = prepare_eta(data[i]['Vm_s_ds'], data[i]['Vm_ds_ts'],
                                 data[i][states[l]['id']+'_start'],
                                 states[l]['t_win'])
        t_sp = prepare_eta_times(data[i]['sp_times'],
                                    data[i][states[l]['id']+'_start'],
                                    states[l]['t_win'])
        t_spike = prepare_eta_times(data[i]['spike_times'],
                                    data[i][states[l]['id']+'_start'],
                                    states[l]['t_win'])
        spikelet_times = data[i]['sp_times'][data[i]['spikelets_ind'].astype(int)]
        t_spikelet = prepare_eta_times(spikelet_times,
                                    data[i][states[l]['id']+'_start'],
                                    states[l]['t_win'])
        single_times = data[i]['sp_times'][data[i]['singles_ind'].astype(int)]
        t_single = prepare_eta_times(single_times,
                                     data[i][states[l]['id']+'_start'],
                                     states[l]['t_win'])
        if data[i]['doublets_ind'].size > 0:
            doublet_times = data[i]['sp_times'][data[i]['doublets_ind'][0]]
        else:
            doublet_times = np.empty(0)
        t_doublet = prepare_eta_times(doublet_times,
                                     data[i][states[l]['id']+'_start'],
                                     states[l]['t_win'])
        burst_times = np.array([d[0] for d in data[i]['bursts']])
        t_burst = prepare_eta_times(burst_times,
                                     data[i][states[l]['id']+'_start'],
                                     states[l]['t_win'])
        t_CS = prepare_eta_times(data[i]['CS_start'],
                                     data[i][states[l]['id']+'_start'],
                                     states[l]['t_win'])
        data[i][states[l]['id']+'_Vm'] = t_Vm
        data[i][states[l]['id']+'_sp'] = t_sp  # all spikes and spikelets
        data[i][states[l]['id']+'_spike'] = t_spike  # all spikes (no spikelets)
        data[i][states[l]['id']+'_spikelet'] = t_spikelet
        data[i][states[l]['id']+'_single'] = t_single
        data[i][states[l]['id']+'_doublet'] = t_doublet
        data[i][states[l]['id']+'_burst'] = t_burst
        data[i][states[l]['id']+'_CS'] = t_CS
    states[l]['t_ts'] = t_ts

# add windows triggered by start of some brain states
# collect relative times for (true) spikes, singles, doublets, bursts, and CS
for l in np.arange(len(states)):
    for i in np.arange(len(data)):
        single_times = data[i]['sp_times'][data[i]['singles_ind'].astype(int)]
        if data[i]['doublets_ind'].size > 0:
            doublet_times = np.concatenate(data[i]['sp_times'][data[i]['doublets_ind']])
        else:
            doublet_times = np.empty(0)
        nonCS_times = np.sort(np.concatenate((single_times, doublet_times)))
        t_nonCS = prepare_eta_times(nonCS_times,
                                     data[i][states[l]['id']+'_start'],
                                     states[l]['t_win'])
        if len(data[i]['CS_ind']) > 0:
            CS_times = data[i]['sp_times'][np.concatenate(data[i]['CS_ind'])]
        else:
            CS_times = np.empty(0)
        t_CS = prepare_eta_times(CS_times,
                                     data[i][states[l]['id']+'_start'],
                                     states[l]['t_win'])
        data[i][states[l]['id']+'_CS_spikes'] = t_CS
        data[i][states[l]['id']+'_nonCS_spikes'] = t_nonCS

# for each event in each cell, calculate the CS rate and CS index
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        CS_rate = np.full(data[i][ntl[l]+'_start'].size, np.nan)
        CS_perc = np.full(data[i][ntl[l]+'_start'].size, np.nan)
        for j in np.arange(data[i][ntl[l]+'_start'].size):
            start = data[i][ntl[l]+'_start'][j]
            stop = data[i][ntl[l]+'_stop'][j]
            num_spikes = np.sum(np.logical_and(data[i]['spike_times'] > start,
                                               data[i]['spike_times'] < stop))
            if len(data[i]['CS_ind']) > 0:
                CS_spike_times = data[i]['sp_times'][np.concatenate(data[i]['CS_ind'])]
                num_CS_spikes = np.sum(np.logical_and(CS_spike_times > start,
                                                      CS_spike_times < stop))
                num_CS = np.sum(np.logical_and(data[i]['CS_start'] > start,
                                               data[i]['CS_start'] < stop))
            else:
                num_CS_spikes = 0
                num_CS = 0
            CS_perc[j] = num_CS_spikes/num_spikes
            CS_rate[j] = num_CS/(stop-start)
        data[i][ntl[l]+'_CS_rate'] = CS_rate
        data[i][ntl[l]+'_CS_perc'] = CS_perc

# %% event-based organization for dVm vs dFR

# make a dictionary to hold values collapsed over all cells
events = [{}  for k in np.arange(len(states))]
# find Vm0, dVm and significance for each run, excluding when Ih is changed
for l in np.arange(len(states)):
    all_c_p = np.empty(0)
    all_Ih = np.empty(0)
    all_Vm0 = np.empty(0)
    all_dVm = np.empty(0)
    all_dVm_p = np.empty(0)
    for i in np.arange(len(data)):
        samp_freq = 1/(data[i]['Vm_ds_ts'][1] - data[i]['Vm_ds_ts'][0])
        num_ind = int(states[l]['samp_time']*samp_freq)
        # find index of dIh_times
        dIh_ind = data[i]['dIh_times']*samp_freq
        dIh_ind = dIh_ind.astype(int)
        c_p = np.zeros(data[i][states[l]['id']+'_start'].size)
        Ih = np.zeros(data[i][states[l]['id']+'_start'].size)
        Vm0 = np.zeros(data[i][states[l]['id']+'_start'].size)
        dVm = np.zeros(data[i][states[l]['id']+'_start'].size)
        dVm_p = np.zeros(data[i][states[l]['id']+'_start'].size)
        for j in np.arange(data[i][states[l]['id']+'_start'].size):
            # find indices
            bef_ind = int(np.sum(data[i]['Vm_ds_ts'] <
                          (data[i][states[l]['id']+'_start'][j] + states[l]['bef'])))
            aft_ind = int(np.sum(data[i]['Vm_ds_ts'] <
                          (data[i][states[l]['id']+'_start'][j] + states[l]['aft'])))
            # put nan if times are straddling a time when dIh is changed
            dIh_true = np.where((dIh_ind > bef_ind) & (dIh_ind < aft_ind + num_ind))[0]
            if dIh_true.size > 0:
                Ih[j] = np.nan
                Vm0[j] = np.nan
                dVm[j] = np.nan
                dVm_p[j] = np.nan
            else:
                if np.logical_or(l==0, l==1):
                    c_p[j] = data[i][states[l]['id']+'_cell_p']
                else:
                    c_p[j] = data[i]['theta_cell_p']
                Ih_ind = np.searchsorted(data[i]['Vm_Ih_ts'],
                                         data[i][states[l]['id']+'_start'][j])
                Ih[j] = data[i]['Vm_Ih'][Ih_ind]
                # test whether Vm values are significantly different
                # Welch's t-test: normal, unequal variances, independent samp
                t, p = stats.ttest_ind(data[i]['Vm_ds'][bef_ind:bef_ind+num_ind],
                                       data[i]['Vm_ds'][aft_ind:aft_ind+num_ind],
                                       equal_var=False, nan_policy='omit')
                dVm_p[j] = p
                if (np.nanmean(data[i]['Vm_ds'][aft_ind:aft_ind+num_ind]) -
                    np.nanmean(data[i]['Vm_ds'][bef_ind:bef_ind+num_ind])) > 0:
                    Vm0[j] = np.nanmin(data[i]['Vm_s_ds'][bef_ind:bef_ind+num_ind])
                    dVm[j] = (np.nanmax(data[i]['Vm_s_ds'][aft_ind:aft_ind+num_ind]) -
                              np.nanmin(data[i]['Vm_s_ds'][bef_ind:bef_ind+num_ind]))
                else:
                    Vm0[j] = np.nanmax(data[i]['Vm_s_ds'][bef_ind:bef_ind+num_ind])
                    dVm[j] = (np.nanmin(data[i]['Vm_s_ds'][aft_ind:aft_ind+num_ind]) -
                              np.nanmax(data[i]['Vm_s_ds'][bef_ind:bef_ind+num_ind]))
        data[i][states[l]['id']+'_c_p'] = c_p
        data[i][states[l]['id']+'_Ih'] = Ih
        data[i][states[l]['id']+'_Vm0'] = Vm0
        data[i][states[l]['id']+'_dVm'] = dVm
        data[i][states[l]['id']+'_dVm_p'] = dVm_p
        all_c_p = np.append(all_c_p, c_p)
        all_Ih = np.append(all_Ih, Ih)
        all_Vm0 = np.append(all_Vm0, Vm0)
        all_dVm = np.append(all_dVm, dVm)
        all_dVm_p = np.append(all_dVm_p, dVm_p)
    events[l]['c_p'] = all_c_p
    events[l]['Ih'] = all_Ih
    events[l]['Vm0'] = all_Vm0
    events[l]['dVm'] = all_dVm
    events[l]['dVm_p'] = all_dVm_p

# add windows triggered by start of some brain states
for l in np.arange(len(states)):
    for i in np.arange(len(data)):
        t_Vm, t_ts = prepare_eta(data[i]['Vm_s_ds'], data[i]['Vm_ds_ts'],
                                 data[i][states[l]['id']+'_start'],
                                 states[l]['t_win'])
        t_sp = prepare_eta_times(data[i]['sp_times'],
                                 data[i][states[l]['id']+'_start'],
                                 states[l]['t_win'])
        data[i][states[l]['id']+'_Vm'] = t_Vm
        data[i][states[l]['id']+'_sp'] = t_sp
    states[l]['t_ts'] = t_ts

# add triggered windows to event dictionary
for l in np.arange(len(events)):
    raster_sp = []
    psth_sp = np.empty(0)
    Vm = np.empty((states[l]['t_ts'].shape[0], 0))
    duration = np.empty(0)
    cell_id = np.empty(0)
    for i in np.arange(len(data)):
        cell_psth_sp = np.empty(0)
        if data[i][states[l]['id']+'_start'].size > 0:
            Vm = np.append(Vm, data[i][states[l]['id']+'_Vm'], axis=1)
            duration = np.append(duration, (data[i][states[l]['id']+'_stop'] -
                                     data[i][states[l]['id']+'_start']))
            if isinstance(data[i]['cell_id'], str):
                ind = data[i]['cell_id'].find('_')
                cell_int = int(data[i]['cell_id'][:ind])*np.ones(data[i][states[l]['id']+'_start'].size)
                cell_id = np.append(cell_id, cell_int)
            else:
                cell_int = data[i]['cell_id']*np.ones(data[i][states[l]['id']+'_start'].size)
                cell_id = np.append(cell_id, cell_int)
            for j in np.arange(data[i][states[l]['id']+'_start'].size):
                psth_sp = np.append(psth_sp, data[i][states[l]['id']+'_sp'][j])
                cell_psth_sp = np.append(cell_psth_sp, data[i][states[l]['id']+'_sp'][j])
                raster_sp.append(data[i][states[l]['id']+'_sp'][j])
            data[i][states[l]['id']+'_psth_sp'] = cell_psth_sp
    # remove nans
    no_nan = np.logical_and([~np.isnan(Vm).any(axis=0)],
                            [~np.isnan(events[l]['Vm0'])]).flatten()
    events[l]['Vm'] = Vm[:, no_nan]
    events[l]['cell_id'] = cell_id[no_nan]
    events[l]['duration'] = duration[no_nan]
    events[l]['raster_sp'] = list(compress(raster_sp, no_nan))
    events[l]['c_p'] = events[l]['c_p'][no_nan]
    events[l]['Ih'] = events[l]['Ih'][no_nan]
    events[l]['Vm0'] = events[l]['Vm0'][no_nan]
    events[l]['dVm'] = events[l]['dVm'][no_nan]
    events[l]['dVm_p'] = events[l]['dVm_p'][no_nan]

# %% process data - for CS/burst analysis

# for each (true) spike, determine which state it occurs in (and those in no state)
# Version: all spikes, not just those used for spike threshold analysis
for i in np.arange(len(data)):
    nost_sp = np.ones(data[i]['spike_times'].size, dtype=bool)
    for l in np.arange(len(states)):
        state_sp = np.zeros(data[i]['spike_times'].size, dtype=bool)
        for j in np.arange(data[i][states[l]['id']+'_start'].size):
            # find the spikes that occur in that event
            temp_bool = np.all((data[i]['spike_times'] > data[i][states[l]['id']+'_start'][j],
                                data[i]['spike_times'] < data[i][states[l]['id']+'_stop'][j]),
                               axis=0)
            state_sp = state_sp + temp_bool
        data[i][states[l]['id']+'_spike_bool'] = np.squeeze(state_sp)
        nost_sp = nost_sp*[state_sp == False]
    data[i]['nost_spike_bool'] = np.squeeze(nost_sp)

# for each burst, determine which state it occurs in (and those in no state)
for i in np.arange(len(data)):
    burst_start = np.array([d[0] for d in data[i]['bursts']])
    nost_bst = np.ones(burst_start.size, dtype=bool)
    for l in np.arange(len(states)):
        state_bst = np.zeros(burst_start.size, dtype=bool)
        for j in np.arange(data[i][states[l]['id']+'_start'].size):
            # find the bursts that start during that event
            temp_bool = np.all((burst_start > data[i][states[l]['id']+'_start'][j],
                                burst_start < data[i][states[l]['id']+'_stop'][j]),
                               axis=0)
            state_bst = state_bst + temp_bool
        data[i][states[l]['id']+'_bst_bool'] = np.squeeze(state_bst)
        nost_bst = nost_bst*[state_bst == False]
    data[i]['nost_bst_bool'] = np.squeeze(nost_bst)

# for each cell, determine the % of spikes in bursts for theta, LIA, nost
ntl = ['nost', 'theta', 'LIA']
for i in np.arange(len(data)):
    burst_perc = np.full(3, np.nan)
    sp_times = data[i]['spike_times']
    if len(data[i]['bursts']) > 0:
        burst_times = np.concatenate(data[i]['bursts'])
    else:
        burst_times = 0
    for l in np.arange(len(ntl)):
        total_spikes = 0
        burst_spikes = 0
        for j in np.arange(data[i][ntl[l]+'_start'].size):
            start = data[i][ntl[l]+'_start'][j]
            stop = data[i][ntl[l]+'_stop'][j]
            spikes = np.sum(np.logical_and(sp_times > start, sp_times < stop))
            bursts = np.sum(np.logical_and(burst_times > start, burst_times < stop))
            total_spikes = total_spikes + spikes
            burst_spikes = burst_spikes + bursts
        if total_spikes != 0:
            burst_perc[l] = burst_spikes/total_spikes
    data[i]['burst_perc'] = burst_perc

# for each CS, determine which state it occurs in (and those in no state)
for i in np.arange(len(data)):
    nost_CS = np.ones(data[i]['CS_start'].size, dtype=bool)
    for l in np.arange(len(states)):
        state_CS = np.zeros(data[i]['CS_start'].size, dtype=bool)
        for j in np.arange(data[i][states[l]['id']+'_start'].size):
            # find the bursts that start during that event
            temp_bool = np.all((data[i]['CS_start'] > data[i][states[l]['id']+'_start'][j],
                                data[i]['CS_start'] < data[i][states[l]['id']+'_stop'][j]),
                               axis=0)
            state_CS = state_CS + temp_bool
        data[i][states[l]['id']+'_CS_bool'] = np.squeeze(state_CS)
        nost_CS = nost_CS*[state_CS == False]
    data[i]['nost_CS_bool'] = np.squeeze(nost_CS)

# collect the CS features divided by state
keep_cells = np.where([isinstance(d['cell_id'], int) for d in data])[0]
CS_ntl = [{} for l in np.arange(len(ntl))]
for l in np.arange(len(ntl)):
    num_sp = np.empty(0)
    CS_dur = np.empty(0)
    CS_height_Vm = np.empty(0)
    CS_rel_ahp_Vm = np.empty(0)
    for k in np.arange(keep_cells.size):
        i = keep_cells[k]
        num_sp = np.append(num_sp, np.array([d.size for d in data[i]['CS_ind']])[data[i][ntl[l]+'_CS_bool']])
        CS_dur = np.append(CS_dur, (data[i]['CS_stop'] - data[i]['CS_start'])[data[i][ntl[l]+'_CS_bool']])
        CS_height_Vm = np.append(CS_height_Vm, (data[i]['CS_max_Vm'][data[i][ntl[l]+'_CS_bool']] -
                                 data[i]['CS_start_Vm'][data[i][ntl[l]+'_CS_bool']]))
        CS_rel_ahp_Vm = np.append(CS_rel_ahp_Vm, (data[i]['CS_stop_Vm'][data[i][ntl[l]+'_CS_bool']] -
                                 data[i]['CS_start_Vm'][data[i][ntl[l]+'_CS_bool']]))
    CS_ntl[l]['num_sp'] = num_sp
    CS_ntl[l]['CS_dur'] = CS_dur
    CS_ntl[l]['CS_height_Vm'] = CS_height_Vm
    CS_ntl[l]['CS_rel_ahp_Vm'] = CS_rel_ahp_Vm

# %% set figure parameters

# set colors
# states
c_run_theta = [0.398, 0.668, 0.547]
c_nonrun_theta = [0.777, 0.844, 0.773]
c_LIA = [0.863, 0.734, 0.582]
# response type
c_hyp = [0.184, 0.285, 0.430]
c_dep = [0.629, 0.121, 0.047]
c_no = [1, 1, 1]
c_lhyp = [0.62, 0.71, 0.84]
c_ldep = [0.97, 0.71, 0.67]
# dependent variables
c_sp = [0.398, 0.461, 0.703]
c_Vm = [0.398, 0.461, 0.703]
# other
c_lgry = [0.75, 0.75, 0.75]
c_mgry = [0.5, 0.5, 0.5]
c_dgry = [0.25, 0.25, 0.25]
c_wht = [1, 1, 1]
c_blk = [0, 0, 0]
c_bwn = [0.340, 0.242, 0.125]
c_lbwn = [0.645, 0.484, 0.394]
c_grn = [0.148, 0.360, 0.000]

c_dVm = [c_hyp, c_mgry, c_dep]
c_state = [c_mgry, c_run_theta, c_lbwn]
c_state_dark = [c_dgry, c_grn, c_bwn]
c_tl = [c_run_theta, c_lbwn]
c_tnl = [c_run_theta, c_blk, c_lbwn]

# set style defaults
mpl.rcParams['font.size'] = 8
mpl.rcParams['savefig.dpi'] = 1200
mpl.rcParams['lines.linewidth'] = 1.5
mpl.rcParams['font.sans-serif'] = "Arial"
mpl.rcParams['font.family'] = "sans-serif"
mpl.rcParams['axes.spines.right'] = False
mpl.rcParams['axes.spines.top'] = False
mpl.rcParams['axes.linewidth'] = 1
mpl.rcParams['xtick.major.size'] = 4
mpl.rcParams['xtick.major.width'] = 1
mpl.rcParams['ytick.major.size'] = 4
mpl.rcParams['ytick.major.width'] = 1
mpl.rcParams['boxplot.whiskerprops.linestyle'] = '-'
mpl.rcParams['patch.force_edgecolor'] = True
mpl.rcParams['patch.facecolor'] = 'b'

# set figure output folder
fig_folder = r'C:\Users\akees\Documents\Ashley\Figures\2020-05_Paper_MIND1\FigS3'

# set which states to plot
## all states
#d_l = [0, 1, 2]
# theta only
d_l = [0, 1]
## LIA only
#d_l = [0, 2]

# %% make hist isi figure

keep_cells = [isinstance(d['cell_id'], int) for d in data]
theta_cell_p = np.array([d['theta_cell_p'] for d in data])[keep_cells]
LIA_cell_p = np.array([d['LIA_cell_p'] for d in data])[keep_cells]

c_state_hist = [c_mgry, c_grn, c_bwn]
c_state_fill = [c_lgry, c_run_theta, c_lbwn]

# prep numbers for mean hist isi - divided between states
bins = np.arange(0, 200, 1)
H = np.full([len(data), bins.size-1, len(ntl)], np.nan)
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        H[i, :, l] = np.histogram(1000*data[i]['isi'][data[i][ntl[l]+'_spike_bool']],
                                  bins=bins, density=True)[0]
# remove extra recordings from cells
H = H[keep_cells, :, :]
# define the 95% CI for each bin by randomly selecting (with replacement) over cells
H_mean = np.full([H.shape[1], H.shape[2]], np.nan)
H_CI_high = np.full([H.shape[1], H.shape[2]], np.nan)
H_CI_low = np.full([H.shape[1], H.shape[2]], np.nan)
CI_perc = 95
num_b = 1000
for l in np.arange(len(ntl)):
    real_H, CI_high, CI_low = CI_avg_hist(H[:, :, l], num_b, CI_perc)
    H_mean[:, l] = real_H
    H_CI_high[:, l] = CI_high
    H_CI_low[:, l] = CI_low
# plot the mean hist isi
fig, ax = plt.subplots(1, figsize=[4.5, 2.2])
for l in d_l:
    ax.plot(bins[:-1], H_mean[:, l], color=c_state_hist[l], zorder=2)
    ax.fill_between(bins[:-1], H_CI_low[:, l], H_CI_high[:, l],
                    facecolor=c_state_fill[l], linewidth=0, zorder=1, alpha=0.25)
#ax.axvline(6, color=c_blk, linestyle='--')
#ax.set_xlim([0, 200])
ax.set_ylim([0, 0.27])
ax.set_yticks([0, 0.05, 0.1, 0.15, 0.2, 0.25])
ax.set_yticklabels([0, '', 0.1, '', 0.2, ''])
ax.set_ylabel('proportion')
ax.set_xscale('log')
ax.set_xlim([1, 100])
ax.set_xlabel('inter-spike interval (ms)')
fig.tight_layout()
plt.savefig(os.path.join(fig_folder, ntl[d_l[-1]]+'_hist_isi.png'), transparent=True)

# do the stats for the above figure
S = np.full([len(data), len(ntl)], np.nan)
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        isi = 1000*data[i]['isi'][data[i][ntl[l]+'_spike_bool']]
        if isi.size > 10:
            S[i, l] = np.nanmedian(isi)
# remove extra recordings from cells
S = S[keep_cells, :]

# do the paired boot stats
num_b = 1000
p = np.full(len(ntl) - 1, np.nan)
d = np.full(len(ntl) - 1, np.nan)
for l in np.arange(len(ntl) - 1):
    dif = S[:, l+1] - S[:, 0]
    # remove nans
    dif = dif[~np.isnan(dif)]
    d[l], p[l] = boot_pair_t(dif, num_b)
    print(dif.size)
print(d)
print(p)

# do the paired boot stats for theta hyp and LIA dep cells only
num_b = 1000
# theta hyp cells
dif = S[:, 1][theta_cell_p < 0.05] - S[:, 0][theta_cell_p < 0.05]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)
# LIA dep cells
dif = S[:, 2][LIA_cell_p > 0.95] - S[:, 0][LIA_cell_p > 0.95]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)

# descriptive numbers
l = 1
np.sum(~np.isnan(S[:, l]))
np.nanmedian(S[:, l])
np.nanstd(S[:, l])
MADAM(S[:, l], np.nanmedian(S[:, l]))

# %% make CS features figures - number of spikes per CS

keep_cells = [isinstance(d['cell_id'], int) for d in data]
theta_cell_p = np.array([d['theta_cell_p'] for d in data])[keep_cells]
LIA_cell_p = np.array([d['LIA_cell_p'] for d in data])[keep_cells]

c_state_hist = [c_mgry, c_grn, c_bwn]
c_state_fill = [c_lgry, c_run_theta, c_lbwn]

# prep numbers for mean hist of # spikes in CS - divided between states
bins = np.arange(0.5, 51.5, 1)
H = np.full([len(data), bins.size-1, len(ntl)], np.nan)
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        num_sp = np.array([d.size for d in data[i]['CS_ind']])[data[i][ntl[l]+'_CS_bool']]
        H[i, :, l] = np.histogram(num_sp, bins=bins)[0]
        # normalize to total number of CS
        H[i, :, l] = H[i, :, l]/np.sum(H[i, :, l])
# remove extra recordings from cells
H = H[keep_cells, :, :]
# define the 95% CI for each bin by randomly selecting (with replacement) over cells
H_mean = np.full([H.shape[1], H.shape[2]], np.nan)
H_CI_high = np.full([H.shape[1], H.shape[2]], np.nan)
H_CI_low = np.full([H.shape[1], H.shape[2]], np.nan)
CI_perc = 95
num_b = 1000
for l in np.arange(len(ntl)):
    real_H, CI_high, CI_low = CI_avg_hist(H[:, :, l], num_b, CI_perc)
    H_mean[:, l] = real_H
    H_CI_high[:, l] = CI_high
    H_CI_low[:, l] = CI_low
# plot the mean hist isi
# create a figure with axes of defined size
fig = plt.figure(figsize=[2, 2])
# The first items are for padding and the second items are for the axes.
# sizes are in inch.
h = [Size.Fixed(0.5), Size.Fixed(1.2)]
v = [Size.Fixed(0.5), Size.Fixed(1.2)]
divider = Divider(fig, (0.0, 0.0, 1., 1.), h, v, aspect=False)
ax = Axes(fig, divider.get_position())
ax.set_axes_locator(divider.new_locator(nx=1, ny=1))
fig.add_axes(ax)
#for l in np.arange(len(ntl)):
for l in d_l:
    ax.plot(np.arange(1, bins.size), H_mean[:, l], color=c_state_hist[l], zorder=2)
    ax.fill_between(np.arange(1, bins.size), H_CI_low[:, l], H_CI_high[:, l],
                    facecolor=c_state_fill[l], linewidth=0, zorder=1, alpha=0.25)
ax.set_ylim([0, 0.4])
ax.set_yticks([0, 0.2, 0.4])
ax.set_yticklabels([0, '', 0.4])
ax.set_xlim([3, 12])
ax.set_xticks([3, 6, 9, 12])
ax.set_xlabel('number of spikes')
ax.set_ylabel('proportion')
plt.savefig(os.path.join(fig_folder, ntl[d_l[-1]]+'_CS_num_spikes.png'), transparent=True)

# do the stats for the above figure
S = np.full([len(data), len(ntl)], np.nan)
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        num_sp = np.array([d.size for d in data[i]['CS_ind']])[data[i][ntl[l]+'_CS_bool']]
        if num_sp.size > 0:
            #S[i, l] = stats.mode(num_sp)[0]
            S[i, l] = np.nanmedian(num_sp)
# remove extra recordings from cells
S = S[keep_cells, :]

# do the paired boot stats
num_b = 1000
p = np.full(len(ntl) - 1, np.nan)
d = np.full(len(ntl) - 1, np.nan)
for l in np.arange(len(ntl) - 1):
    dif = S[:, l+1] - S[:, 0]
    # remove nans
    dif = dif[~np.isnan(dif)]
    d[l], p[l] = boot_pair_t(dif, num_b)
    print(dif.size)
print(d)
print(p)

# do the paired boot stats for theta hyp and LIA dep cells only
num_b = 1000
# theta hyp cells
dif = S[:, 1][theta_cell_p < 0.05] - S[:, 0][theta_cell_p < 0.05]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)
# LIA dep cells
dif = S[:, 2][LIA_cell_p > 0.95] - S[:, 0][LIA_cell_p > 0.95]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)

# descriptive numbers
l = 1
np.sum(~np.isnan(S[:, l]))
np.nanmedian(S[:, l])
np.nanstd(S[:, l])
MADAM(S[:, l], np.nanmedian(S[:, l]))

# %% make CS features figures - CS duration

# prep numbers for mean hist CS duration - divided between states
bins = np.arange(0, 0.5, 0.02)
H = np.full([len(data), bins.size-1, len(ntl)], np.nan)
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        CS_dur = (data[i]['CS_stop'] - data[i]['CS_start'])[data[i][ntl[l]+'_CS_bool']]
        H[i, :, l] = np.histogram(CS_dur, bins=bins)[0]
        # normalize to total number of CS
        H[i, :, l] = H[i, :, l]/np.sum(H[i, :, l])
# define the 95% CI for each bin by randomly selecting (with replacement) over cells
H_mean = np.full([H.shape[1], H.shape[2]], np.nan)
H_CI_high = np.full([H.shape[1], H.shape[2]], np.nan)
H_CI_low = np.full([H.shape[1], H.shape[2]], np.nan)
CI_perc = 95
num_b = 1000
for l in np.arange(len(ntl)):
    real_H, CI_high, CI_low = CI_avg_hist(H[:, :, l], num_b, CI_perc)
    H_mean[:, l] = real_H
    H_CI_high[:, l] = CI_high
    H_CI_low[:, l] = CI_low
# plot the mean hist CS duration
# create a figure with axes of defined size
fig = plt.figure(figsize=[2, 2])
# The first items are for padding and the second items are for the axes.
# sizes are in inch.
h = [Size.Fixed(0.5), Size.Fixed(1.2)]
v = [Size.Fixed(0.5), Size.Fixed(1.2)]
divider = Divider(fig, (0.0, 0.0, 1., 1.), h, v, aspect=False)
ax = Axes(fig, divider.get_position())
ax.set_axes_locator(divider.new_locator(nx=1, ny=1))
fig.add_axes(ax)
for l in d_l:
    ax.plot(bins[:-1], H_mean[:, l], color=c_state_hist[l], zorder=2)
    ax.fill_between(bins[:-1], H_CI_low[:, l], H_CI_high[:, l],
                    facecolor=c_state_fill[l], linewidth=0, zorder=1, alpha=0.25)
ax.set_xlim([0, 0.2])
ax.set_ylim([0, 0.5])
ax.set_xticks([0, 0.1, 0.2])
ax.set_yticks([0, 0.25, 0.5])
ax.set_yticklabels([0, '', 0.5])
ax.set_xlabel('duration (ms)')
plt.savefig(os.path.join(fig_folder, ntl[d_l[-1]]+'_CS_dur.png'), transparent=True)

# do the stats for the above figure
S = np.full([len(data), len(ntl)], np.nan)
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        CS_dur = 1000*(data[i]['CS_stop'] - data[i]['CS_start'])[data[i][ntl[l]+'_CS_bool']]
        if CS_dur.size > 0:
            S[i, l] = np.nanmedian(CS_dur)
# remove extra recordings from cells
S = S[keep_cells, :]

## do the kruskall-wallace
#H, p_kw = stats.kruskal(S[:, 0], S[:, 1], S[:, 2], nan_policy='omit')

# do the paired boot stats
num_b = 1000
p = np.full(len(ntl) - 1, np.nan)
d = np.full(len(ntl) - 1, np.nan)
for l in np.arange(len(ntl) - 1):
    dif = S[:, l+1] - S[:, 0]
    # remove nans
    dif = dif[~np.isnan(dif)]
    d[l], p[l] = boot_pair_t(dif, num_b)
    print(dif.size)
print(d)
print(p)

# do the paired boot stats for theta hyp and LIA dep cells only
num_b = 1000
# theta hyp cells
dif = S[:, 1][theta_cell_p < 0.05] - S[:, 0][theta_cell_p < 0.05]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)
# LIA dep cells
dif = S[:, 2][LIA_cell_p > 0.95] - S[:, 0][LIA_cell_p > 0.95]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)
# LIA nondep cells
dif = S[:, 2][LIA_cell_p < 0.95] - S[:, 0][LIA_cell_p < 0.95]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)

# descriptive numbers
l = 1
np.sum(~np.isnan(S[:, l]))
np.nanmedian(S[:, l])
np.nanstd(S[:, l])
MADAM(S[:, l], np.nanmedian(S[:, l]))

# %% make CS features figures - subthreshold depolarization during CS

# prep numbers for mean hist CS max-start Vm - divided between states
bins = np.arange(0, 40, 2)
H = np.full([len(data), bins.size-1, len(ntl)], np.nan)
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        CS_height_Vm = (data[i]['CS_max_Vm'][data[i][ntl[l]+'_CS_bool']] -
                        data[i]['CS_start_Vm'][data[i][ntl[l]+'_CS_bool']])
        H[i, :, l] = np.histogram(CS_height_Vm, bins=bins)[0]
        # normalize to total number of CS
        H[i, :, l] = H[i, :, l]/np.sum(H[i, :, l])
# remove extra recordings from cells
H = H[keep_cells, :, :]
# define the 95% CI for each bin by randomly selecting (with replacement) over cells
H_mean = np.full([H.shape[1], H.shape[2]], np.nan)
H_CI_high = np.full([H.shape[1], H.shape[2]], np.nan)
H_CI_low = np.full([H.shape[1], H.shape[2]], np.nan)
CI_perc = 95
num_b = 1000
for l in np.arange(len(ntl)):
    real_H, CI_high, CI_low = CI_avg_hist(H[:, :, l], num_b, CI_perc)
    H_mean[:, l] = real_H
    H_CI_high[:, l] = CI_high
    H_CI_low[:, l] = CI_low
# plot the mean hist isi
# create a figure with axes of defined size
fig = plt.figure(figsize=[2, 2])
# The first items are for padding and the second items are for the axes.
# sizes are in inch.
h = [Size.Fixed(0.5), Size.Fixed(1.2)]
v = [Size.Fixed(0.5), Size.Fixed(1.2)]
divider = Divider(fig, (0.0, 0.0, 1., 1.), h, v, aspect=False)
ax = Axes(fig, divider.get_position())
ax.set_axes_locator(divider.new_locator(nx=1, ny=1))
fig.add_axes(ax)
for l in d_l:
    ax.plot(bins[:-1], H_mean[:, l], color=c_state_hist[l], zorder=2)
    ax.fill_between(bins[:-1], H_CI_low[:, l], H_CI_high[:, l],
                    facecolor=c_state_fill[l], linewidth=0, zorder=1, alpha=0.25)
#ax.set_xlabel('CS height (mV)')
#ax.set_ylabel('proportion of CS')
ax.set_xlim([0, 35])
ax.set_xticks([0, 10, 20, 30])
ax.set_xlabel('subthreshold depolarization (mV)')
ax.set_ylim([0, 0.3])
ax.set_yticks([0, 0.1, 0.2, 0.3])
ax.set_yticklabels([0, '', '', 0.3])
plt.savefig(os.path.join(fig_folder, ntl[d_l[-1]]+'_CS_height.png'), transparent=True)

# do the stats for the above figure
S = np.full([len(data), len(ntl)], np.nan)
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        CS_height_Vm = (data[i]['CS_max_Vm'][data[i][ntl[l]+'_CS_bool']] -
                        data[i]['CS_start_Vm'][data[i][ntl[l]+'_CS_bool']])
        if CS_height_Vm.size > 0:
            S[i, l] = np.nanmedian(CS_height_Vm)
# remove extra recordings from cells
S = S[keep_cells, :]

# do the paired boot stats
num_b = 1000
p = np.full(len(ntl) - 1, np.nan)
d = np.full(len(ntl) - 1, np.nan)
for l in np.arange(len(ntl) - 1):
    dif = S[:, l+1] - S[:, 0]
    # remove nans
    dif = dif[~np.isnan(dif)]
    d[l], p[l] = boot_pair_t(dif, num_b)
    print(dif.size)
print(d)
print(p)

# do the paired boot stats for theta hyp and LIA dep cells only
num_b = 1000
# theta hyp cells
dif = S[:, 1][theta_cell_p < 0.05] - S[:, 0][theta_cell_p < 0.05]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)
# LIA dep cells
dif = S[:, 2][LIA_cell_p > 0.95] - S[:, 0][LIA_cell_p > 0.95]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)
## LIA nondep cells
#dif = S[:, 2][LIA_cell_p < 0.95] - S[:, 0][LIA_cell_p < 0.95]
## remove nans
#dif = dif[~np.isnan(dif)]
#d, p = boot_pair_t(dif, num_b)
#print(dif.size)
#print(d)
#print(p)

# descriptive numbers
l = 1
np.sum(~np.isnan(S[:, l]))
np.nanmedian(S[:, l])
np.nanstd(S[:, l])
MADAM(S[:, l], np.nanmedian(S[:, l]))

# %% make CS features figures - after-CS hyperpolarization

# prep numbers for mean hist CS relative ahp - divided between states
bins = np.arange(-25, 10, 2)
H = np.full([len(data), bins.size-1, len(ntl)], np.nan)
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        CS_rel_ahp_Vm = (data[i]['CS_stop_Vm'][data[i][ntl[l]+'_CS_bool']] -
                        data[i]['CS_start_Vm'][data[i][ntl[l]+'_CS_bool']])
        H[i, :, l] = np.histogram(CS_rel_ahp_Vm, bins=bins)[0]
        # normalize to total number of CS
        H[i, :, l] = H[i, :, l]/np.sum(H[i, :, l])
# remove extra recordings from cells
H = H[keep_cells, :, :]
# define the 95% CI for each bin by randomly selecting (with replacement) over cells
H_mean = np.full([H.shape[1], H.shape[2]], np.nan)
H_CI_high = np.full([H.shape[1], H.shape[2]], np.nan)
H_CI_low = np.full([H.shape[1], H.shape[2]], np.nan)
CI_perc = 95
num_b = 1000
for l in np.arange(len(ntl)):
    real_H, CI_high, CI_low = CI_avg_hist(H[:, :, l], num_b, CI_perc)
    H_mean[:, l] = real_H
    H_CI_high[:, l] = CI_high
    H_CI_low[:, l] = CI_low
# plot the mean hist isi
# create a figure with axes of defined size
fig = plt.figure(figsize=[2, 2])
# The first items are for padding and the second items are for the axes.
# sizes are in inch.
h = [Size.Fixed(0.5), Size.Fixed(1.2)]
v = [Size.Fixed(0.5), Size.Fixed(1.2)]
divider = Divider(fig, (0.0, 0.0, 1., 1.), h, v, aspect=False)
ax = Axes(fig, divider.get_position())
ax.set_axes_locator(divider.new_locator(nx=1, ny=1))
fig.add_axes(ax)
for l in d_l:
    ax.plot(bins[:-1], H_mean[:, l], color=c_state_hist[l], zorder=2)
    ax.fill_between(bins[:-1], H_CI_low[:, l], H_CI_high[:, l],
                    facecolor=c_state_fill[l], linewidth=0, zorder=1, alpha=0.25)
#ax.set_xlabel('CS relative afterhyperpolarization (mV)')
#ax.set_ylabel('proportion of CS')
ax.set_xlim([-20, 3])
ax.set_xticks([-20, -10, 0])
ax.set_xlabel('hyperpolarization (mV)')
ax.set_ylim([0, 0.5])
ax.set_yticks([0, 0.25, 0.5])
ax.set_yticklabels([0, '', 0.5])
plt.savefig(os.path.join(fig_folder, ntl[d_l[-1]]+'_CS_ahp.png'), transparent=True)

# do the stats for the above figure
S = np.full([len(data), len(ntl)], np.nan)
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        CS_rel_ahp_Vm = (data[i]['CS_stop_Vm'][data[i][ntl[l]+'_CS_bool']] -
                        data[i]['CS_start_Vm'][data[i][ntl[l]+'_CS_bool']])
        if CS_rel_ahp_Vm.size > 0:
            S[i, l] = np.nanmedian(CS_rel_ahp_Vm)
# remove extra recordings from cells
S = S[keep_cells, :]

## do the friedman test (nonparametric repeated measures anova)
## remove cells that have any nans
#S_nonan = S[np.all(~np.isnan(S), axis=1), :]
#X2, p_fried = stats.friedmanchisquare(S_nonan[:, 0], S_nonan[:, 1], S_nonan[:, 2])
#X2, p_fried = stats.friedmanchisquare(S[:, 0], S[:, 1], S[:, 2])

# do the paired boot stats
num_b = 1000
p = np.full(len(ntl) - 1, np.nan)
d = np.full(len(ntl) - 1, np.nan)
for l in np.arange(len(ntl) - 1):
    dif = S[:, l+1] - S[:, 0]
    # remove nans
    dif = dif[~np.isnan(dif)]
    d[l], p[l] = boot_pair_t(dif, num_b)
    print(dif.size)
print(d)
print(p)

# do the paired boot stats for theta hyp and LIA dep cells only
num_b = 1000
# theta hyp cells
dif = S[:, 1][theta_cell_p < 0.05] - S[:, 0][theta_cell_p < 0.05]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)
# LIA dep cells
dif = S[:, 2][LIA_cell_p > 0.95] - S[:, 0][LIA_cell_p > 0.95]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)
# theta nonhyp cells
dif = dif = S[:, 1][theta_cell_p > 0.05] - S[:, 0][theta_cell_p > 0.05]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)

# descriptive numbers
l = 1
np.sum(~np.isnan(S[:, l]))
np.nanmedian(S[:, l])
np.nanstd(S[:, l])
MADAM(S[:, l], np.nanmedian(S[:, l]))

#%% CS-based stats for the average histograms

measure = 'num_sp'
measure = 'CS_dur'
measure = 'CS_height_Vm'
measure = 'CS_rel_ahp_Vm'

num_b = 1000
g0 = CS_ntl[0][measure]
g1 = CS_ntl[1][measure]
g2 = CS_ntl[2][measure]
groups_list = [g0, g1, g2]
real_F, p_boot = boot_anova(groups_list, num_b)
# try the stats test again with a kruskal-wallace (nonparametric 1-way anova)
H, p_kw = stats.kruskal(g0, g1, g2, nan_policy='omit')

# do the pairwise t-tests
boot_t(g0, g1, 1000)
boot_t(g0, g2, 1000)
boot_t(g1, g2, 1000)

# do the 2-sample Kolmogorov–Smirnov test (good for bimodal distributions?)
stats.ks_2samp(g0, g1)
stats.ks_2samp(g0, g2)
stats.ks_2samp(g1, g2)

# some numbers from the histogram
l = 1
CS_ntl[l][measure].size
np.nanmedian(CS_ntl[l][measure])
np.nanstd(CS_ntl[l][measure])
MADAM(CS_ntl[l][measure], np.nanmedian(CS_ntl[l][measure]))

# %% make CS rate and index figures

keep_cells = [isinstance(d['cell_id'], int) for d in data]
theta_cell_p = np.array([d['theta_cell_p'] for d in data])[keep_cells]
LIA_cell_p = np.array([d['LIA_cell_p'] for d in data])[keep_cells]

## find which cells have a significant change - boot
#anova_cells = np.full(len(data), np.nan)
#t_boot_cells = np.full([len(data), len(states)], np.nan)
#real_d_cells = np.full([len(data), len(states)], np.nan)
#num_b = 1000
#for i in np.arange(len(data)):
#    groups_list = [data[i]['nost_CS_rate'], data[i]['theta_CS_rate'],
#                   data[i]['LIA_CS_rate']]
#    real_F, anova_cells[i] = boot_anova(groups_list, num_b)
#    # if the anova is significant, do the adhoc stats
#    if anova_cells[i] < 0.05:
#        for l in np.arange(len(states)):
#            real_d_cells[i, l], t_boot_cells[i, l] = boot_t(groups_list[0], groups_list[l+1], num_b)
## remove extra recordings
#anova_cells = anova_cells[keep_cells]
#t_boot_cells = t_boot_cells[keep_cells, :]
#real_d_cells = real_d_cells[keep_cells, :]

# find which cells have a significant change - nonparametric stats
p_kw = np.full(len(data), np.nan)
p_mw = np.full([len(data), len(states)], np.nan)
num_b = 1000
for i in np.arange(len(data)):
    groups_list = [data[i]['nost_CS_rate'], data[i]['theta_CS_rate'],
                   data[i]['LIA_CS_rate']]
    try:
        H, p_kw[i] = stats.kruskal(groups_list[0], groups_list[1], groups_list[2],
                                   nan_policy='omit')
    except ValueError:
        p_kw[i] = np.nan
    # if the anova is significant, do the adhoc stats
    if p_kw[i] < 0.05:
        for l in np.arange(len(states)):
            U, p_mw[i, l] = stats.mannwhitneyu(groups_list[0], groups_list[l+1],
                                               alternative='two-sided')
# remove extra recordings
p_kw = p_kw[keep_cells]
p_mw = p_mw[keep_cells, :]

# each cells' average frequency of CS during theta, LIA, and no state
# Version: theta and LIA separate
# prep numbers
# only take first recording from each cell
S = np.full([len(data), len(ntl)], np.nan)
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        num_CS = np.sum([data[i][ntl[l]+'_CS_bool']])
        total_time = np.sum(data[i][ntl[l]+'_stop'] - data[i][ntl[l]+'_start'])
        S[i, l] = num_CS/total_time
# remove extra recordings from cells
S = S[keep_cells, :]
# plot the stack plot for cell values for each state
fig, ax = plt.subplots(1, figsize=[2.3, 2])
line_x = np.array([1.75, 3.25])
bar_x = np.array([1, 4])
y = S[:, d_l]
for i in np.arange(y.shape[0]):
    ax.plot(line_x, y[i, :], color=c_lgry, zorder=1)
    if d_l == [0, 1]:
        if theta_cell_p[i] < 0.05:
            ax.plot(line_x, y[i, :], color=rgb2hex(c_hyp), zorder=2)
        if theta_cell_p[i] > 0.95:
            ax.plot(line_x, y[i, :], color=rgb2hex(c_dep), zorder=2)
    elif d_l == [0, 2]:
        if LIA_cell_p[i] < 0.05:
            ax.plot(line_x, y[i, :], color=rgb2hex(c_hyp), zorder=2)
        if LIA_cell_p[i] > 0.95:
            ax.plot(line_x, y[i, :], color=rgb2hex(c_dep), zorder=2)
for l in np.arange(y.shape[1]):
    # remove nans
    no_nan = y[:, l]
    no_nan = no_nan[~np.isnan(no_nan)]
    bp = ax.boxplot(no_nan, sym='', patch_artist=True,
                         whis=[5, 95], widths=0.75, positions=[bar_x[l]])
    for element in ['boxes', 'whiskers', 'fliers', 'means', 'medians', 'caps']:
        plt.setp(bp[element], color=c_state[d_l[l]], linewidth=1.5)
    for patch in bp['boxes']:
        patch.set(facecolor=c_wht)
ax.set_xticks(bar_x)
ax.xaxis.set_tick_params(length=0)
ax.set_xticklabels(['unlabeled', 'theta'])
ax.set_yticks([0, 0.2, 0.4, 0.6, 0.8])
ax.set_yticklabels([0, '', 0.4, '', 0.8])
ax.set_ylabel('Cs rate (Hz)')
ax.set_xlim([0, bar_x[1]+1])
ax.spines['bottom'].set_visible(False)
fig.tight_layout()
plt.savefig(os.path.join(fig_folder, ntl[d_l[-1]]+'_CS_rate.png'), transparent=True)

# do the paired boot stats
num_b = 1000
p = np.full(len(ntl) - 1, np.nan)
d = np.full(len(ntl) - 1, np.nan)
for l in np.arange(len(ntl) - 1):
    dif = S[:, l+1] - S[:, 0]
    # remove nans
    dif = dif[~np.isnan(dif)]
    d[l], p[l] = boot_pair_t(dif, num_b)
print(d)
print(p)

# do the paired boot stats for theta hyp and LIA dep cells only
num_b = 1000
# theta hyp cells
dif = S[:, 1][theta_cell_p < 0.05] - S[:, 0][theta_cell_p < 0.05]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)
# LIA dep cells
dif = S[:, 2][LIA_cell_p > 0.95] - S[:, 0][LIA_cell_p > 0.95]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)

# descriptive numbers
l = 1
np.sum(~np.isnan(S[:, l]))
np.nanmedian(S[:, l])
np.nanstd(S[:, l])
MADAM(S[:, l], np.nanmedian(S[:, l]))

# find which cells have a significant change - nonparametric stats
p_kw = np.full(len(data), np.nan)
p_mw = np.full([len(data), len(states)], np.nan)
num_b = 1000
for i in np.arange(len(data)):
    groups_list = [data[i]['nost_CS_perc'], data[i]['theta_CS_perc'],
                   data[i]['LIA_CS_perc']]
    # do the kruskall-wallace if not all the CS_perc values are nan
    if ~np.all(np.isnan(np.concatenate(groups_list))):
        try:
            H, p_kw[i] = stats.kruskal(groups_list[0], groups_list[1], groups_list[2],
                                       nan_policy='omit')
        except ValueError:
            p_kw[i] = np.nan
    # if the anova is significant, do the adhoc stats
    if p_kw[i] < 0.05:
        for l in np.arange(len(states)):
            # remove nans before running the test
            g0 = groups_list[0]
            g0 = g0[~np.isnan(g0)]
            g1 = groups_list[l+1]
            g1 = g1[~np.isnan(g1)]
            U, p_mw[i, l] = stats.mannwhitneyu(g0, g1,
                                               alternative='two-sided')
# remove extra recordings
p_kw = p_kw[keep_cells]
p_mw = p_mw[keep_cells, :]

# %% CS index

# each cells' CS index during theta, LIA, and no state
# Version: theta and LIA separate
# prep numbers
# only take first recording from each cell
S = np.full([len(data), len(ntl)], np.nan)
for i in np.arange(len(data)):
    for l in np.arange(len(ntl)):
        num_CS_spikes = np.sum(np.array([c.size for c in data[i]['CS_ind']])[data[i][ntl[l]+'_CS_bool']])
        total_spikes = np.sum([data[i][ntl[l]+'_spike_bool']])
        CS_perc = num_CS_spikes/total_spikes
        if CS_perc > 1:
            CS_perc = 1
        S[i, l] = CS_perc
# remove extra recordings from cells
S = S[keep_cells, :]
# plot the stack plot for cell values for each state
fig, ax = plt.subplots(1, figsize=[2.3, 2])
line_x = np.array([1.75, 3.25])
bar_x = np.array([1, 4])
y = S[:, d_l]
for i in np.arange(y.shape[0]):
    ax.plot(line_x, y[i, :], color=c_lgry, zorder=1)
    if d_l == [0, 1]:
        if theta_cell_p[i] < 0.05:
            ax.plot(line_x, y[i, :], color=rgb2hex(c_hyp), zorder=2)
        if theta_cell_p[i] > 0.95:
            ax.plot(line_x, y[i, :], color=rgb2hex(c_dep), zorder=2)
    elif d_l == [0, 2]:
        if LIA_cell_p[i] < 0.05:
            ax.plot(line_x, y[i, :], color=rgb2hex(c_hyp), zorder=2)
        if LIA_cell_p[i] > 0.95:
            ax.plot(line_x, y[i, :], color=rgb2hex(c_dep), zorder=2)
for l in np.arange(y.shape[1]):
    # remove nans
    no_nan = y[:, l]
    no_nan = no_nan[~np.isnan(no_nan)]
    bp = ax.boxplot(no_nan, sym='', patch_artist=True,
                         whis=[5, 95], widths=0.75, positions=[bar_x[l]])
    for element in ['boxes', 'whiskers', 'fliers', 'means', 'medians', 'caps']:
        plt.setp(bp[element], color=c_state[d_l[l]], linewidth=1.5)
    for patch in bp['boxes']:
        patch.set(facecolor=c_wht)
ax.set_xticks(bar_x)
ax.xaxis.set_tick_params(length=0)
ax.set_xticklabels(['unlabeled', 'theta'])
ax.set_yticks([0, 0.25, 0.5, 0.75, 1])
ax.set_yticklabels([0, '', 0.5, '', 1])
ax.set_ylabel('Cs index')
ax.set_xlim([0, bar_x[1]+1])
ax.spines['bottom'].set_visible(False)
fig.tight_layout()
plt.savefig(os.path.join(fig_folder, ntl[d_l[-1]]+'_CS_index.png'), transparent=True)

# do the paired boot stats
num_b = 1000
p = np.full(len(ntl) - 1, np.nan)
d = np.full(len(ntl) - 1, np.nan)
for l in np.arange(len(ntl) - 1):
    dif = S[:, l+1] - S[:, 0]
    # remove nans
    dif = dif[~np.isnan(dif)]
    d[l], p[l] = boot_pair_t(dif, num_b)
print(d)
print(p)

# do the paired boot stats for theta hyp and LIA dep cells only
num_b = 1000
# theta hyp cells
dif = S[:, 1][theta_cell_p < 0.05] - S[:, 0][theta_cell_p < 0.05]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)
# LIA dep cells
dif = S[:, 2][LIA_cell_p > 0.95] - S[:, 0][LIA_cell_p > 0.95]
# remove nans
dif = dif[~np.isnan(dif)]
d, p = boot_pair_t(dif, num_b)
print(dif.size)
print(d)
print(p)

# descriptive numbers
l = 1
np.sum(~np.isnan(S[:, l]))
np.nanmedian(S[:, l])
np.nanstd(S[:, l])
MADAM(S[:, l], np.nanmedian(S[:, l]))

# %% make figures - dCSI vs dVm

l = 0
state='theta'

# Event-based correlation between dVm and change in CS index
unique_cells = [isinstance(d['cell_id'], int) for d in data]
fig, ax = plt.subplots(1, figsize=[2.25, 2.25])
n = 0
for i in np.arange(len(data)):
    for j in np.arange(data[i][state+'_start'].size):
        x = data[i][state+'_dVm'][j]
        z = data[i][state+'_dVm_p'][j]
        # calculate the CS index in the before window
        CS_bef = np.logical_and(data[i][state+'_CS_spikes'][j] > states[l]['bef'],
                                data[i][state+'_CS_spikes'][j] < states[l]['bef'] + states[l]['samp_time'])
        CS_bef = np.sum(CS_bef)
        nonCS_bef = np.logical_and(data[i][state+'_nonCS_spikes'][j] > states[l]['bef'],
                                   data[i][state+'_nonCS_spikes'][j] < states[l]['bef'] + states[l]['samp_time'])
        nonCS_bef = np.sum(nonCS_bef)
        CSindex_bef = CS_bef/(CS_bef+nonCS_bef)
        # calculate the CS index in the after window
        CS_aft = np.logical_and(data[i][state+'_CS_spikes'][j] > states[l]['aft'],
                                data[i][state+'_CS_spikes'][j] < states[l]['aft'] + states[l]['samp_time'])
        CS_aft = np.sum(CS_aft)
        nonCS_aft = np.logical_and(data[i][state+'_nonCS_spikes'][j] > states[l]['aft'],
                                   data[i][state+'_nonCS_spikes'][j] < states[l]['aft'] + states[l]['samp_time'])
        nonCS_aft = np.sum(nonCS_aft)
        CSindex_aft = CS_aft/(CS_aft+nonCS_aft)
        if np.logical_and(CS_bef+nonCS_bef == 0, CS_aft+nonCS_aft == 0):
            y = np.nan
        else:
            y = CSindex_aft-CSindex_bef
        if np.isnan(y) == False:
            n = n+1
        if z > 0.05:
            ax.scatter(x, y, s=5, facecolors='none', edgecolors=c_mgry, alpha=1, zorder=1)
        elif x < 0:
            ax.scatter(x, y, s=5, facecolors=c_lhyp, edgecolors=c_lhyp, alpha=1, zorder=2)
        elif x > 0:
            ax.scatter(x, y, s=5, facecolors=c_ldep, edgecolors=c_ldep, alpha=1, zorder=2)
ax.axhline(0, linestyle='--', color=c_blk, zorder=1)
ax.axvline(0, linestyle='--', color=c_blk, zorder=1)
ax.set_ylim([-1.1, 1.1])
ax.set_xlim([-18, 18])
# cell-based dVm vs change in CS index
# prep numbers for dVm
all_dVm = np.array([d[state+'_mean_dVm'] for d in data])[[isinstance(d['cell_id'], int) for d in data]]
all_cell_p = np.array([d[state+'_cell_p'] for d in data])[[isinstance(d['cell_id'], int) for d in data]]
keep_cells = np.logical_or(np.isnan(all_dVm), np.isnan(all_cell_p))==0
all_dVm = all_dVm[keep_cells]
all_cell_p = all_cell_p[keep_cells]
cell_hyp_sig = all_dVm[all_cell_p < 0.05]
cell_hyp_no = all_dVm[(all_dVm < 0) & (all_cell_p >= 0.05)]
cell_dep_sig = all_dVm[all_cell_p > 0.95]
cell_dep_no = all_dVm[(all_dVm > 0) & (all_cell_p <= 0.95)]
# prep number for CS index
dCSI = np.full(len(data), np.nan)
for i in np.arange(len(data)):
    dCSI_cell = np.full(data[i][state+'_start'].size, np.nan)
    for j in np.arange(data[i][state+'_start'].size):
        # calculate the CS index in the before window
        CS_bef = np.logical_and(data[i][state+'_CS_spikes'][j] > states[l]['bef'],
                                data[i][state+'_CS_spikes'][j] < states[l]['bef'] + states[l]['samp_time'])
        CS_bef = np.sum(CS_bef)
        nonCS_bef = np.logical_and(data[i][state+'_nonCS_spikes'][j] > states[l]['bef'],
                                   data[i][state+'_nonCS_spikes'][j] < states[l]['bef'] + states[l]['samp_time'])
        nonCS_bef = np.sum(nonCS_bef)
        CSindex_bef = CS_bef/(CS_bef+nonCS_bef)
        # calculate the CS index in the after window
        CS_aft = np.logical_and(data[i][state+'_CS_spikes'][j] > states[l]['aft'],
                                data[i][state+'_CS_spikes'][j] < states[l]['aft'] + states[l]['samp_time'])
        CS_aft = np.sum(CS_aft)
        nonCS_aft = np.logical_and(data[i][state+'_nonCS_spikes'][j] > states[l]['aft'],
                                   data[i][state+'_nonCS_spikes'][j] < states[l]['aft'] + states[l]['samp_time'])
        nonCS_aft = np.sum(nonCS_aft)
        CSindex_aft = CS_aft/(CS_aft+nonCS_aft)
        if np.logical_and(CS_bef+nonCS_bef == 0, CS_aft+nonCS_aft == 0):
            dCSI_cell[j] = np.nan
        else:
            dCSI_cell[j] = CSindex_aft-CSindex_bef
    dCSI[i] = np.nanmean(dCSI_cell)
dCSI = dCSI[unique_cells]
dCSI = dCSI[keep_cells]
cell_hyp_sig_dCSI = dCSI[all_cell_p < 0.05]
cell_hyp_no_dCSI = dCSI[(all_dVm < 0) & (all_cell_p >= 0.05)]
cell_dep_sig_dCSI = dCSI[all_cell_p > 0.95]
cell_dep_no_dCSI = dCSI[(all_dVm > 0) & (all_cell_p <= 0.95)]
# add the cell dots on top
s_cell = 20
ax.scatter(cell_hyp_sig, cell_hyp_sig_dCSI, s=s_cell, facecolors=c_hyp,
           edgecolors=c_blk, zorder=3, alpha=1)
ax.scatter(cell_hyp_no, cell_hyp_no_dCSI, s=s_cell, facecolors='none',
           edgecolors=c_blk, zorder=3, alpha=1)
ax.scatter(cell_dep_sig, cell_dep_sig_dCSI, s=s_cell, facecolors=rgb2hex(c_dep),
           edgecolors=rgb2hex(c_blk), zorder=3, alpha=1)
ax.scatter(cell_dep_no, cell_dep_no_dCSI, s=s_cell, facecolors='none',
           edgecolors=c_blk, zorder=3, alpha=1)
ax.set_xlabel(r'$\Delta$'+' Vm (mV)')
ax.set_ylabel(r'$\Delta$'+' CS index')
fig.tight_layout()
plt.savefig(os.path.join(fig_folder, 'dCSI_vs_dVm_'+state+'.png'), transparent=True)


from selenium import webdriver
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.remote.webelement import WebElement
from selenium.common.exceptions import NoSuchElementException
from selenium.common.exceptions import TimeoutException
import time
import json
import os
from course import Course

# Web Driver configuration
PATH = "C:\Program Files (x86)\chromedriver.exe"
driver = webdriver.Chrome(PATH)
coursesList = []

# getting pages URLs
f = open("URL.txt", "r")
URLs = []
for x in f:
    URLs.append(x)
f.close()

# searching through each page from file and through each subpage (< 1 2 3 ... 7 >)
for URL in URLs:
    emptyPage = False # means that the page number is out of range and there is no more content on this page
    subpageCounter = 1
    while not emptyPage:
        print(URL+'&p='+str(subpageCounter))
        driver.get(URL+'&p='+str(subpageCounter))
        subpageCounter += 1
        try: # element with this class name is a big container for all smaller divs. If it is not present then there is no content on the page
            WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.CLASS_NAME, 'course-list--container--3zXPS')))
            container = driver.find_element_by_class_name('course-list--container--3zXPS')
            coursesBiggerDivs = container.find_elements_by_class_name('browse-course-card--link--3KIkQ')
            courses = container.find_elements_by_class_name('course-card--container--3w8Zm')
            driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
            counter = 0
            for course in courses: # each course we convert into an object of 'Course' class (data extraction)
                title = course.find_element_by_class_name('udlite-heading-md').text
                desc = course.find_element_by_class_name('udlite-text-sm').text
                author = course.find_element_by_class_name('udlite-text-xs').text

                try:
                    spanElement = course.find_element_by_css_selector('span.star-rating--rating-number--3lVe8')
                except NoSuchElementException:
                    ratings = 'Brak ocen'
                else:
                    ratings = spanElement.text

                try:
                    details = course.find_elements_by_css_selector('span.course-card--row--1OMjg')
                    courseLength = details[0].text
                    courseLevel = details[len(details)-1].text
                except NoSuchElementException:
                    print("Brak dodatkowych informacji")
                    courseLength = 'Brak informacji'
                    courseLevel = 'Brak informacji'

                try:
                    image = course.find_element_by_class_name('course-card--course-image--2sjYP')
                    ActionChains(driver).move_to_element(image).perform()
                    imageSourceURL = image.get_attribute('src')
                except NoSuchElementException:
                    print("Brak zdjęcia")
                    imageSourceURL = 'https://www.google.com/url?sa=i&url=https%3A%2F%2Fwww.smarthome.com.au%2Faeotec-z-wave-plug-in-smart-switch-6.html&psig=AOvVaw33Vx1wP6a3B3QAn_6WPe4A&ust=1602514347326000&source=images&cd=vfe&ved=0CAIQjRxqFwoTCNitsanlrOwCFQAAAAAdAAAAABAE'

                try:
                    priceDiv = course.find_element_by_css_selector('div.price-text--price-part--Tu6MH')
                    ActionChains(driver).move_to_element(priceDiv).perform()
                    spans = priceDiv.find_elements_by_tag_name('span')
                    price = spans[len(spans) - 1].text
                except NoSuchElementException:
                    price = 'Brak ceny'

                try:
                    courseLink = coursesBiggerDivs[counter].get_attribute('href')
                except NoSuchElementException:
                    courseLink = None

                counter += 1
                c = Course(title, desc, author, ratings, price, imageSourceURL, courseLength, courseLevel, courseLink)
                coursesList.append(c)

        except TimeoutException:
            print('[INFO] Ostatnia podstrona adresu URL')
            emptyPage = True

os.remove('objectsInJSON.txt')

for course in coursesList: #search through each course page and get some more specific information
    driver.get(course.URL)
    WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.CLASS_NAME, 'topic-menu')))
    topicDiv = driver.find_element_by_class_name('topic-menu')
    elements = topicDiv.find_elements_by_class_name('udlite-heading-sm')
    course.setCategory(elements[0].text)
    course.setSubcategory(elements[1].text)
    courseDescription = driver.find_element_by_class_name('styles--description--3y4KY')
    course.setExtendedDescription(courseDescription.get_attribute('innerHTML'))

    # write converted course object into output file
    string = course.makeJSON()
    with open('objectsInJSON.txt','a',encoding='utf-8') as file:
        json.dump(string, file, ensure_ascii=False)
        file.write("\n")

driver.quit()
file.close()


from selenium import webdriver
from selenium.webdriver.common.by import By
import time
import os

try:
    link = "http://suninjuly.github.io/file_input.html"
    browser = webdriver.Chrome()
    browser.get(link)

    elements = browser.find_elements(By.CSS_SELECTOR, ".form-control")
    for element in elements:
        if element.get_attribute('required') != None:
            element.send_keys("Мой ответ")

    print(os.getcwd())
    current_dir = os.path.realpath(os.path.dirname(__file__))
    file_path = os.path.join(current_dir, 'file.txt')
    print(file_path)
    element = browser.find_element(By.CSS_SELECTOR, "#file")
    element.send_keys(file_path)

    button = browser.find_element(By.CSS_SELECTOR, "button.btn")
    button.click()

finally:
    # ожидание чтобы визуально оценить результаты прохождения скрипта
    time.sleep(10)
    # закрываем браузер после всех манипуляций
    browser.quit()


# Imports
import sqlalchemy as sa
from sqlalchemy.ext.declarative import declarative_base
import users
from datetime import datetime

"""
Модуль для поиска атлетов по парамтрам пользователя
"""
# Variables
Base = declarative_base()

# Class definitions

class Athlette(Base):
    __tablename__ = "Athelete"
    id = sa.Column(sa.INTEGER, primary_key=True)
    age = sa.Column(sa.INTEGER)
    birthdate = sa.Column(sa.TEXT)
    gender = sa.Column(sa.TEXT)
    height = sa.Column(sa.REAL)
    name = sa.Column(sa.TEXT)
    weight = sa.Column(sa.INTEGER)
    gold_medals = sa.Column(sa.INTEGER)
    silver_medals = sa.Column(sa.INTEGER)
    bronze_medals = sa.Column(sa.INTEGER)
    total_medals = sa.Column(sa.INTEGER)
    sport = sa.Column(sa.TEXT)
    country = sa.Column(sa.TEXT)

# Function definitions
def search_id(id, session):
    query_str = session.query(Athlette).filter(Athlette.id == id).first()
    usr = f"{query_str}"
    return usr

def height_compare(id, session, bcolors):
    """
    Сравнение роста атлетов с пользовательским
    """
    # берем из базы рост пользователя
    usr_query = session.query(users.User).filter(users.User.id == id).first()
    usr_height = usr_query.height

    # ищем атлетов по росту пользователя
    ath_query = session.query(Athlette).filter(
        Athlette.height == usr_height)

    ath_count = ath_query.count()
    ath_found = ath_query.all()

    # выводим содержимое объектов Athlete, если найдены
    res = ""
    if ath_found:
        for ath in ath_found:
            res += f"  {ath.name}, {ath.sport} \n"
        res = f"{res}\n  Всего атлетов с ростом {ath.height} метра: {ath_count}"
    else:
        print(bcolors.FAIL +
              f"\nERROR: Атлет с ростом {usr_height}m не найден" + bcolors.ENDC)

    return res

def bday_compare(id, session):
    """
    Ищем атлета, наиболее близкого по дате рождения к пользователю
    """
    dt_format = '%Y-%m-%d'
    usr_query = session.query(users.User).filter(users.User.id == id).first()
    user_bday_str = usr_query.birthdate
    user_bday_dt_obj = datetime.strptime(user_bday_str, dt_format)

    ath_query_all_obj = session.query(Athlette).all()

    ath_bday_all_dt_list = list()
    for ath in ath_query_all_obj:
        ath_bday_all_dt_list.append(
            datetime.strptime(ath.birthdate, dt_format))

    closest_bday_dt_obj = ath_bday_all_dt_list[min(range(len(ath_bday_all_dt_list)),
                                                   key=lambda i: abs(ath_bday_all_dt_list[i]-user_bday_dt_obj))]

    # выбираем всех атлетов по самой ближней дате рождения
    closest_bday_str = closest_bday_dt_obj.strftime(dt_format)
    ath_query_bday_query = session.query(Athlette).filter(
        Athlette.birthdate == closest_bday_str)

    # берем из базы данные и считаем
    ath_bday_obj = ath_query_bday_query.all()
    ath_bday_count = ath_query_bday_query.count()

    # формируем возврат
    res = ""
    for ath in ath_bday_obj:
        res = f"{res}\n  {ath.name}, д.р.: {ath.birthdate}, {ath.sport}"

    return res

if __name__ == "__main__":
    print("ERROR: Запуск скрипта через выполнение модуля start.py \n")

# DEBUG
# print('Info: Module find_athlete.py - imported')


from summarize import *
from rbm_dae.deepAE import *
import rouge

def summarize_sentence_vectors(df, vector_set):
    """
    Function applying the summarization function to get the ranked sentences.
    Parameters:
        df: dataframe containing the data to summarize
        vector_set: the column name of the vector set to rank on
    Returns the ranked sentences
    """
    print('summarizing sentence vectors..')
    sentence_vectors = df[vector_set].tolist()
    sentence_vectors = np.array(sentence_vectors)
    #Create a list of ranked sentences.
    ranked_sentences = summarize_emails(df, sentence_vectors[0])
    # display_summary(df, ranked_sentences)
    return ranked_sentences

def summarize_autoencoder_vectors(df, net, vector_set):
    """
    Function applying the autoencoder to the df vectors and the applying the summarization function to get the ranked sentences.
    Parameters:
        df: dataframe containing the data to summarize
        net: trained autoencoder
        vector_set: the column name of the vector set to rank on
    Returns the ranked sentences
    """
    print('summarizing autoencoder sentence vectors..')
    sentence_vectors = df[vector_set].tolist()
    torch_vectors = torch.tensor(sentence_vectors[0], dtype=torch.float32)
    output_vectors = net(torch_vectors)
    #Create a list of ranked sentences.
    ranked_sentences = summarize_emails(df, output_vectors, True)
    # display_summary(df, ranked_sentences)
    return ranked_sentences

def evaluate_rankings(df_train, df_test, target, sum_lens, corpus_ae=True, vector_set='sentence_vectors'):
    """
    Funtion to evaluate the returned summaries. the summaries are created baased on the raw sentence vectors and the autoencoder vectors
    Parameters:
        df_train: dataframe with the training data
        df_test: dataframe with the test data
        target: string containing the column that should be used as the summary reference
        sum_len: An array holding the number of sentences to include in the summary
        corpus_as: Boolean deciding wether to train the autoencoder on the entire corpus or on each document
        vector_set: column name of the column with the sentence vectors (can be glove vectors or tf vectors)
    Returns: the scores for the rouge parameters (3D matrix)
    """
    evaluator = rouge.Rouge()
    #create and train the autoencoder (see autoencoder module)
    net = None
    if corpus_ae:
        net = train_autoencoder(df_train, vector_set)

    # loop through all docs in the corpus
    print('evaluating summaries..')
    df_len = int(df_test.shape[0])
    sum_scores = np.zeros((len(sum_lens), 3, 3, df_len))
    ae_sum_scores = np.zeros((len(sum_lens), 3, 3, df_len))
    curr_row = 0
    for index, row in df_test.iterrows():
        print('iteration: ', index)
        df_c = pd.DataFrame([row])
        df_c['body'].iloc[0]
        # Only proceed if the vectors of the current row are of correct dimensions (not [])
        if len(df_c[vector_set].tolist()[0]) > 0:
            # train AE on the current document only
            if not corpus_ae :
                net = train_autoencoder(df_c, vector_set)
            reference = df_c[target].iloc[0] # reference that we score against (could be summary or subject)!
            print('reference: ', reference)
            # get the ranked sentences for the original and the ae modified sentence vectors
            ranked_sentences = summarize_sentence_vectors(df_c, vector_set)
            ranked_ae_sentences = summarize_autoencoder_vectors(df_c, net, vector_set)
            # collecting the scores for the specified summary lengths
            for s_len in sum_lens:
                print('s_len: ', s_len)
                if len(ranked_sentences) >= s_len:
                    # get the top ranked sentences
                    sum = []
                    sum_ae = []
                    for i in range(s_len):
                        sum.append(ranked_sentences[i][2])
                        sum_ae.append(ranked_ae_sentences[i][2])

                    sum_str = ' '.join(sum)
                    sum_ae_str = ' '.join(sum_ae)
                    print('summary: ', sum_str)
                    print('ae summary: ', sum_ae_str)

                    # get the ROUGE scores for the ranked sentences and add to plot data
                    sum_score = evaluator.get_scores(sum_str, reference)
                    sum_ae_score = evaluator.get_scores(sum_ae_str, reference)
                    sum_scores[s_len-1, 0, 0, curr_row] = sum_score[0]['rouge-1']['f']
                    sum_scores[s_len-1, 0, 1, curr_row] = sum_score[0]['rouge-1']['p']
                    sum_scores[s_len-1, 0, 2, curr_row] = sum_score[0]['rouge-1']['r']
                    sum_scores[s_len-1, 1, 0, curr_row] = sum_score[0]['rouge-2']['f']
                    sum_scores[s_len-1, 1, 1, curr_row] = sum_score[0]['rouge-2']['p']
                    sum_scores[s_len-1, 1, 2, curr_row] = sum_score[0]['rouge-2']['r']
                    sum_scores[s_len-1, 2, 0, curr_row] = sum_score[0]['rouge-l']['f']
                    sum_scores[s_len-1, 2, 1, curr_row] = sum_score[0]['rouge-l']['p']
                    sum_scores[s_len-1, 2, 2, curr_row] = sum_score[0]['rouge-l']['r']

                    ae_sum_scores[s_len-1, 0, 0, curr_row] = sum_ae_score[0]['rouge-1']['f']
                    ae_sum_scores[s_len-1, 0, 1, curr_row] = sum_ae_score[0]['rouge-1']['p']
                    ae_sum_scores[s_len-1, 0, 2, curr_row] = sum_ae_score[0]['rouge-1']['r']
                    ae_sum_scores[s_len-1, 1, 0, curr_row] = sum_ae_score[0]['rouge-2']['f']
                    ae_sum_scores[s_len-1, 1, 1, curr_row] = sum_ae_score[0]['rouge-2']['p']
                    ae_sum_scores[s_len-1, 1, 2, curr_row] = sum_ae_score[0]['rouge-2']['r']
                    ae_sum_scores[s_len-1, 2, 0, curr_row] = sum_ae_score[0]['rouge-l']['f']
                    ae_sum_scores[s_len-1, 2, 1, curr_row] = sum_ae_score[0]['rouge-l']['p']
                    ae_sum_scores[s_len-1, 2, 2, curr_row] = sum_ae_score[0]['rouge-l']['r']
            curr_row += 1

    sum_scores = sum_scores[:, :, :, 0:curr_row]
    ae_sum_scores = ae_sum_scores[:, :, :, 0:curr_row]
    return sum_scores, ae_sum_scores

# calculating averages
def analyze_and_plot_rouge_scores(sum_scores, ae_sum_scores, metric, dataset_name, summary_len):
    avg_scores = np.mean(sum_scores)
    avg_scores_ae = np.mean(ae_sum_scores)
    print(dataset_name)
    print('Summary length: ', summary_len)
    raw_mean = 'Mean ' + metric + ' for raw vectors: ' + str(round(avg_scores, 3))
    dae_mean = 'Mean ' + metric + ' for DAE vectors: ' + str(round(avg_scores_ae, 3))
    print(raw_mean)
    print(dae_mean)

    # Add to plot graphs for the extracted sentences
    """
    x = np.arange(len(sum_scores)).tolist()
    label_1 = "Raw " + metric
    label_2 = "AE vector " + metric
    plt.plot(x, sum_scores.tolist(), label = label_1)
    plt.plot(x, ae_sum_scores.tolist(), label = label_2)
    plt.xlabel('Sentence')
    plt.ylabel('ROUGE score')
    title = "ROUGE " +metric + " for raw (mean: " + str(round(avg_scores, 3)) +") and AE (mean: "+str(round(avg_scores_ae, 3)) +") for " + dataset_name
    plt.title(title)
    plt.legend()
    plt.show()
    """

def evaluate_bc3():
    """
    Base function to run and plot the ROUGE scores for the bc3 dataset
    """
    BC3_PICKLE_LOC  = "./final_data/BC3_127.pkl"
    BC3_df = pd.read_pickle(BC3_PICKLE_LOC)
    # df contains 127 rows that all have df_vectors representation!
    # Split into training and test set
    BC3_df_train = BC3_df.iloc[:117]
    BC3_df_test = BC3_df.iloc[117:]

    # evaluate on 'summary' or 'subject'
    target = 'summary'
    summary_len = [1]
    # can set to use the df vectors ('df_vectors') or the glove vectors ('sentence_vectors')
    corpus_ae = True
    vector_set = 'sentence_vectors'                #df_vectors
    sum_scores, ae_sum_scores = evaluate_rankings(BC3_df_train, BC3_df_test, target, summary_len, corpus_ae, vector_set)
    plot_all_scores(sum_scores, ae_sum_scores, 'bc3 dataset', summary_len[0])

def evaluate_spotify():
    """
    Base function to run and plot the ROUGE scores for the spotify dataset
    """
    SPOTIFY_PICKLE_TRAIN_LOC  = "./final_data/spotify_train_422.pkl"
    SPOTIFY_PICKLE_TEST_LOC  = "./final_data/spotify_test_45.pkl"
    df_train = pd.read_pickle(SPOTIFY_PICKLE_TRAIN_LOC)
    df_test = pd.read_pickle(SPOTIFY_PICKLE_TEST_LOC)

    # section to get the summary for a specidic episode
    # df_sent = df_train.loc[df_train['episode_id'] == '7DoDuJE4sCBu2jJlOgCrwA']
    # df_test = df_sent

    target = 'episode_desc'
    summary_len = [1]
    corpus_ae = True # if false, the autoencoder is only trained on the sentences in the current document
    # can set to use the df vectors (t-idf) ('df_vectors') or the glove vectors ('sentence_vectors')
    vector_set = 'sentence_vectors'

    sum_scores, ae_sum_scores = evaluate_rankings(df_train, df_test, target, summary_len, corpus_ae, vector_set)
    plot_all_scores(sum_scores, ae_sum_scores, 'spotify dataset', summary_len[0])

def plot_all_scores(sum_scores, ae_sum_scores, dataset, summary_len):
    """
    Base function to plot ROUGE scores.
    Parameters:
        - sum_scores: Matirx of scores for the raw vectors.
        - as_sum_scores: Matrix of scores for the vectors produced by autoencoder
    """
    analyze_and_plot_rouge_scores(sum_scores[0][0][0], ae_sum_scores[0][0][0], 'rouge-1 f-score', dataset, summary_len)
    analyze_and_plot_rouge_scores(sum_scores[0][0][1], ae_sum_scores[0][0][1], 'rouge-1 precision', dataset, summary_len)
    analyze_and_plot_rouge_scores(sum_scores[0][0][2], ae_sum_scores[0][0][2], 'rouge-1 recall', dataset, summary_len)

    # plot rouge-2 scores:
    analyze_and_plot_rouge_scores(sum_scores[0][1][0], ae_sum_scores[0][1][0], 'rouge-2 f-score', dataset, summary_len)
    analyze_and_plot_rouge_scores(sum_scores[0][1][1], ae_sum_scores[0][1][1], 'rouge-2 precision', dataset, summary_len)
    analyze_and_plot_rouge_scores(sum_scores[0][1][2], ae_sum_scores[0][1][2], 'rouge-2 recall', dataset, summary_len)

    # plot rouge-l scores:
    analyze_and_plot_rouge_scores(sum_scores[0][2][0], ae_sum_scores[0][2][0], 'rouge-l f-score', dataset, summary_len)
    analyze_and_plot_rouge_scores(sum_scores[0][2][1], ae_sum_scores[0][2][1], 'rouge-l precision', dataset, summary_len)
    analyze_and_plot_rouge_scores(sum_scores[0][2][2], ae_sum_scores[0][2][2], 'rouge-l recall', dataset, summary_len)

def get_mean(sum_scores, ae_sum_scores):
    """
    Function to get the mean of a vector of scores.
    """
    avg_scores = np.mean(sum_scores)
    avg_scores_ae = np.mean(ae_sum_scores)
    return avg_scores, avg_scores_ae

def evaluate_sentence_length_performance(df_train, df_test, target, summary_len, corpus_ae, vector_set, dataset):
    """
    Function to cumpute the rouge scores for a range of summary lengths.
    """
    averages_p = np.zeros((summary_len, 2))
    averages_ae_p = np.zeros((summary_len, 2))
    averages_r = np.zeros((summary_len, 2))
    averages_ae_r = np.zeros((summary_len, 2))
    summary_lengths = [1, 2, 3, 4, 5, 6]
    sum_scores, ae_sum_scores = evaluate_rankings(df_train, df_test, target, summary_lengths, corpus_ae, vector_set)
    for i in range(1, summary_len):
        print('evaluating rankings for # sentences: ', i)
        for j in range(2): # for rouge-1 and rouge-2
                avg_score_p, avg_score_ae_p = get_mean(sum_scores[i-1][j][1], ae_sum_scores[i-1][j][1])
                avg_score_r, avg_score_ae_r = get_mean(sum_scores[i-1][j][2], ae_sum_scores[i-1][j][2])

                averages_p[i, j] = avg_score_p
                averages_ae_p[i, j] = avg_score_ae_p
                averages_r[i, j] = avg_score_r
                averages_ae_r[i, j] = avg_score_ae_r
    print('averages: ', averages_p)
    print('averages ae: ', averages_ae_p)
    averages_p = averages_p[1:].transpose()
    averages_ae_p = averages_ae_p[1:].transpose()
    averages_r = averages_r[1:].transpose()
    averages_ae_r = averages_ae_r[1:].transpose()
    return averages_p, averages_ae_p, averages_r, averages_ae_r

def plot_sentences(glove_averages, glove_averages_ae, df_averages, df_averages_ae, title, dataset):
    """
    Function to plot the mean scores vs sentence lengths for the different sentence encodings
    """
    x = np.arange(1,7).tolist()
    plt.plot(x, glove_averages.tolist(), label = "Glove vector")
    plt.plot(x, glove_averages_ae.tolist(), label = "Glove DAE vector")
    plt.plot(x, df_averages.tolist(), label = 'tf-idf vector')
    plt.plot(x, df_averages_ae.tolist(), label = 'tf-idf DAE vector')
    plt.xlabel('Number of sentences')
    plt.ylabel(title)
    t = title + ' for ' + dataset
    plt.title(t)
    plt.legend()
    plt.show()

def run_sentence_length_evaluation():
    """
    Main function to compute the mean scores for each summary length for the two datasets.
    """
    BC3_PICKLE_LOC  = "./final_data/BC3_127.pkl"
    BC3_df = pd.read_pickle(BC3_PICKLE_LOC)
    # df contains 127 rows that all have df_vectors representation!
    # Split into training and test set
    bc3_df_train = BC3_df.iloc[:117]
    bc3_df_test = BC3_df.iloc[117:]
    bc3_target = 'summary'

    SPOTIFY_PICKLE_TRAIN_LOC  = "./final_data/spotify_train_422.pkl"
    SPOTIFY_PICKLE_TEST_LOC  = "./final_data/spotify_test_45.pkl"
    s_df_train = pd.read_pickle(SPOTIFY_PICKLE_TRAIN_LOC)
    s_df_test = pd.read_pickle(SPOTIFY_PICKLE_TEST_LOC)
    s_target = 'episode_desc'

    summary_len = 7
    corpus_ae = True
    vector_set = 'sentence_vectors'
    df_vector_set = 'df_vectors'
    # metric = 0 # 0 = f-score, 1 = precision, 2 = recall
    bc3_glove_p, bc3_glove_ae_p, bc3_glove_r, bc3_glove_ae_r = evaluate_sentence_length_performance(bc3_df_train, bc3_df_test, bc3_target, summary_len, corpus_ae, vector_set, 'bc3 dataset')
    bc3_df_p, bc3_df_ae_p, bc3_df_r, bc3_df_ae_r = evaluate_sentence_length_performance(bc3_df_train, bc3_df_test, bc3_target, summary_len, corpus_ae, df_vector_set, 'bc3 dataset')
    plot_sentences(bc3_glove_p[0], bc3_glove_ae_p[0], bc3_df_p[0], bc3_df_ae_p[0], 'ROUGE-1 scores precision', 'BC3 dataset')
    plot_sentences(bc3_glove_p[1], bc3_glove_ae_p[1], bc3_df_p[1], bc3_df_ae_p[1], 'ROUGE-2 scores precision', 'BC3 dataset')
    plot_sentences(bc3_glove_r[0], bc3_glove_ae_r[0], bc3_df_r[0], bc3_df_ae_r[0], 'ROUGE-1 scores recall', 'BC3 dataset')
    plot_sentences(bc3_glove_r[1], bc3_glove_ae_r[1], bc3_df_r[1], bc3_df_ae_r[1], 'ROUGE-2 scores recall', 'BC3 dataset')

    s_glove_p, s_glove_ae_p, s_glove_r, s_glove_ae_r = evaluate_sentence_length_performance(s_df_train, s_df_test, s_target, summary_len, corpus_ae, vector_set, 'Spotify dataset')
    s_df_p, s_df_ae_p, s_df_r, s_df_ae_r = evaluate_sentence_length_performance(s_df_train, s_df_test, s_target, summary_len, corpus_ae, df_vector_set, 'Spotify dataset')
    plot_sentences(s_glove_p[0], s_glove_ae_p[0], s_df_p[0], s_df_ae_p[0], 'ROUGE-1 scores precision', 'Spotify dataset')
    plot_sentences(s_glove_p[1], s_glove_ae_p[1], s_df_p[1], s_df_ae_p[1], 'ROUGE-2 scores precision', 'Spotify dataset')
    plot_sentences(s_glove_r[0], s_glove_ae_r[0], s_df_r[0], s_df_ae_r[0], 'ROUGE-1 scores recall', 'Spotify dataset')
    plot_sentences(s_glove_r[1], s_glove_ae_r[1], s_df_r[1], s_df_ae_r[1], 'ROUGE-2 scores recall', 'Spotify dataset')

# evaluate_bc3()
evaluate_spotify()
# run_sentence_length_evaluation()


from django.shortcuts import render, redirect
from django.contrib import messages
from .models import User
from .forms import RegisterForm, LoginForm
from .utils import require_login

def login_page(request):
	context = {"reg_form": RegisterForm(), "login_form": LoginForm()}
	return render(request, "users/login.html", context)

def login(request):
	curr_user = User.user_manager.login(request.POST["email_address"], request.POST["password"])
	if not curr_user:
		messages.error(request, "E-mail or password incorrect")
		return redirect("login_page")
	else:
		request.session["curr_user"] = curr_user.id
		return redirect("dashboard")

def register(request):
	registered, guy_or_errors = User.user_manager.register(request.POST)

	if not registered:
		for error in guy_or_errors: messages.error(request, error)
		return redirect("login_page")
	else:
		request.session["curr_user"] = guy_or_errors.id
		return redirect("dashboard")

def log_off(request):
	request.session.clear()
	return redirect("login_page")

@require_login
def dashboard(request, curr_user):
	context = {
		"curr_user": curr_user,
		"users": User.user_manager.all(),
	}
	return render(request, "users/dashboard.html", context)

@require_login
def show(request, curr_user, id):
	print("show page")
	context = {
		"curr_user": curr_user,
		"user": User.user_manager.get(id=id),
	}

	return render(request, "users/show.html", context)

@require_login
def edit(request, curr_user, id):
	context = {
		"curr_user": curr_user,
		"user": User.user_manager.get(id=id),
	}
	return render(request, "users/edit.html", context)

@require_login
def update(request, curr_user, id):
	# Logic to check if curr_user is admin or user being updated
	if not (curr_user.admin or curr_user.id == int(id)):
		print(curr_user.id, id, curr_user.id == id)
		return redirect("/dashboard")

	# Logic to actually update
	errors = User.user_manager.update(id, request.POST)

	if errors:
		for error in errors:
			messages.error(request, error)
		return redirect("edit", id=id)
	else:
		return redirect("show", id=id)


import re
import os
import sys
import nltk
import json
import wandb
import joblib
import datasets
import numpy as np
import pandas as pd
from time import process_time
from nltk import word_tokenize
from nltk.stem import WordNetLemmatizer
from sklearn.svm import LinearSVC
from sklearn.pipeline import make_pipeline
from sklearn.feature_extraction.text import TfidfVectorizer

np.random.seed(42)

class LemmaTokenizer:
        ignore_tokens = [',', '.', ';', ':', '"', '``', "''", '`']
        def __init__(self):
            self.wnl = WordNetLemmatizer()
        def __call__(self, doc):
            return [self.wnl.lemmatize(t) for t in word_tokenize(doc) if t not in self.ignore_tokens]

def prepare_dataset(data_folder, label2id, data_types, max_length):
    def combine_data(example):
        temp_text = ""
        for data_type in data_types:
            temp_text += example[data_type] + " "
        example["text"] = temp_text[:max_length]
        return example

    dataset = datasets.load_from_disk(data_folder + "dataset/")
    dataset = dataset["train"]
    dataset_encoded = dataset.class_encode_column("category")
    dataset_aligned = dataset_encoded.align_labels_with_mapping(label2id, "category")
    dataset_cleaned = dataset_aligned.map(combine_data)

    dataset = dataset_cleaned.remove_columns(["title", "body"])
    dataset = dataset.rename_column("category", "label")
    return dataset

def main():
    hps = {
        "data_types": ["title", "body"],
        "loss_function": "squared_hinge",
        "ngram_range": 3,
        "max_length": 512,
    }

    wandb_id = wandb.util.generate_id()

    run = wandb.init(
        project="DMOZ-classification",
        config=hps,
        job_type="training",
        name="SVM_DMOZ_" + str(wandb_id),
        tags=["SVM", "DMOZ"],
    )

    data_folder = "/ceph/csedu-scratch/other/jbrons/thesis-web-classification/"

    id2label = {0: "Arts", 1: "Business", 2: "Computers", 3: "Health", 4: "Home", 5: "News", 6: "Recreation", 7: "Reference", 8: "Science", 9: "Shopping", 10: "Society", 11: "Sports", 12: "Games"}
    label2id = {v: k for k, v in id2label.items()}
    labels = label2id.keys()

    dataset = prepare_dataset(data_folder, label2id, hps["data_types"], hps["max_length"])
    X_train, y_train = dataset["text"], dataset["label"]

    tokenizer=LemmaTokenizer()

    pipeline = make_pipeline(
        TfidfVectorizer(
            ngram_range=(1, hps["ngram_range"]),
            tokenizer=tokenizer,
            token_pattern=None
        ),
        LinearSVC(loss=hps["loss_function"])
    )

    t0 = process_time()
    pipeline.fit(X_train, y_train)
    training_time = process_time() - t0

    print("Training time {:5.2f}s for {:0d} samples.".format(training_time, len(y_train)))
    run.summary["training_time"] = training_time

    filename = data_folder + "models/SVM/model.pkl"
    joblib.dump(pipeline, filename, compress=3)

    model_artifact = wandb.Artifact(
        name="model_SVM_DMOZ",
        type="model"
    )
    model_artifact.add_file(thesis_folder + "models/SVM/model.pkl")
    run.log_artifact(model_artifact)

if __name__ == "__main__":
    main()


"""Tutorial 8. Putting two plots on the screen

This tutorial sets up for showing how Chaco allows easily opening multiple
views into a single dataspace, which is demonstrated in later tutorials.
"""

from scipy import arange
from scipy.special import jn

from enable.api import ComponentEditor
from traits.api import HasTraits, Instance
from traitsui.api import Item, View

from chaco.api import create_line_plot, HPlotContainer
from chaco.tools.api import PanTool

class PlotExample(HasTraits):
    container = Instance(HPlotContainer)

    traits_view = View(
        Item(
            "container",
            editor=ComponentEditor(),
            show_label=False,
            width=800,
            height=600,
        ),
        title="Chaco Tutorial",
    )

    def _container_default(self):
        x = arange(-5.0, 15.0, 20.0 / 100)

        y = jn(0, x)
        left_plot = create_line_plot(
            (x, y), bgcolor="white", add_grid=True, add_axis=True
        )
        left_plot.tools.append(PanTool(left_plot))
        self.left_plot = left_plot

        y = jn(1, x)
        right_plot = create_line_plot(
            (x, y), bgcolor="white", add_grid=True, add_axis=True
        )
        right_plot.tools.append(PanTool(right_plot))
        right_plot.y_axis.orientation = "right"
        self.right_plot = right_plot

        # Tone down the colors on the grids
        right_plot.hgrid.line_color = (0.3, 0.3, 0.3, 0.5)
        right_plot.vgrid.line_color = (0.3, 0.3, 0.3, 0.5)
        left_plot.hgrid.line_color = (0.3, 0.3, 0.3, 0.5)
        left_plot.vgrid.line_color = (0.3, 0.3, 0.3, 0.5)

        container = HPlotContainer(spacing=20, padding=50, bgcolor="lightgray")
        container.add(left_plot)
        container.add(right_plot)
        return container

demo = PlotExample()

if __name__ == "__main__":
    demo.configure_traits()