import numpy as np
import pandas as pd
from scipy.signal import correlate
import Edmond.VR_grid_analysis.analysis_settings as Settings
import PostSorting.parameters
import PostSorting.vr_stop_analysis
import PostSorting.vr_time_analysis
import PostSorting.vr_make_plots
import PostSorting.vr_cued
import PostSorting.theta_modulation
import PostSorting.vr_spatial_data
from matplotlib.markers import TICKDOWN
from Edmond.VR_grid_analysis.remake_position_data import syncronise_position_data
from Edmond.VR_grid_analysis.vr_grid_stability_plots import add_hit_miss_try3, add_avg_track_speed, get_avg_correlation, \
    get_reconstructed_trial_signal, plot_firing_rate_maps_per_trial_by_hmt_aligned, plot_firing_rate_maps_per_trial_by_hmt_aligned_other_neuron, get_shifts
from scipy import stats
from scipy import signal
from scipy.interpolate import interp1d
from astropy.convolution import convolve, Gaussian1DKernel, Gaussian2DKernel
import os
import traceback
import warnings
import matplotlib.ticker as ticker
import sys
import Edmond.plot_utility2
import Edmond.VR_grid_analysis.hit_miss_try_firing_analysis
import settings
import matplotlib.pylab as plt
import matplotlib as mpl
import control_sorting_analysis
import PostSorting.post_process_sorted_data_vr
from astropy.timeseries import LombScargle
from Edmond.utility_functions.array_manipulations import *
from Edmond.VR_grid_analysis.vr_grid_cells import add_lomb_classifier
from joblib import Parallel, delayed
import multiprocessing
import open_ephys_IO
warnings.filterwarnings('ignore')
from scipy.stats.stats import pearsonr
from scipy.stats import shapiro
import samplerate
plt.rc('axes', linewidth=3)

def get_p_text(p, ns=True):

    if p is not None:
        if np.isnan(p):
            return " "
        if p<0.00001:
            return 'p < 1e'+('{:.1e}'.format(p)).split("e")[-1]
        if p<0.0001:
            return "****"
        elif p<0.001:
            return "***"
        elif p<0.01:
            return "**"
        elif p<0.05:
            return "*"
        elif ns:
            return "ns"
        else:
            return " "
    else:
        return " "

def add_avg_trial_speed(processed_position_data):
    avg_trial_speeds = []
    for trial_number in np.unique(processed_position_data["trial_number"]):
        trial_processed_position_data = processed_position_data[processed_position_data["trial_number"] == trial_number]
        speeds = np.asarray(trial_processed_position_data['speeds_binned_in_time'])[0]
        avg_speed = np.nanmean(speeds)
        avg_trial_speeds.append(avg_speed)
    processed_position_data["avg_trial_speed"] = avg_trial_speeds
    return processed_position_data

def add_hit_miss_try(processed_position_data, track_length):
    reward_zone_start = track_length-60-30-20
    reward_zone_end = track_length-60-30

    rewarded_processed_position_data = processed_position_data[(processed_position_data["rewarded"] == True)]
    minimum_speeds_in_rz = []
    speeds_in_rz = []
    for trial_number in np.unique(rewarded_processed_position_data["trial_number"]):
        trial_rewarded_processed_position_data = rewarded_processed_position_data[rewarded_processed_position_data["trial_number"] == trial_number]

        rewarded_speeds_in_space = Edmond.plot_utility2.pandas_collumn_to_numpy_array(trial_rewarded_processed_position_data['speeds_binned_in_space'])
        rewarded_bin_centres = Edmond.plot_utility2.pandas_collumn_to_numpy_array(trial_rewarded_processed_position_data['position_bin_centres'])
        in_rz_mask = (rewarded_bin_centres > reward_zone_start) & (rewarded_bin_centres <= reward_zone_end)
        rewarded_speeds_in_space_in_reward_zone = rewarded_speeds_in_space[in_rz_mask]
        rewarded_speeds_in_space_in_reward_zone = rewarded_speeds_in_space_in_reward_zone[~np.isnan(rewarded_speeds_in_space_in_reward_zone)]
        minimum_speed = min(rewarded_speeds_in_space_in_reward_zone)
        minimum_speeds_in_rz.append(minimum_speed)
        speeds_in_rz.extend(rewarded_speeds_in_space_in_reward_zone.tolist())

    speeds_in_rz = np.array(speeds_in_rz)
    minimum_speeds_in_rz = np.array(minimum_speeds_in_rz)
    mean, sigma = np.nanmean(speeds_in_rz), np.nanstd(speeds_in_rz)
    #mean, sigma = np.nanmean(minimum_speeds_in_rz), np.nanstd(minimum_speeds_in_rz)
    interval = stats.norm.interval(0.95, loc=mean, scale=sigma)
    upper = interval[1]
    lower = interval[0]

    hit_miss_try =[]
    avg_speed_in_rz =[]
    for i, trial_number in enumerate(processed_position_data.trial_number):
        trial_process_position_data = processed_position_data[(processed_position_data.trial_number == trial_number)]
        trial_speeds_in_space = Edmond.plot_utility2.pandas_collumn_to_numpy_array(trial_process_position_data['speeds_binned_in_space'])
        trial_bin_centres = Edmond.plot_utility2.pandas_collumn_to_numpy_array(trial_process_position_data['position_bin_centres'])
        in_rz_mask = (trial_bin_centres > reward_zone_start) & (trial_bin_centres <= reward_zone_end)
        trial_speeds_in_reward_zone = trial_speeds_in_space[in_rz_mask]
        trial_speeds_in_reward_zone = trial_speeds_in_reward_zone[~np.isnan(trial_speeds_in_reward_zone)]
        avg_trial_speed_in_reward_zone = np.mean(trial_speeds_in_reward_zone)

        if trial_process_position_data["rewarded"].iloc[0] == True:
            hit_miss_try.append("hit")
        elif (avg_trial_speed_in_reward_zone >= lower) and (avg_trial_speed_in_reward_zone <= upper):
            hit_miss_try.append("try")
        elif (avg_trial_speed_in_reward_zone < lower) or (avg_trial_speed_in_reward_zone > upper):
            hit_miss_try.append("miss")
        else:
            hit_miss_try.append("miss")

        avg_speed_in_rz.append(avg_trial_speed_in_reward_zone)

    processed_position_data["hit_miss_try"] = hit_miss_try
    processed_position_data["avg_speed_in_rz"] = avg_speed_in_rz
    return processed_position_data, upper

def find_paired_recording(recording_path, of_recording_path_list):
    mouse=recording_path.split("/")[-1].split("_")[0]
    training_day=recording_path.split("/")[-1].split("_")[1]

    for paired_recording in of_recording_path_list:
        paired_mouse=paired_recording.split("/")[-1].split("_")[0]
        paired_training_day=paired_recording.split("/")[-1].split("_")[1]

        if (mouse == paired_mouse) and (training_day == paired_training_day):
            return paired_recording, True
    return None, False

def find_set(a,b):
    return set(a) & set(b)

def moving_sum(array, window):
    ret = np.cumsum(array, dtype=float)
    ret[window:] = ret[window:] - ret[:-window]
    return ret[window:]

def downsample(array, npts):
    interpolated = interp1d(np.arange(len(array)), array, axis = 0, fill_value = 'extrapolate')
    downsampled = interpolated(np.linspace(0, len(array), npts))
    return downsampled

def reduce_digits(numeric_float, n_digits=6):
    #if len(list(str(numeric_float))) > n_digits:
    scientific_notation = "{:.1e}".format(numeric_float)
    return scientific_notation

def get_tt_color(tt):
    if tt == 0:
        return "black"
    elif tt==1:
        return "red"
    elif tt ==2:
        return "blue"

def get_hmt_linestyle(hmt):
    if hmt == "hit":
        return "solid"
    elif hmt=="miss":
        return "dashed"
    elif hmt =="try":
        return "dotted"

def distance_from_integer(frequencies):
    distance_from_zero = np.asarray(frequencies)%1
    distance_from_one = 1-(np.asarray(frequencies)%1)
    tmp = np.vstack((distance_from_zero, distance_from_one))
    return np.min(tmp, axis=0)

def min_max_normalize(x):
    """
        argument
            - x: input image data in numpy array [32, 32, 3]
        return
            - normalized x
    """
    min_val = np.min(x)
    max_val = np.max(x)
    x = (x-min_val) / (max_val-min_val)
    return x

def min_max_normlise(array, min_val, max_val):
    normalised_array = ((max_val-min_val)*((array-min(array))/(max(array)-min(array))))+min_val
    return normalised_array

def get_track_length(recording_path):
    parameter_file_path = control_sorting_analysis.get_tags_parameter_file(recording_path)
    stop_threshold, track_length, cue_conditioned_goal = PostSorting.post_process_sorted_data_vr.process_running_parameter_tag(parameter_file_path)
    return track_length

def get_trial_color(trial_type):
    if trial_type == 0:
        return "black"
    elif trial_type == 1:
        return "red"
    elif trial_type == 2:
        return "blue"
    else:
        print("invalid trial-type passed to get_trial_color()")

def add_time_elapsed_collumn(position_data):
    time_seconds = np.array(position_data['time_seconds'].to_numpy())
    time_elapsed = np.diff(time_seconds)
    time_elapsed = np.append(time_elapsed[0], time_elapsed)
    position_data["time_in_bin_seconds"] = time_elapsed.tolist()
    return position_data

def get_max_SNR(spatial_frequency, powers):
    max_SNR = np.abs(np.max(powers)/np.min(powers))
    max_SNR = powers[np.argmax(powers)]
    max_SNR_freq = spatial_frequency[np.argmax(powers)]
    return max_SNR, max_SNR_freq

def plot_joint_cell_cross_correlations(spike_data, output_path):
    spike_data = add_lomb_classifier(spike_data)
    print('plotting joint cell correlations...')
    save_path = output_path + '/Figures/joint_correlations'
    if os.path.exists(save_path) is False:
        os.makedirs(save_path)

    cluster_ids = pandas_collumn_to_numpy_array(spike_data["cluster_id"])
    peak_powers = pandas_collumn_to_2d_numpy_array(spike_data["peak_powers"])
    lomb_classes = pandas_collumn_to_numpy_array(spike_data["Lomb_classifier_"])

    cluster_ids = cluster_ids[np.argsort(lomb_classes)]
    peak_powers = peak_powers[np.argsort(lomb_classes)]
    lomb_classes = lomb_classes[np.argsort(lomb_classes)]

    cross_correlation_matrix = np.zeros((len(cluster_ids), len(cluster_ids)))
    for i, cluster_id_i in enumerate(cluster_ids):
        for j, cluster_id_j in enumerate(cluster_ids):
            if (len(peak_powers[i])>1) and (len(peak_powers[j])>1):
                corr = pearsonr(peak_powers[i], peak_powers[j])[0]
                cross_correlation_matrix[i, j] = corr
            else:
                corr = np.nan
                cross_correlation_matrix[i, j] = corr

    fig, ax = plt.subplots()
    im= ax.imshow(cross_correlation_matrix, vmin=0, vmax=1)
    ax.set_xticks(np.arange(len(cluster_ids)))
    ax.set_yticks(np.arange(len(cluster_ids)))
    ax.set_yticklabels(cluster_ids)
    ax.set_xticklabels(cluster_ids)

    colors = []
    for i in range(len(lomb_classes)):
        colors.append(get_lomb_class_color(lomb_classes[i]))

    for xtick, color in zip(ax.get_xticklabels(), colors):
        xtick.set_color(color)
    for ytick, color in zip(ax.get_yticklabels(), colors):
        ytick.set_color(color)

    ax.set_ylabel("Cluster ID", fontsize=5)
    ax.set_xlabel("Cluster IDs", fontsize=5)
    ax.tick_params(axis='both', which='major', labelsize=15)
    fig.tight_layout()
    fig.colorbar(im, ax=ax)
    plt.savefig(save_path + '/' + spike_data.session_id.iloc[0] + '_joint_peak_powers_cross_correlations.png', dpi=300)
    plt.close()
    return

def plot_joint_jitter_correlations(spike_data, of_spike_data, processed_position_data, position_data, output_path, track_length, matched_recording_df):

    spike_data = add_lomb_classifier(spike_data)
    print('plotting joint cell correlations...')
    save_path = output_path + '/Figures/joint_correlations'
    if os.path.exists(save_path) is False:
        os.makedirs(save_path)

    spike_data = pd.merge(spike_data, of_spike_data[["cluster_id", "grid_cell"]], on="cluster_id")
    spike_data = spike_data.sort_values(by=["grid_cell"], ascending=False)
    cluster_ids = pandas_collumn_to_numpy_array(spike_data["cluster_id"])
    ids = cluster_ids
    n_shuffles = 1
    shuffled = np.ones((n_shuffles+1), dtype=bool); shuffled[0] = False

    for m, hmt in enumerate(["hit", "try", "miss"]):
        for n, tt in enumerate([0,1]):
            hmt_processed_position_data = processed_position_data[processed_position_data["hit_miss_try"] == hmt]
            subset_processed_position_data = hmt_processed_position_data[hmt_processed_position_data["trial_type"] == tt]
            subset_trial_numbers = pandas_collumn_to_numpy_array(subset_processed_position_data["trial_number"])
            if len(subset_trial_numbers)>0:

                cross_correlations = np.zeros((len(shuffled), len(ids), len(ids)))
                for si, shuffle in enumerate(shuffled):
                    for i in range(len(cross_correlations[0])):
                        for j in range(len(cross_correlations[0][0])):
                            cluster_j_df = spike_data[spike_data["cluster_id"] == ids[j]]
                            cluster_i_df = spike_data[spike_data["cluster_id"] == ids[i]]
                            shifts_i = get_shifts(cluster_i_df, hmt=hmt, tt=tt)
                            if shuffle:
                                #np.random.seed(j*i)
                                np.random.shuffle(shifts_i)
                                #shifts_i = np.random.randint(-100, 100, size=len(shifts_i))

                            cluster_firing_maps_j = np.array(cluster_j_df["fr_binned_in_space"].iloc[0])
                            where_are_NaNs2 = np.isnan(cluster_firing_maps_j)
                            cluster_firing_maps_j[where_are_NaNs2] = 0
                            cluster_firing_maps_j = min_max_normalize(cluster_firing_maps_j)
                            hmt_cluster_firing_maps_j = cluster_firing_maps_j[subset_trial_numbers-1]
                            avg_correlation = get_avg_correlation(hmt_cluster_firing_maps_j)

                            # reconstruct avg spatial correlation using the newly aligned trials
                            reconstructed_signal = []
                            for ti, tn in enumerate(subset_trial_numbers):
                                reconstructed_trial = get_reconstructed_trial_signal(shifts_i[ti], hmt_cluster_firing_maps_j[ti].flatten(),
                                                                                min_shift=min(shifts_i), max_shift=max(shifts_i))
                                reconstructed_signal.append(reconstructed_trial.tolist())
                            reconstructed_signal = np.array(reconstructed_signal)
                            reconstructed_signal_corr = get_avg_correlation(reconstructed_signal)

                            cross_correlations[si, i, j] = reconstructed_signal_corr - avg_correlation

                fig, ax = plt.subplots()
                cross_correlation = cross_correlations[0]-np.nanmean(cross_correlations[1:], axis=0)
                np.fill_diagonal(cross_correlation, np.nan)

                cross_correlation_rank = np.zeros((len(cross_correlation), len(cross_correlation[0])))
                cross_correlation_values = cross_correlation.flatten()
                cross_correlation_values = cross_correlation_values[~np.isnan(cross_correlation_values)]
                for ii in range(len(cross_correlation)):
                    for jj in range(len(cross_correlation[0])):
                        cross_correlation_rank[ii, jj] = stats.percentileofscore(cross_correlation_values, cross_correlation[ii,jj])

                im= ax.imshow(cross_correlation_rank, cmap="Purples", vmin=50, vmax=100)
                ax.set_xticks(np.arange(len(ids)))
                ax.set_yticks(np.arange(len(ids)))
                ax.set_yticklabels(ids)
                ax.set_xticklabels(ids)
                ax.set_ylabel("Cluster ID: Reference", fontsize=20)
                ax.set_xlabel("Cluster ID: Shifted", fontsize=20)
                ax.tick_params(axis='both', which='major', labelsize=8)
                fig.tight_layout()
                fig.colorbar(im, ax=ax)
                plt.savefig(save_path + '/' + spike_data.session_id.iloc[0] + 'joint_alignment_cross_correlations_'+hmt+'_'+str(tt)+'.png', dpi=300)
                plt.close()
                print(hmt+" "+ str(np.nanmean(cross_correlations[0])))

                fig, ax = plt.subplots(figsize=(6,6))
                ax.axhline(y=0, linestyle="dashed", color="gray")
                ax.axvline(x=0, linestyle="dashed", color="gray")
                ax.plot(np.arange(-2, 2), np.arange(-2, 2), linestyle="solid", color="black")

                grid_pairs_vs_shuffle = []
                grid_non_grid_pairs_vs_shuffle = []
                non_grid_non_grid_pairs_vs_shuffle = []
                tetrode_level = []
                for i in range(len(cross_correlations[0])):
                    for j in range(len(cross_correlations[0][0])):
                        cluster_j_classifier = spike_data[spike_data["cluster_id"] == ids[j]]["grid_cell"].iloc[0]
                        cluster_i_classifier = spike_data[spike_data["cluster_id"] == ids[i]]["grid_cell"].iloc[0]
                        tetrode_j = spike_data[spike_data["cluster_id"] == ids[j]]["tetrode"].iloc[0]
                        tetrode_i = spike_data[spike_data["cluster_id"] == ids[i]]["tetrode"].iloc[0]

                        if tetrode_i == tetrode_j:
                            marker="x"
                            tetrode_level.append("same")
                        else:
                            marker="o"
                            tetrode_level.append("different")

                        if ((cluster_j_classifier == True) and (cluster_i_classifier == True)):
                            ax.scatter(cross_correlations[0, i, j], np.nanmean(cross_correlations[1:], axis=0)[i, j], marker=marker, color="red", zorder=10)
                            grid_pairs_vs_shuffle.append(cross_correlations[0, i, j]-np.nanmean(cross_correlations[1:]))
                        elif ((cluster_j_classifier == True) or (cluster_i_classifier == True)):
                            ax.scatter(cross_correlations[0, i, j], np.nanmean(cross_correlations[1:], axis=0)[i, j], marker=marker, color="blue", alpha=0.3, zorder=5)
                            grid_non_grid_pairs_vs_shuffle.append(cross_correlations[0, i, j]-np.nanmean(cross_correlations[1:]))
                        else:
                            ax.scatter(cross_correlations[0, i, j], np.nanmean(cross_correlations[1:], axis=0)[i, j], marker=marker, color="black", alpha=0.3, zorder=1)
                            non_grid_non_grid_pairs_vs_shuffle.append(cross_correlations[0, i, j]-np.nanmean(cross_correlations[1:]))

                ax.set_ylabel("Change in R (Shuffle)", fontsize=20)
                ax.set_xlabel("Change in R (Real)", fontsize=20)
                ax.spines['top'].set_visible(False)
                ax.spines['right'].set_visible(False)
                ax.set_ylim([-0.5, 0.5])
                ax.set_xlim([-0.5, 0.5])
                ax.tick_params(axis='both', which='major', labelsize=15)
                fig.tight_layout()
                plt.savefig(save_path + '/' + spike_data.session_id.iloc[0] + 'joint_alignment_correlations_shuffle_vs_real_'+hmt+'_'+str(tt)+'png', dpi=300)
                plt.close()

                tetrode_level = np.array(tetrode_level)
                grid_pairs_vs_shuffle = np.array(grid_pairs_vs_shuffle); gg_p = stats.wilcoxon(grid_pairs_vs_shuffle)[1]; gg_p= get_p_text(gg_p, ns=False)
                grid_non_grid_pairs_vs_shuffle = np.array(grid_non_grid_pairs_vs_shuffle); gng_p = stats.wilcoxon(grid_non_grid_pairs_vs_shuffle)[1]; gng_p= get_p_text(gng_p, ns=False)
                non_grid_non_grid_pairs_vs_shuffle = np.array(non_grid_non_grid_pairs_vs_shuffle); ngng_p = stats.wilcoxon(non_grid_non_grid_pairs_vs_shuffle)[1]; ngng_p= get_p_text(ngng_p, ns=False)
                data = [grid_pairs_vs_shuffle, grid_non_grid_pairs_vs_shuffle, non_grid_non_grid_pairs_vs_shuffle];

                fig, ax = plt.subplots(figsize=(6,6))
                vp = ax.violinplot(data, [2, 4, 6], widths=2,
                                   showmeans=False, showmedians=True, showextrema=False)
                height =ax.get_ylim()[1]
                ax.text(2, height, gg_p, ha="center", fontsize=15)
                ax.text(4, height, gng_p, ha="center", fontsize=15)
                ax.text(6, height, ngng_p, ha="center", fontsize=15)
                ax.axhline(y=0, linestyle="dashed", color="black")
                ax.set_xticks([2,4,6])
                ax.set_yticks([-0.2, 0, 0.2, 0.4])
                ax.set_ylim([-0.4, 0.6])
                ax.set_xticklabels(["G-G", "G-NG", "NG-NG"])
                ax.set_ylabel("Change in spatial\ncorrelation vs shuffle", fontsize=20, labelpad=10)
                ax.spines['top'].set_visible(False)
                ax.spines['right'].set_visible(False)
                ax.spines['bottom'].set_visible(False)
                ax.xaxis.set_tick_params(length=0)
                ax.tick_params(axis='both', which='major', labelsize=25)
                fig.tight_layout()
                plt.savefig(save_path + '/' + spike_data.session_id.iloc[0] + 'joint_alignment_correlations_violin_shuffle_vs_real_'+hmt+'_'+str(tt)+'png', dpi=300)
                plt.close()

                session_df = pd.DataFrame()
                session_df["session_id"] = [spike_data.session_id.iloc[0]]
                session_df["trial_type"] = [tt]
                session_df["hit_miss_try"] = [hmt]
                session_df["tetrode_level"] = [tetrode_level]
                session_df["grid_pairs_vs_shuffle"] = [grid_pairs_vs_shuffle]
                session_df["grid_non_grid_pairs_vs_shuffle"] = [grid_non_grid_pairs_vs_shuffle]
                session_df["non_grid_non_grid_pairs_vs_shuffle"] = [non_grid_non_grid_pairs_vs_shuffle]
                # save the session_df and append to the global matched recording df
                matched_recording_df = pd.concat([matched_recording_df, session_df], ignore_index=True)

    return matched_recording_df

def plot_joint_cell_correlations(spike_data, of_spike_data, processed_position_data, position_data, raw_position_data, output_path, track_length):
    spike_data = add_lomb_classifier(spike_data)
    print('plotting joint cell correlations...')
    save_path = output_path + '/Figures/joint_correlations'
    if os.path.exists(save_path) is False:
        os.makedirs(save_path)

    gauss_kernel_peak_powers = Gaussian1DKernel(stddev=30)

    # get trial numbers to use from processed_position_data
    trial_number_to_use = np.unique(processed_position_data["trial_number"])
    trial_numbers = np.array(position_data["trial_number"])
    x_positions = np.array(position_data["x_position_cm"])
    x_positions_elapsed = x_positions+(trial_numbers*track_length)-track_length
    n_trials = max(trial_numbers)

    # get distances from raw position data which has been smoothened
    rpd = np.asarray(raw_position_data["x_position_cm"])
    tn = np.asarray(raw_position_data["trial_number"])
    elapsed_distance30 = rpd+(tn*track_length)-track_length

    # get denominator and handle nans
    denominator, _ = np.histogram(elapsed_distance30, bins=int(track_length/1)*n_trials, range=(0, track_length*n_trials))

    fig = plt.figure(figsize=(4,4))
    ax = fig.add_subplot(1, 1, 1)  # specify (nrows, ncols, axnum)

    peak_powers = []
    for cluster_index, cluster_id in enumerate(spike_data.cluster_id):
        cluster_spike_data = spike_data[spike_data["cluster_id"] == cluster_id]
        cluster_spike_data_of = of_spike_data[of_spike_data["cluster_id"] == cluster_id]

        firing_times_cluster = np.array(cluster_spike_data["firing_times"].iloc[0])
        firing_locations_cluster = np.array(cluster_spike_data["x_position_cm"].iloc[0])
        firing_trial_numbers = np.array(cluster_spike_data["trial_number"].iloc[0])
        lomb_class = cluster_spike_data["Lomb_classifier_"].iloc[0]
        #cell_type = cluster_spike_data_of["classifier"].iloc[0]
        firing_locations_cluster_elapsed = firing_locations_cluster+(firing_trial_numbers*track_length)-track_length

        if len(firing_times_cluster)>1:
            numerator, bin_edges = np.histogram(firing_locations_cluster_elapsed, bins=int(track_length/1)*n_trials, range=(0, track_length*n_trials))
            fr = numerator/denominator
            elapsed_distance = 0.5*(bin_edges[1:]+bin_edges[:-1])/track_length
            trial_numbers_by_bin=((0.5*(bin_edges[1:]+bin_edges[:-1])//track_length)+1).astype(np.int32)
            gauss_kernel = Gaussian1DKernel(stddev=1)

            # remove nan values that coincide with start and end of the track before convolution
            fr[fr==np.inf] = np.nan
            nan_mask = ~np.isnan(fr)
            fr = fr[nan_mask]
            trial_numbers_by_bin = trial_numbers_by_bin[nan_mask]
            elapsed_distance = elapsed_distance[nan_mask]

            fr = convolve(fr, gauss_kernel)
            fr = moving_sum(fr, window=2)/2
            fr = np.append(fr, np.zeros(len(elapsed_distance)-len(fr)))

            # make and apply the set mask
            set_mask = np.isin(trial_numbers_by_bin, trial_number_to_use)
            fr = fr[set_mask]
            elapsed_distance = elapsed_distance[set_mask]

            # construct the lomb-scargle periodogram
            step = Settings.frequency_step
            frequency = Settings.frequency
            sliding_window_size=track_length*Settings.window_length_in_laps

            powers = []
            centre_distances = []
            indices_to_test = np.arange(0, len(fr)-sliding_window_size, 1, dtype=np.int64)[::10]
            for m in indices_to_test:
                ls = LombScargle(elapsed_distance[m:m+sliding_window_size], fr[m:m+sliding_window_size])
                power = ls.power(frequency)
                power[np.isnan(power)] = 0
                powers.append(power.tolist())
                centre_distances.append(np.nanmean(elapsed_distance[m:m+sliding_window_size]))
            powers = np.array(powers)
            centre_trials = np.round(np.array(centre_distances)).astype(np.int64)

            avg_power = np.nanmean(powers, axis=0)
            max_SNR, max_SNR_freq = get_max_SNR(frequency, avg_power)
            max_SNR_text = "SNR: " + reduce_digits(np.round(max_SNR, decimals=2), n_digits=6)
            max_SNR_freq_test = "Freq: " + str(np.round(max_SNR_freq, decimals=1))

            peak_powers_cluster = powers[:, np.argmin(np.abs(frequency-max_SNR_freq))]
            peak_powers_cluster = convolve(peak_powers_cluster, gauss_kernel_peak_powers)
            peak_powers_cluster = moving_sum(peak_powers_cluster, window=20)/20
            peak_powers_cluster = np.append(peak_powers_cluster, np.zeros(len(centre_trials)-len(peak_powers_cluster)))

            peak_powers_cluster = min_max_normalize(peak_powers_cluster)
            peak_powers.append(peak_powers_cluster.tolist())

            ax.plot(centre_trials, peak_powers_cluster, color=get_lomb_class_color(lomb_class), alpha=0.3, linewidth=1)
        else:
            peak_powers.append(np.nan)

    spike_data["peak_powers"] = peak_powers
    plt.xlabel('Centre Trial', fontsize=20, labelpad = 10)
    plt.ylabel('Peak Power', fontsize=20, labelpad = 10)
    plt.xlim(0,max(centre_trials))
    plt.ylim(bottom=0)
    plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=1.0)
    plt.savefig(save_path + '/' + spike_data.session_id.iloc[cluster_index] + '_joint_peak_powers.png', dpi=300)
    plt.close()
    return spike_data

def get_lomb_class_color(lomb_class):
    if lomb_class == "Position":
        return "turquoise"
    elif lomb_class == "Distance":
        return "orange"
    elif lomb_class == "Null":
        return "gray"
    else:
        return ""

def get_cell_type_color(cell_type):
    if cell_type == "G":
        return "solid"
    else:
        return "dashed"

def add_percentage_for_lomb_classes(combined_df):
    # this function calculates the percentage of concurrently recorded cells

    NG_props_D = []
    NG_props_P = []
    NG_props_N = []
    G_props_D = []
    G_props_P = []
    G_props_N = []
    for index, cluster_row in combined_df.iterrows():
        cluster_row = cluster_row.to_frame().T.reset_index(drop=True)
        session_id = cluster_row["session_id"].iloc[0]
        cluster_id = cluster_row["cluster_id"].iloc[0]
        NG_cells_in_same_recordings = combined_df[(combined_df["session_id"] == session_id) & (combined_df["classifier"] != "G") & (combined_df["cluster_id"] != cluster_id)]
        G_cells_in_same_recordings = combined_df[(combined_df["session_id"] == session_id) & (combined_df["classifier"] == "G") & (combined_df["cluster_id"] != cluster_id)]
        NG_lomb_classes = NG_cells_in_same_recordings["Lomb_classifier_"]
        G_lomb_classes = G_cells_in_same_recordings["Lomb_classifier_"]

        if len(NG_lomb_classes)>0:
            NG_proportion_D = np.sum(NG_lomb_classes=="Distance")/len(NG_lomb_classes)
            NG_proportion_P = np.sum(NG_lomb_classes=="Position")/len(NG_lomb_classes)
            NG_proportion_N = np.sum(NG_lomb_classes=="Null")/len(NG_lomb_classes)
        else:
            NG_proportion_D = np.nan
            NG_proportion_P = np.nan
            NG_proportion_N = np.nan

        if len(G_lomb_classes)>0:
            G_proportion_D = np.sum(G_lomb_classes=="Distance")/len(G_lomb_classes)
            G_proportion_P = np.sum(G_lomb_classes=="Position")/len(G_lomb_classes)
            G_proportion_N = np.sum(G_lomb_classes=="Null")/len(G_lomb_classes)
        else:
            G_proportion_D = np.nan
            G_proportion_P = np.nan
            G_proportion_N = np.nan

        NG_props_D.append(NG_proportion_D)
        NG_props_P.append(NG_proportion_P)
        NG_props_N.append(NG_proportion_N)
        G_props_D.append(G_proportion_D)
        G_props_P.append(G_proportion_P)
        G_props_N.append(G_proportion_N)

    combined_df["NG_props_D"] = NG_props_D
    combined_df["NG_props_P"] = NG_props_P
    combined_df["NG_props_N"] = NG_props_N
    combined_df["G_props_D"] = G_props_D
    combined_df["G_props_P"] = G_props_P
    combined_df["G_props_N"] = G_props_N
    return combined_df

def get_grid_cells_from_same_recording(spatial_firing):

    grid_cells = pd.DataFrame()
    for index, cluster_row in spatial_firing.iterrows():
        cluster_row = cluster_row.to_frame().T.reset_index(drop=True)
        session_id = cluster_row["session_id"].iloc[0]

        same_session_id_cells = spatial_firing[spatial_firing["session_id"] == session_id]
        if len(same_session_id_cells)>1:
            grid_cells = pd.concat([grid_cells, cluster_row], ignore_index=True)
    return grid_cells

def plot_class_prection_credence(spatial_firing, save_path):
    distance_cells = spatial_firing[spatial_firing["Lomb_classifier_"] == "Distance"]
    position_cells = spatial_firing[spatial_firing["Lomb_classifier_"] == "Position"]
    null_cells = spatial_firing[spatial_firing["Lomb_classifier_"] == "Null"]
    errorbarwidth =2

    fig, ax = plt.subplots(figsize=(6,6))
    ax.bar(x=0.4, height=np.nanmean(np.asarray(position_cells["G_props_P"]), dtype=np.float64), width=0.1, edgecolor="black", color="turquoise")
    ax.errorbar(x=0.4, y=np.nanmean(np.asarray(position_cells["G_props_P"]), dtype=np.float64), yerr=stats.sem(np.asarray(position_cells["G_props_P"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)
    ax.bar(x=0.5, height=np.nanmean(np.asarray(position_cells["G_props_D"]), dtype=np.float64), width=0.1, edgecolor="black", color="orange")
    ax.errorbar(x=0.5, y=np.nanmean(np.asarray(position_cells["G_props_D"]), dtype=np.float64), yerr=stats.sem(np.asarray(position_cells["G_props_D"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)
    ax.bar(x=0.6, height=np.nanmean(np.asarray(position_cells["G_props_N"]), dtype=np.float64), width=0.1, edgecolor="black", color="gray")
    ax.errorbar(x=0.6, y=np.nanmean(np.asarray(position_cells["G_props_N"]), dtype=np.float64), yerr=stats.sem(np.asarray(position_cells["G_props_N"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)

    ax.bar(x=0.9, height=np.nanmean(np.asarray(distance_cells["G_props_P"]), dtype=np.float64), width=0.1, edgecolor="black", color="turquoise")
    ax.errorbar(x=0.9, y=np.nanmean(np.asarray(distance_cells["G_props_P"]), dtype=np.float64), yerr=stats.sem(np.asarray(distance_cells["G_props_P"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)
    ax.bar(x=1.0, height=np.nanmean(np.asarray(distance_cells["G_props_D"]), dtype=np.float64), width=0.1, edgecolor="black", color="orange")
    ax.errorbar(x=1.0, y=np.nanmean(np.asarray(distance_cells["G_props_D"]), dtype=np.float64), yerr=stats.sem(np.asarray(distance_cells["G_props_D"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)
    ax.bar(x=1.1, height=np.nanmean(np.asarray(distance_cells["G_props_N"]), dtype=np.float64), width=0.1, edgecolor="black", color="gray")
    ax.errorbar(x=1.1, y=np.nanmean(np.asarray(distance_cells["G_props_N"]), dtype=np.float64), yerr=stats.sem(np.asarray(distance_cells["G_props_N"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)

    ax.bar(x=1.4, height=np.nanmean(np.asarray(null_cells["G_props_P"]), dtype=np.float64), width=0.1, edgecolor="black", color="turquoise")
    ax.errorbar(x=1.4, y=np.nanmean(np.asarray(null_cells["G_props_P"]), dtype=np.float64), yerr=stats.sem(np.asarray(null_cells["G_props_P"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)
    ax.bar(x=1.5, height=np.nanmean(np.asarray(null_cells["G_props_D"]), dtype=np.float64), width=0.1, edgecolor="black", color="orange")
    ax.errorbar(x=1.5, y=np.nanmean(np.asarray(null_cells["G_props_D"]), dtype=np.float64), yerr=stats.sem(np.asarray(null_cells["G_props_D"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)
    ax.bar(x=1.6, height=np.nanmean(np.asarray(null_cells["G_props_N"]), dtype=np.float64), width=0.1, edgecolor="black", color="gray")
    ax.errorbar(x=1.6, y=np.nanmean(np.asarray(null_cells["G_props_N"]), dtype=np.float64), yerr=stats.sem(np.asarray(null_cells["G_props_N"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)

    ax.axhline(y=0.33, color="black", linestyle="dashed", linewidth=2)
    plt.ylabel('Probability', fontsize=20, labelpad = 10)
    plt.xlabel("Cell Class", fontsize=20, labelpad = 10)
    ax.set_xticks([0.5, 1, 1.5])
    ax.set_xticklabels(["P", "D", "N"])
    ax.set_yticks([0, 0.5,  1])
    ax.set_ylim(bottom=0, top=1)
    ax.set_xlim(left=0, right=2)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    plt.subplots_adjust(hspace = .35, wspace = .35,  bottom = 0.2, left = 0.32, right = 0.87, top = 0.92)
    plt.savefig(save_path + '/prediction_credence_G.png', dpi=200)
    plt.close()

    fig, ax = plt.subplots(figsize=(6,6))
    ax.bar(x=0.4, height=np.nanmean(np.asarray(position_cells["NG_props_P"]), dtype=np.float64), width=0.1, edgecolor="black", color="turquoise", hatch="/")
    ax.errorbar(x=0.4, y=np.nanmean(np.asarray(position_cells["NG_props_P"]), dtype=np.float64), yerr=stats.sem(np.asarray(position_cells["NG_props_P"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)
    ax.bar(x=0.5, height=np.nanmean(np.asarray(position_cells["NG_props_D"]), dtype=np.float64), width=0.1, edgecolor="black", color="orange", hatch="/")
    ax.errorbar(x=0.5, y=np.nanmean(np.asarray(position_cells["NG_props_D"]), dtype=np.float64), yerr=stats.sem(np.asarray(position_cells["NG_props_D"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)
    ax.bar(x=0.6, height=np.nanmean(np.asarray(position_cells["NG_props_N"]), dtype=np.float64), width=0.1, edgecolor="black", color="gray", hatch="/")
    ax.errorbar(x=0.6, y=np.nanmean(np.asarray(position_cells["NG_props_N"]), dtype=np.float64), yerr=stats.sem(np.asarray(position_cells["NG_props_N"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)

    ax.bar(x=0.9, height=np.nanmean(np.asarray(distance_cells["NG_props_P"]), dtype=np.float64), width=0.1, edgecolor="black", color="turquoise", hatch="/")
    ax.errorbar(x=0.9, y=np.nanmean(np.asarray(distance_cells["NG_props_P"]), dtype=np.float64), yerr=stats.sem(np.asarray(distance_cells["NG_props_P"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)
    ax.bar(x=1.0, height=np.nanmean(np.asarray(distance_cells["NG_props_D"]), dtype=np.float64), width=0.1, edgecolor="black", color="orange", hatch="/")
    ax.errorbar(x=1.0, y=np.nanmean(np.asarray(distance_cells["NG_props_D"]), dtype=np.float64), yerr=stats.sem(np.asarray(distance_cells["NG_props_D"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)
    ax.bar(x=1.1, height=np.nanmean(np.asarray(distance_cells["NG_props_N"]), dtype=np.float64), width=0.1, edgecolor="black", color="gray", hatch="/")
    ax.errorbar(x=1.1, y=np.nanmean(np.asarray(distance_cells["NG_props_N"]), dtype=np.float64), yerr=stats.sem(np.asarray(distance_cells["NG_props_N"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)

    ax.bar(x=1.4, height=np.nanmean(np.asarray(null_cells["NG_props_P"]), dtype=np.float64), width=0.1, edgecolor="black", color="turquoise", hatch="/")
    ax.errorbar(x=1.4, y=np.nanmean(np.asarray(null_cells["NG_props_P"]), dtype=np.float64), yerr=stats.sem(np.asarray(null_cells["NG_props_P"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)
    ax.bar(x=1.5, height=np.nanmean(np.asarray(null_cells["NG_props_D"]), dtype=np.float64), width=0.1, edgecolor="black", color="orange", hatch="/")
    ax.errorbar(x=1.5, y=np.nanmean(np.asarray(null_cells["NG_props_D"]), dtype=np.float64), yerr=stats.sem(np.asarray(null_cells["NG_props_D"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)
    ax.bar(x=1.6, height=np.nanmean(np.asarray(null_cells["NG_props_N"]), dtype=np.float64), width=0.1, edgecolor="black", color="gray", hatch="/")
    ax.errorbar(x=1.6, y=np.nanmean(np.asarray(null_cells["NG_props_N"]), dtype=np.float64), yerr=stats.sem(np.asarray(null_cells["NG_props_N"], dtype=np.float64), nan_policy="omit"), color="black",elinewidth=errorbarwidth)

    ax.axhline(y=0.33, color="black", linestyle="dashed", linewidth=2)
    plt.ylabel('Probability', fontsize=20, labelpad = 10)
    plt.xlabel("Cell Class", fontsize=20, labelpad = 10)
    ax.set_xticks([0.5, 1, 1.5])
    ax.set_xticklabels(["P", "D", "N"])
    ax.set_yticks([0, 0.5,  1])
    ax.set_ylim(bottom=0, top=1)
    ax.set_xlim(left=0, right=2)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    plt.subplots_adjust(hspace = .35, wspace = .35,  bottom = 0.2, left = 0.32, right = 0.87, top = 0.92)
    plt.savefig(save_path + '/prediction_credence_NG.png', dpi=200)
    plt.close()

    return

def significance_bar(ax, start,end,height,displaystring,linewidth = 1.2,markersize = 8,boxpad  =0.3,fontsize = 15,color = 'k'):
    # draw a line with downticks at the ends
    ax.plot([start,end],[height]*2,'-',color = color,lw=linewidth,marker = TICKDOWN,markeredgewidth=linewidth,markersize = markersize)
    # draw the text with a bounding box covering up the line
    ax.text(0.5*(start+end),height,displaystring,ha = 'center',va='center',bbox=dict(facecolor='1.', edgecolor='none',boxstyle='Square,pad='+str(boxpad)),size = fontsize)

def plot_all_paired_vs_shuffle(df, output_path):
    for hmt in ["hit", "miss", "try"]:
        for tt in [0, 1]:
            subset_df = df[((df["hit_miss_try"]==hmt) & (df["trial_type"]==tt))]
            grid_pairs_vs_shuffle = np.array([])
            grid_non_grid_pairs_vs_shuffle = np.array([])
            non_grid_non_grid_pairs_vs_shuffle = np.array([])

            for index, session_row in subset_df.iterrows():
                session_row = session_row.to_frame().T.reset_index(drop=True)
                grid_pairs_vs_shuffle = np.concatenate((grid_pairs_vs_shuffle, session_row.iloc[0]["grid_pairs_vs_shuffle"]))
                grid_non_grid_pairs_vs_shuffle = np.concatenate((grid_non_grid_pairs_vs_shuffle, session_row.iloc[0]["grid_non_grid_pairs_vs_shuffle"]))
                non_grid_non_grid_pairs_vs_shuffle = np.concatenate((non_grid_non_grid_pairs_vs_shuffle, session_row.iloc[0]["non_grid_non_grid_pairs_vs_shuffle"]))

            grid_pairs_vs_shuffle = grid_pairs_vs_shuffle[~np.isnan(grid_pairs_vs_shuffle)]
            grid_non_grid_pairs_vs_shuffle = grid_non_grid_pairs_vs_shuffle[~np.isnan(grid_non_grid_pairs_vs_shuffle)]
            non_grid_non_grid_pairs_vs_shuffle = non_grid_non_grid_pairs_vs_shuffle[~np.isnan(non_grid_non_grid_pairs_vs_shuffle)]

            data = [grid_pairs_vs_shuffle, grid_non_grid_pairs_vs_shuffle, non_grid_non_grid_pairs_vs_shuffle];
            gg_p = stats.wilcoxon(grid_pairs_vs_shuffle)[1]; gg_p= get_p_text(gg_p, ns=False)
            gng_p = stats.wilcoxon(grid_non_grid_pairs_vs_shuffle)[1]; gng_p= get_p_text(gng_p, ns=False)
            ngng_p = stats.wilcoxon(non_grid_non_grid_pairs_vs_shuffle)[1]; ngng_p= get_p_text(ngng_p, ns=False)

            gg_ngng_p = stats.mannwhitneyu(grid_pairs_vs_shuffle, non_grid_non_grid_pairs_vs_shuffle)[1]
            gg_gng_p = stats.mannwhitneyu(grid_pairs_vs_shuffle, grid_non_grid_pairs_vs_shuffle)[1]
            gng_ngng_p = stats.mannwhitneyu(grid_non_grid_pairs_vs_shuffle, non_grid_non_grid_pairs_vs_shuffle)[1]

            fig, ax = plt.subplots(figsize=(6,6))
            vp = ax.violinplot(data, [2, 4, 6], widths=2,
                               showmeans=False, showmedians=True, showextrema=False)
            ax.axhline(y=0, linestyle="dashed", color="black")
            significance_bar(ax, start=2, end=4, height=0.5, displaystring=get_p_text(gg_gng_p))
            significance_bar(ax, start=4, end=6, height=0.55, displaystring=get_p_text(gng_ngng_p))
            significance_bar(ax, start=2, end=6, height=0.6, displaystring=get_p_text(gg_ngng_p))
            ax.set_xticks([2,4,6])
            ax.set_yticks([-0.4, -0.2, 0, 0.2, 0.4, 0.6])
            ax.set_ylim([-0.4, 0.6])
            ax.set_xticklabels(["G-G", "G-NG", "NG-NG"])
            ax.set_ylabel("Change in spatial\ncorrelation vs shuffle", fontsize=20, labelpad=10)
            ax.spines['top'].set_visible(False)
            ax.spines['right'].set_visible(False)
            ax.spines['bottom'].set_visible(False)
            ax.xaxis.set_tick_params(length=0)
            ax.tick_params(axis='both', which='major', labelsize=25)
            fig.tight_layout()
            plt.savefig(output_path + '/joint_alignment_correlations_violin_shuffle_vs_real_'+hmt+'_'+str(tt)+'png', dpi=300)
            plt.close()

def plot_paired_vs_shuffle_by_tt(df, output_path):
    df = df[df["hit_miss_try"]=="hit"]
    b_df = df[df["trial_type"]==0]
    nb_df = df[df["trial_type"]==1]
    for collumn in ["grid_pairs_vs_shuffle", "grid_non_grid_pairs_vs_shuffle", "non_grid_non_grid_pairs_vs_shuffle"]:
        b_pairs_vs_shuffle = np.array([])
        nb_pairs_vs_shuffle = np.array([])

        for index, session_row in b_df.iterrows():
            session_row = session_row.to_frame().T.reset_index(drop=True)
            b_pairs_vs_shuffle = np.concatenate((b_pairs_vs_shuffle, session_row.iloc[0][collumn]))
        for index, session_row in nb_df.iterrows():
            session_row = session_row.to_frame().T.reset_index(drop=True)
            nb_pairs_vs_shuffle = np.concatenate((nb_pairs_vs_shuffle, session_row.iloc[0][collumn]))

        b_pairs_vs_shuffle = b_pairs_vs_shuffle[~np.isnan(b_pairs_vs_shuffle)]
        nb_pairs_vs_shuffle = nb_pairs_vs_shuffle[~np.isnan(nb_pairs_vs_shuffle)]

        data = [b_pairs_vs_shuffle, nb_pairs_vs_shuffle];
        b_p = stats.wilcoxon(b_pairs_vs_shuffle)[1]; b_p= get_p_text(b_p, ns=False)
        nb_p = stats.wilcoxon(nb_pairs_vs_shuffle)[1]; nb_p= get_p_text(nb_p, ns=False)

        bnb_p = stats.mannwhitneyu(b_pairs_vs_shuffle, nb_pairs_vs_shuffle)[1]

        fig, ax = plt.subplots(figsize=(6,6))
        vp = ax.violinplot(data, [2, 4], widths=2,
                           showmeans=False, showmedians=True, showextrema=False)
        ax.axhline(y=0, linestyle="dashed", color="black")
        significance_bar(ax, start=2, end=4, height=0.6, displaystring=get_p_text(bnb_p))
        ax.set_xticks([2,4])
        ax.set_yticks([-0.4, -0.2, 0, 0.2, 0.4, 0.6])
        ax.set_ylim([-0.4, 0.6])
        ax.set_xticklabels(["Cued", "PI"])
        ax.set_ylabel("Change in spatial\ncorrelation vs shuffle", fontsize=20, labelpad=10)
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['bottom'].set_visible(False)
        ax.xaxis.set_tick_params(length=0)
        ax.tick_params(axis='both', which='major', labelsize=25)
        fig.tight_layout()
        plt.savefig(output_path + '/joint_alignment_correlations_violin_shuffle_vs_real_'+collumn+'_by_tt.png', dpi=300)
        plt.close()

def plot_paired_vs_shuffle_by_hmt(df, output_path):
    df = df[df["trial_type"]==1]
    hit_df = df[df["hit_miss_try"]=="hit"]
    try_df = df[df["hit_miss_try"]=="try"]
    miss_df = df[df["hit_miss_try"]=="miss"]
    for collumn in ["grid_pairs_vs_shuffle", "grid_non_grid_pairs_vs_shuffle", "non_grid_non_grid_pairs_vs_shuffle"]:
        hit_pairs_vs_shuffle = np.array([])
        try_pairs_vs_shuffle = np.array([])
        miss_pairs_vs_shuffle = np.array([])

        for index, session_row in hit_df.iterrows():
            session_row = session_row.to_frame().T.reset_index(drop=True)
            hit_pairs_vs_shuffle = np.concatenate((hit_pairs_vs_shuffle, session_row.iloc[0][collumn]))
        for index, session_row in try_df.iterrows():
            session_row = session_row.to_frame().T.reset_index(drop=True)
            try_pairs_vs_shuffle = np.concatenate((try_pairs_vs_shuffle, session_row.iloc[0][collumn]))
        for index, session_row in miss_df.iterrows():
            session_row = session_row.to_frame().T.reset_index(drop=True)
            miss_pairs_vs_shuffle = np.concatenate((miss_pairs_vs_shuffle, session_row.iloc[0][collumn]))

        hit_pairs_vs_shuffle = hit_pairs_vs_shuffle[~np.isnan(hit_pairs_vs_shuffle)]
        try_pairs_vs_shuffle = try_pairs_vs_shuffle[~np.isnan(try_pairs_vs_shuffle)]
        miss_pairs_vs_shuffle = miss_pairs_vs_shuffle[~np.isnan(miss_pairs_vs_shuffle)]

        data = [hit_pairs_vs_shuffle, try_pairs_vs_shuffle, miss_pairs_vs_shuffle];
        h_p = stats.wilcoxon(hit_pairs_vs_shuffle)[1]; h_p= get_p_text(h_p, ns=False)
        t_p = stats.wilcoxon(try_pairs_vs_shuffle)[1]; t_p= get_p_text(t_p, ns=False)
        m_p = stats.wilcoxon(miss_pairs_vs_shuffle)[1]; m_p= get_p_text(m_p, ns=False)

        hm_p = stats.mannwhitneyu(hit_pairs_vs_shuffle, miss_pairs_vs_shuffle)[1]
        ht_p = stats.mannwhitneyu(hit_pairs_vs_shuffle, try_pairs_vs_shuffle)[1]
        mt_p = stats.mannwhitneyu(miss_pairs_vs_shuffle, try_pairs_vs_shuffle)[1]

        fig, ax = plt.subplots(figsize=(6,6))
        vp = ax.violinplot(data, [2, 4, 6], widths=2,
                           showmeans=False, showmedians=True, showextrema=False)
        ax.axhline(y=0, linestyle="dashed", color="black")
        significance_bar(ax, start=2, end=4, height=0.5, displaystring=get_p_text(ht_p))
        significance_bar(ax, start=4, end=6, height=0.55, displaystring=get_p_text(mt_p))
        significance_bar(ax, start=2, end=6, height=0.6, displaystring=get_p_text(hm_p))
        ax.set_xticks([2,4,6])
        ax.set_yticks([-0.4, -0.2, 0, 0.2, 0.4, 0.6])
        ax.set_ylim([-0.4, 0.6])
        ax.set_xticklabels(["Hit", "Try", "Miss"])
        ax.set_ylabel("Change in spatial\ncorrelation vs shuffle", fontsize=20, labelpad=10)
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['bottom'].set_visible(False)
        ax.xaxis.set_tick_params(length=0)
        ax.tick_params(axis='both', which='major', labelsize=25)
        fig.tight_layout()
        plt.savefig(output_path + '/joint_alignment_correlations_violin_shuffle_vs_real_'+collumn+'_by_hmt.png', dpi=300)
        plt.close()

def get_n_simultaneously_recorded_cells(grid_cells, group1, group2):
    #how many group2 cells have been simulatenously recorded with group1 cells
    n =0
    session_ids = np.unique(grid_cells["session_id"])
    for i in range(len(session_ids)):
        session_id = session_ids[i]
        single_session = grid_cells[grid_cells["session_id"]==session_id]
        n_group1 = len(single_session[single_session["Lomb_classifier_"] == group1])
        n_group2 = len(single_session[single_session["Lomb_classifier_"] == group2])

        if group1==group2:
            if n_group1>1:
               n+=n_group1
        else:
            if n_group1>0:
                n+=n_group2
    return n

def plot_n_cells_simulatenously_recorded(concantenated_dataframe,  save_path, normalised=False):
    grid_cells = concantenated_dataframe[concantenated_dataframe["classifier"] == "G"]
    grid_cells = add_lomb_classifier(grid_cells, suffix="")

    P_proportion = len(grid_cells[grid_cells["Lomb_classifier_"] == "Position"])/len(grid_cells)
    D_proportion = len(grid_cells[grid_cells["Lomb_classifier_"] == "Distance"])/len(grid_cells)
    N_proportion = len(grid_cells[grid_cells["Lomb_classifier_"] == "Null"])/len(grid_cells)
    if not normalised:
        P_proportion = 1
        D_proportion = 1
        N_proportion = 1

    objects = ["Position", "Distance", "Null"]
    x_pos = np.arange(len(objects))
    fig, axes = plt.subplots(3, 1, figsize=(6,6), sharex=True)
    for ax, group, color in zip(axes, objects, [Settings.allocentric_color, Settings.egocentric_color, Settings.null_color]):
        n_p = get_n_simultaneously_recorded_cells(grid_cells, group1=group, group2="Position")/P_proportion
        n_d = get_n_simultaneously_recorded_cells(grid_cells, group1=group, group2="Distance")/D_proportion
        n_n = get_n_simultaneously_recorded_cells(grid_cells, group1=group, group2="Null")/N_proportion

        ax.bar(x_pos[0], n_p, color=Settings.allocentric_color, edgecolor="black")
        ax.bar(x_pos[1], n_d, color=Settings.egocentric_color, edgecolor="black")
        ax.bar(x_pos[2], n_n, color=Settings.null_color, edgecolor="black")

        Edmond.plot_utility2.style_vr_plot(ax)
        ax.tick_params(axis='both', which='both', labelsize=20)
        ax.set_ylabel(group, fontsize=25)
        ax.set_yticks([0, int(np.round(ax.get_ylim()[1]))])
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')
    ax.set_xticks(x_pos)
    ax.set_xticklabels(["P", "D", "N"])
    fig.tight_layout(pad=2.0)
    plt.subplots_adjust(hspace = .35, wspace = .35,  bottom = 0.2, left = 0.3, right = 0.87, top = 0.92)
    if normalised:
        plt.savefig(save_path + '/simulatenous_recorded_cell_number_normalised.png', dpi=300)
    else:
        plt.savefig(save_path + '/simulatenous_recorded_cell_numbers.png', dpi=300)
    plt.close()

    return

def process_recordings(vr_recording_path_list, of_recording_path_list):

    matched_recording_df = pd.DataFrame()
    for recording in vr_recording_path_list:
        print("processing ", recording)
        paired_recording, found_paired_recording = find_paired_recording(recording, of_recording_path_list)
        try:
            output_path = recording+'/'+settings.sorterName
            position_data = pd.read_pickle(recording+"/MountainSort/DataFrames/position_data.pkl")
            raw_position_data, position_data = syncronise_position_data(recording, get_track_length(recording))

            position_data = add_time_elapsed_collumn(position_data)
            spike_data = pd.read_pickle(recording+"/MountainSort/DataFrames/spatial_firing.pkl")
            processed_position_data = pd.read_pickle(recording+"/MountainSort/DataFrames/processed_position_data.pkl")
            processed_position_data = add_avg_track_speed(processed_position_data, track_length=get_track_length(recording))
            processed_position_data, _ = add_hit_miss_try3(processed_position_data, track_length=get_track_length(recording))

            if paired_recording is not None:
                of_spike_data = pd.read_pickle(paired_recording+"/MountainSort/DataFrames/spatial_firing.pkl")
                spike_data = plot_joint_cell_correlations(spike_data, of_spike_data, processed_position_data, position_data, raw_position_data, output_path, get_track_length(recording))
                plot_firing_rate_maps_per_trial_by_hmt_aligned_other_neuron(spike_data=spike_data, processed_position_data=processed_position_data, output_path=output_path, track_length=get_track_length(recording), trial_types=[1])
                matched_recording_df = plot_joint_jitter_correlations(spike_data, of_spike_data, processed_position_data, position_data, output_path, get_track_length(recording), matched_recording_df)
                matched_recording_df.to_pickle("/mnt/datastore/Harry/cohort8_may2021/matched_grid_recording_df.pkl")
                #plot_joint_cell_cross_correlations(spike_data, output_path)

            #spike_data.to_pickle(recording+"/MountainSort/DataFrames/spatial_firing.pkl")

            print("successfully processed and saved vr_grid analysis on "+recording)
        except Exception as ex:
            print('This is what Python says happened:')
            print(ex)
            exc_type, exc_value, exc_traceback = sys.exc_info()
            traceback.print_tb(exc_traceback)
            print("couldn't process vr_grid analysis on "+recording)

def main():
    print('-------------------------------------------------------------')

    # give a path for a directory of recordings or path of a single recording
    vr_path_list = [f.path for f in os.scandir("/mnt/datastore/Harry/cohort8_may2021/vr") if f.is_dir()]
    of_path_list = [f.path for f in os.scandir("/mnt/datastore/Harry/cohort8_may2021/of") if f.is_dir()]
    #vr_path_list = [f.path for f in os.scandir("/mnt/datastore/Harry/cohort7_october2020/vr") if f.is_dir()]
    #of_path_list = [f.path for f in os.scandir("/mnt/datastore/Harry/cohort7_october2020/of") if f.is_dir()]
    #vr_path_list = [f.path for f in os.scandir("/mnt/datastore/Harry/cohort6_july2020/vr") if f.is_dir()]
    #of_path_list = [f.path for f in os.scandir("/mnt/datastore/Harry/cohort6_july2020/of") if f.is_dir()]

    # all of these recordings have at least 2 grid cells recorded
    vr_path_list = ['/mnt/datastore/Harry/cohort8_may2021/vr/M11_D11_2021-05-24_10-00-53', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D12_2021-05-25_09-49-23',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D13_2021-05-26_09-46-36', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D15_2021-05-28_10-42-15',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D16_2021-05-31_10-21-05', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D17_2021-06-01_10-36-53',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D18_2021-06-02_10-36-39', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D22_2021-06-08_10-55-28',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D26_2021-06-14_10-34-14', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D29_2021-06-17_10-35-48',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D30_2021-06-18_10-46-48', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D33_2021-06-23_11-08-03',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D35_2021-06-25_12-02-52', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D36_2021-06-28_12-04-36',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D37_2021-06-29_11-50-02', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D38_2021-06-30_11-54-56',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D3_2021-05-12_09-37-41', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D41_2021-07-05_12-05-02',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D44_2021-07-08_12-03-21', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D45_2021-07-09_11-39-02',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M12_D6_2021-05-17_10-26-15', '/mnt/datastore/Harry/cohort8_may2021/vr/M13_D17_2021-06-01_11-45-20',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M13_D24_2021-06-10_12-01-54', '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D12_2021-05-25_11-03-39',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D15_2021-05-28_12-29-15', '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D16_2021-05-31_12-01-35',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D20_2021-06-04_12-20-57', '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D27_2021-06-15_12-21-58',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D31_2021-06-21_12-07-01', '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D35_2021-06-25_12-41-16',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D37_2021-06-29_12-33-24']

    vr_path_list = ['/mnt/datastore/Harry/cohort8_may2021/vr/M11_D35_2021-06-25_12-02-52', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D36_2021-06-28_12-04-36',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D37_2021-06-29_11-50-02', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D38_2021-06-30_11-54-56',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D3_2021-05-12_09-37-41', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D41_2021-07-05_12-05-02',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D44_2021-07-08_12-03-21', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D45_2021-07-09_11-39-02',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M12_D6_2021-05-17_10-26-15', '/mnt/datastore/Harry/cohort8_may2021/vr/M13_D17_2021-06-01_11-45-20',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M13_D24_2021-06-10_12-01-54', '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D12_2021-05-25_11-03-39',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D15_2021-05-28_12-29-15', '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D16_2021-05-31_12-01-35',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D20_2021-06-04_12-20-57', '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D27_2021-06-15_12-21-58',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D31_2021-06-21_12-07-01', '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D35_2021-06-25_12-41-16',
                    '/mnt/datastore/Harry/cohort8_may2021/vr/M14_D37_2021-06-29_12-33-24']
    #vr_path_list = ['/mnt/datastore/Harry/cohort8_may2021/vr/M11_D36_2021-06-28_12-04-36']
    #vr_path_list = ['/mnt/datastore/Harry/cohort8_may2021/vr/M11_D22_2021-06-08_10-55-28', '/mnt/datastore/Harry/cohort8_may2021/vr/M11_D36_2021-06-28_12-04-36']
    process_recordings(vr_path_list, of_path_list)

    combined_df = pd.read_pickle("/mnt/datastore/Harry/Vr_grid_cells/combined_cohort8.pkl")
    combined_df = add_lomb_classifier(combined_df,suffix="")
    combined_df = add_percentage_for_lomb_classes(combined_df)

    # load df for plot_all_paired_vs_shuffle
    plot_all_paired_vs_shuffle(pd.read_pickle("/mnt/datastore/Harry/cohort8_may2021/matched_grid_recording_df.pkl"), output_path= "/mnt/datastore/Harry/Vr_grid_cells/joint_activity")
    plot_paired_vs_shuffle_by_hmt(pd.read_pickle("/mnt/datastore/Harry/cohort8_may2021/matched_grid_recording_df.pkl"), output_path= "/mnt/datastore/Harry/Vr_grid_cells/joint_activity")
    plot_paired_vs_shuffle_by_tt(pd.read_pickle("/mnt/datastore/Harry/cohort8_may2021/matched_grid_recording_df.pkl"), output_path= "/mnt/datastore/Harry/Vr_grid_cells/joint_activity")

    #grid_cells = combined_df[combined_df["classifier"] == "G"]
    #grid_cells_from_same_recording = get_grid_cells_from_same_recording(grid_cells)
    #plot_class_prection_credence(grid_cells_from_same_recording, save_path="/mnt/datastore/Harry/Vr_grid_cells/joint_activity")
    plot_n_cells_simulatenously_recorded(combined_df,  save_path="/mnt/datastore/Harry/Vr_grid_cells/joint_activity")
    plot_n_cells_simulatenously_recorded(combined_df,  save_path="/mnt/datastore/Harry/Vr_grid_cells/joint_activity", normalised=True)

    print("look now")

if __name__ == '__main__':
    main()


from django.urls import path
from .views import *
from django.views.decorators.cache import cache_page

urlpatterns = [
    path('ad/<int:pk>', AdDetail.as_view(), name = 'ad_detail'),
    path('ad/<int:pk>/update', UpdateAd.as_view(), name = 'update_ad'),
    path('ad/create', CreateAd.as_view(), name = 'create_ad'),
    path('', AdList.as_view(), name = 'ad_list'),
    path('ad/<int:pk>/create_answer', CreateAnswer.as_view(), name = 'create_answer'),
    path('answer/tomyad', AnswersToMyAd.as_view(), name = 'answer_to_my_ad'),
    path('answer/my', MyAnswers.as_view(), name = 'my_answer'),
    path('answer/<int:pk>', AnswerDetail.as_view(), name = 'answer_detail'),
    path('answer/<int:pk>/confirm_answer', ConfirmAnswer.as_view(), name = 'confirm_answer'),
    path('answer/<int:pk>/delete_answer', DeleteAnswer.as_view(), name = 'delete_answer'),
]


from django import template

from questions.models import QuestionStatus

register = template.Library()

@register.filter
def question_status_class(question):
    return {
        QuestionStatus.DRAFT: 'primary',
        QuestionStatus.DEFERRED: 'warning',
        QuestionStatus.PENDING: 'info',
        QuestionStatus.APPROVED: 'success',
        QuestionStatus.REJECTED: 'danger',
        QuestionStatus.PUBLISHED: 'primary',
        QuestionStatus.ANSWERED: 'primary',
        QuestionStatus.CLOSED: 'secondary',
    }[question.status]

@register.filter
def question_status_icon(question):
    return {
        QuestionStatus.DRAFT: 'icons/clock.png',
        QuestionStatus.DEFERRED: 'icons/band.png',
        QuestionStatus.PENDING: 'icons/clock.png',
        QuestionStatus.APPROVED: 'icons/checkmark.png',
        QuestionStatus.REJECTED: 'icons/cross.png',
        QuestionStatus.PUBLISHED: 'icons/clock.png',
        QuestionStatus.ANSWERED: 'icons/clock.png',
    }.get(question.status, 'icons/question.png')


# -*- coding:utf-8 -*-

import pytest

from pageloader.url import UrlConverter

def test_to_filename():
    expected = 'hexlet-io-courses.html'
    url = UrlConverter('https://hexlet.io/courses')
    result = url.get_filename()
    assert expected == result

def test_to_resource_dirname():
    expected = 'hexlet-io-courses_files'
    url = UrlConverter('https://hexlet.io/courses')
    result = url.get_dirname()
    assert expected == result

def test_to_resource():
    expected = 'assets-application.css'
    url = UrlConverter('/assets/application.css')
    result = url.get_resourcename()
    assert expected == result


""" xx_investigate_dags.py

Playground to analyze / look at adgs

"""
import json
import yaml
import argparse
import pickle
from functools import partial
from pathlib import Path

import numpy as np
import ms_pred.common as common
import ms_pred.magma.fragmentation as fragmentation

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--num-workers", default=0, action="store", type=int)
    parser.add_argument("--dag-folder", action="store")
    parser.add_argument(
        "--min-inten",
        type=float,
        default=0,
        help="Minimum intensity to call a peak in prediction",
    )
    parser.add_argument(
        "--max-peaks", type=int, default=100, help="Max num peaks to call"
    )
    parser.add_argument("--num-bins", type=int, default=3000, help="Num bins")
    parser.add_argument("--upper-limit", type=int, default=1500, help="upper lim")
    return parser.parse_args()

def bin_dag(
    dag: dict, max_peaks: int, upper_limit: int, num_bins: int, min_inten: float
) -> np.ndarray:
    """bin_dag.

    Args:
        dag (dict): dag
        max_peaks (int): max_peaks
        upper_limit (int): upper_limit
        num_bins (int): num_bins
        min_inten (float): min_inten

    Returns:
        np.ndarray:
    """
    root = dag["root_inchi"]
    frags = dag["frags"]
    engine = fragmentation.FragmentEngine(mol_str=root, mol_str_type="inchi")

    mz, intens = [], []
    for k, v in frags.items():
        base_mass = engine.single_mass(v["frag"])
        new_masses = engine.shift_bucket_masses + base_mass
        new_intens = v["intens"]
        mz.extend(new_masses.tolist())
        intens.extend(new_intens)

    print(len(set(mz)), len(mz))

    # Cap at min
    mz, intens = np.array(mz), np.array(intens)
    min_inten_mask = intens > min_inten
    mz = mz[min_inten_mask]
    intens = intens[min_inten_mask]

    # Keep 100 max
    argsort_intens = np.argsort(intens)[::-1][:max_peaks]
    mz = mz[argsort_intens]
    intens = intens[argsort_intens]
    spec_ar = np.vstack([mz, intens]).transpose(1, 0)

    # Bin intensities
    binned = common.bin_spectra([spec_ar], num_bins, upper_limit, pool_fn="max")
    return binned, root

def dag_stats(
    dag_file: dict, max_peaks: int, upper_limit: int, num_bins: int, min_inten: float
):
    """bin_dag_file.

    Args:
        dag_file (dict): dag_file
        max_peaks (int): max_peaks
        upper_limit (int): upper_limit
        num_bins (int): num_bins
        min_inten (float): min_inten

    Returns:
        np.ndarray:
        str
    """
    dag = json.load(open(dag_file, "r"))
    root = dag["root_inchi"]
    frags = dag["frags"]
    engine = fragmentation.FragmentEngine(mol_str=root, mol_str_type="inchi")
    subforms = []
    for i in list(dag["frags"].values()):
        # subforms.append(i['form'])
        subforms.append(i["base_mass"])

    print(f"Len of subforms: {len(subforms)}")
    print(f"Len of unique subforms: {len(set(subforms))}")
    if (len((set(subforms))) - len(subforms)) < -5:
        import pdb

        pdb.set_trace()

    # mz, intens = [], []
    # for k, v in frags.items():
    #    base_mass = engine.single_mass(v['frag'])
    #    new_masses = engine.shift_bucket_masses + base_mass
    #    new_intens = v['intens']
    #    mz.extend(new_masses.tolist())
    #    intens.extend(new_intens)

    ## Cap at min
    # mz, intens = np.array(mz), np.array(intens)
    # min_inten_mask = intens > min_inten
    # mz = mz[min_inten_mask]
    # intens = intens[min_inten_mask]

    ## Keep 100 max
    # argsort_intens = np.argsort(intens)[::-1][:max_peaks]
    # mz = mz[argsort_intens]
    # intens = intens[argsort_intens]
    # spec_ar = np.vstack([mz, intens]).transpose(1, 0)

    ## Bin intensities
    # binned = common.bin_spectra([spec_ar], num_bins, upper_limit,
    #                            pool_fn="max")
    ##

def main():
    """main."""
    args = get_args()

    max_peaks, min_inten = args.max_peaks, args.min_inten
    num_bins, upper_limit = args.num_bins, args.upper_limit
    dag_folder = Path(args.dag_folder)
    dag_files = list(dag_folder.glob("*.json"))

    read_dag_file = partial(
        dag_stats,
        max_peaks=max_peaks,
        upper_limit=upper_limit,
        num_bins=num_bins,
        min_inten=min_inten,
    )

    dag_outs = [read_dag_file(i) for i in dag_files]

if __name__ == "__main__":
    main()


#find leaves
# Definition for a binary tree node.
from collections import List, defaultdict

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:
    def findLeaves(self, root: TreeNode) -> List[List[int]]:
        def order(root, dic):
            if not root:
                return 0
            left = order(root.left, dic)
            right = order(root.right, dic)
            level = max(left, right) + 1
            dic[level].append(root.val)
            return level
        dic = defaultdict(list)
        order(root, dic)
        idx = sorted(dic)
        return [dic[i] for i in idx]


# valor do dolar em relação ao real

import requests
from bs4 import BeautifulSoup

headers = {'user-agent': 'user-agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, '
                         'like Gecko) Chrome/99.0.4844.82 Safari/537.36'}
page = requests.get(
    'https://www.google.com/search?q=cota%C3%A7%C3%A3o+dolar&oq=cot&aqs=chrome.0'
    '.69i59j0i131i433i512j69i57j69i59j69i65j69i60l2j69i61.1077j0j7&sourceid=chrome&ie=UTF-8', headers=headers)

soup = BeautifulSoup(page.content, 'html.parser')
atributos = {'class': 'g'}

#class="DFlfde SwHCTb"

valor_dolar = soup.find_all('span', class_ = 'DFlfde SwHCTb')[0]

cotacao_dollar = float(valor_dolar['data-value'])

#print(cotacao_dollar)


from network import *
from frequency import Frequency
from media import RectangularWaveguide

import os
import pylab as plb
import numpy as npy
from scipy.constants import mil
from datetime import datetime

# pre-initialized classes

#frequency bands
f_wr10 	= Frequency(75,110,201, 'ghz')
f_wr3  	= Frequency(220,325,201, 'ghz')
f_wr2p2 = Frequency(330,500,201, 'ghz')
f_wr1p5 = Frequency(500,750,201, 'ghz')
f_wr1	= Frequency(750,1100,201, 'ghz')

# rectangular waveguides
wr10 	= RectangularWaveguide(Frequency(75,110,201, 'ghz'), 100*mil,z0=50)
wr3 	= RectangularWaveguide(Frequency(220,325,201, 'ghz'), 30*mil,z0=50)
wr2p2	= RectangularWaveguide(Frequency(330,500,201, 'ghz'), 22*mil,z0=50)
wr1p5 	= RectangularWaveguide(Frequency(500,750,201, 'ghz'), 15*mil,z0=50)
wr1 	= RectangularWaveguide(Frequency(750,1100,201, 'ghz'), 10*mil,z0=50)

## Functions
# Ploting
def save_all_figs(dir = './', format=['eps','pdf','png']):
	if dir[-1] != '/':
		dir = dir + '/'
	for fignum in plb.get_fignums():
		plb.figure(fignum)
		fileName = plb.gca().get_title()
		if fileName == '':
				fileName = 'unamedPlot'
		for fmt in format:
			plb.savefig(dir+fileName+'.'+fmt, format=fmt)
			print (dir+fileName+'.'+fmt)

def legend_off(ax=None):
	'''
	turn off the legend for a given axes. if no axes is given then
	it will use current axes.
	'''
	if ax is None:
		plb.gca().legend_.set_visible(0)
	else:
		ax.lengend_.set_visible(0)

def plot_complex(z,*args, **kwargs):
	'''
	plots a complex array or list in real vs imaginary.
	'''
	plb.plot(npy.array(z).real,npy.array(z).imag, *args, **kwargs)

# other
def now_string():
	return datetime.now().__str__().replace('-','.').replace(':','.').replace(' ','.')

def find_nearest(array,value):
	'''
	find nearest value in array.
	taken from  http://stackoverflow.com/questions/2566412/find-nearest-value-in-numpy-array

	'''
	idx=(npy.abs(array-value)).argmin()
	return array[idx]

def find_nearest_index(array,value):
	'''
	find nearest value in array.
	taken from  http://stackoverflow.com/questions/2566412/find-nearest-value-in-numpy-array

	'''
	return (npy.abs(array-value)).argmin()


# -*- coding: utf-8 -*-
# __author__ = 'qingchen.zhang'

import os,json
import linecache
import pandas as pd
from biocluster.agent import Agent
from biocluster.tool import Tool
from biocluster.core.exceptions import OptionError

class SeqStatAgent(Agent):
    """
    传入一个group表，以及是否进行样本合并的参数生成一张丰度表并对并依照group表OTU表进行筛选合并
    """
    def __init__(self, parent):
        super(SeqStatAgent, self).__init__(parent)
        options = [
            {"name": "total_genome", "type": "infile", "format": "sequence.profile_table"},  # 对应的参数列表
            {"name": "sequence_dir", "type": "string"},  # 序列seq的dir
            {"name": "raw_dir", "type": "string"},  # 序列的total的dir
        ]
        self.add_option(options)
        self.step.add_steps("seq_stat")
        self.on('start', self.start_seq_stat)
        self.on('end', self.end_seq_stat)

    def start_seq_stat(self):
        self.step.seq_stat.start()
        self.step.update()

    def end_seq_stat(self):
        self.step.seq_stat.finish()
        self.step.update()

    def check_options(self):
        """
        参数检测
        """
        if not self.option("total_genome").is_set:
            raise OptionError("必须提供输入的参数列表")

    def end(self):
        super(SeqStatAgent, self).end()

    def set_resource(self):
        """
        设置所需的资源
        """
        self._cpu = 2
        self._memory = "10G"

class SeqStatTool(Tool):
    def __init__(self, config):
        super(SeqStatTool, self).__init__(config)
        self.logger.info("开始进行seq_stat的序列统计")
        self.stat_seq = os.path.join(self.config.PACKAGE_DIR, "bac_comp_genome/seq_stat.py")
        self.python = "/miniconda2/bin/python"

    def seq_stat(self):
        """
        统计基因组水平的基础信息
        :return:
        """
        self.logger.info("开始用package进行计算统计")
        out_file = os.path.join(self.work_dir, "upload_file.xls")
        cmd = "{} {} -i {} -seq {} -total {} -o {}".format(self.python, self.stat_seq, self.option("total_genome").prop["path"], self.option("sequence_dir"), self.option("raw_dir"), out_file)
        self.logger.info(cmd)
        command_name = "seq_stat"
        command = self.add_command(command_name, cmd).run()
        self.wait(command)
        if command.return_code == 0:
            self.logger.info("统计序列成功")
            self.set_output()
        else:
            self.set_error("统计序列失败!")

    def set_output(self):
        """
        设置结果文件目录
        :return:
        """
        self.logger.info("正在生成结果文件目录")
        outfile = os.path.join(self.output_dir, "upload_file.xls")
        if os.path.exists(outfile):
            os.remove(outfile)
        os.link(os.path.join(self.work_dir, "upload_file.xls"), outfile)
        self.logger.info("完成结果文件的设置")

    def run(self):
        super(SeqStatTool, self).run()
        self.logger.info("开始运行seq_stat的tool！")
        self.seq_stat()
        self.set_output()
        self.end()


from src.common.admin.interfaces import IAdminUpdateUseCase
from src.common.exceptions.repository_exceptions import NotFoundException
from src.common.exceptions.use_case_exceptions import NotFoundHTTPException
from src.common.repository import ModelEntity, BaseRepo
from src.common.types.python_types import IdType, SchemaInType

class AdminUpdateUseCase(IAdminUpdateUseCase[IdType, SchemaInType, ModelEntity]):
    async def __call__(self, object_id: IdType, updated_object: SchemaInType) -> ModelEntity:
        async with self.uow:
            repository: BaseRepo = getattr(self.uow, self.repository_attr_name)
            try:
                entity = await repository.retrieve(self.filter_schema_class(id=object_id))
            except NotFoundException:
                raise NotFoundHTTPException

            for attr_name, attr_value in updated_object.dict(exclude_unset=True).items():
                setattr(entity, attr_name, attr_value)

            await self.uow.commit()
            await self.uow.refresh(entity)

        return entity


from flask import Flask, session
from pyhtml import html, title, body, p, table, tr, td, strong, head

app = Flask(__name__)

app.config['SECRET_KEY'] = 'Og5WDE6A8OidZG-fqsLgo2afWs18VIhmOVdMaYOCr3E'

@app.route('/')
def main():

    if 'visits' not in session:
        session['visits'] = 0
    session['visits'] += 1

    code = html(
        body(
            p(f"Welcome to your {session['visits']} visit.")
        )
    )
    return str(code)

if __name__ == "__main__":
    app.run(debug=True)


# import tensorflow as tf

import numpy as np
from tensorflow.keras import Input, Model  # , Sequential
# from tensorflow.keras.initializers import he_uniform
#import convolutional and Max pooling layers
from tensorflow.keras.layers import (Conv2D, Dense, Dropout, Flatten,
                                     MaxPooling2D)
from tensorflow.keras.regularizers import l2

def generate_model(num_classes,
					input_shape,
					num_filter = 32,
					act_func = "relu",
					kernel_size = (3, 3),
					layer_num = 3,
					drop_rate = 0.25):
	"""
	(num_classes,
				input_shape,
				num_filter = 32,
				act_func = "relu",
				kernel_size = (3, 3),
				layer_num = 3,
				drop_rate = 0.25,
				bath_norm = False)
	"""

	inputs = Input(input_shape)

	for i in range(1, layer_num+1):
		if i == 1:
			inptt = inputs
			x = Conv2D(i*num_filter, (5, 5),
							 strides=1,
							 input_shape = input_shape,
							 padding = 'same',
							 activation = act_func,
							 kernel_regularizer = l2(l=0.1),
							 kernel_initializer = 'he_normal')(inptt)
		else:
			inptt = x
			x = Conv2D((2*i-2)*num_filter, kernel_size,
							 strides=1,
							 input_shape = input_shape,
							 padding = 'same',
							 activation = act_func,
							 kernel_regularizer = l2(l=0.1),
							 kernel_initializer = 'he_normal')(inptt)
		# x = Conv2D(i*num_filter, kernel_size,
		# 					 input_shape = input_shape,
		# 					 padding = 'same',
		# 					 activation = act_func,
		# 					 kernel_initializer = he_uniform())(inptt)

		x = MaxPooling2D((2, 2), padding = 'same')(x)

		x = Dropout(drop_rate)(x, training=True)

	x = Flatten()(x)

	x = Dense((2*layer_num-2)*num_filter,
						 activation = act_func,
						 kernel_regularizer = l2(l=0.1),
						kernel_initializer = 'he_normal')(x)
	x = Dropout(drop_rate)(x, training=True)

	output = Dense(num_classes, kernel_regularizer = l2(l=0.1), activation = 'softmax')(x)

	return Model(inputs=inputs, outputs=output)

def get_epistemic_uncertainty(model, image, T=10):
	dim = image.shape
	p_hat = []
	epistemic = np.zeros(dim[0])
	for t in range(T):
		p_hat.append( model.predict(image) )

    # estimate uncertainties (eq. 4 )
    # eq.4 in https://openreview.net/pdf?id=Sk_P2Q9sG
	p_hat = np.array(p_hat)
    # mean prediction
	prediction = []
	p_bar = np.mean(p_hat, axis=0)
	for i in range(dim[0]):
		prediction.append(np.where( p_bar[i] > 0.5))
		for j in range(T):
			epistemic[i] += np.dot((p_hat[j,i] -p_bar[i]), (p_hat[j,i]-p_bar[i]))/T

	return prediction, epistemic


def linear_search(search_list, target_value):
  for idx in range(len(search_list)):
    if search_list[idx] == target_value:
      return idx
  raise ValueError("{0} not in list".format(target_value))

number_list = [ 10, 14, 19, 26, 27, 31, 33, 35, 42, 44]
target_number = 33

try:
  result = linear_search(number_list, target_number)
  print(result)

  result2 = linear_search(number_list, 100)
  print(result2)

except ValueError as error_message:
  print("{0}".format(error_message))

def linear_search_multiple(search_list, target_value):
  matches = []
  for idx in range(len(search_list)):
    if search_list[idx] == target_value:
      matches.append(idx)
  if matches:
    return matches
  else:
    raise ValueError("{0} not in list".format(target_value))

#Function call
tour_locations = [ "New York City", "Los Angeles", "Bangkok", "Istanbul", "London", "New York City", "Toronto"]
target_city = "New York City"

tour_stops = linear_search(tour_locations, target_city)
print(tour_stops)

# Max
def linear_search_max(search_list):
  maximum_score_index = None
  for idx in range(len(search_list)):
    print(search_list[idx])
    if (maximum_score_index is None) or (search_list[idx] > search_list[maximum_score_index]):
      maximum_score_index = idx

  return maximum_score_index

# Function call
test_scores = [88, 93, 75, 100, 80, 67, 71, 92, 90, 83]
highest_score = linear_search(test_scores)
print(highest_score)


# SIR as a compartmented model
#
# Copyright (C) 2017--2020 Simon Dobson
#
# This file is part of epydemic, epidemic network simulations in Python.
#
# epydemic is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# epydemic is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with epydemic. If not, see <http://www.gnu.org/licenses/gpl.html>.

from epydemic import CompartmentedModel
import sys
import random as rnd
import numpy as np
from networkx import community

if sys.version_info >= (3, 8):
    from typing import Final, Dict, Any
else:
    # backport compatibility with older typing
    from typing import Dict, Any
    from typing_extensions import Final

class SIR_Q(CompartmentedModel):
    '''The Susceptible-Infected-Removed :term:`compartmented model of disease`.
    Susceptible nodes are infected by infected neighbours, and recover to
    removed.'''

    # Model parameters
    P_INFECTED: Final[str] = 'epydemic.SIR.pInfected'  #: Parameter for probability of initially being infected.
    P_INFECT: Final[str] = 'epydemic.SIR.pInfect'  #: Parameter for probability of infection on contact.
    P_REMOVE: Final[str] = 'epydemic.SIR.pRemove'  #: Parameter for probability of removal (recovery).

    # Possible dynamics states of a node for SIR dynamics
    SUSCEPTIBLE: Final[str] = 'epydemic.SIR.S'  #: Compartment for nodes susceptible to infection.
    INFECTED: Final[str] = 'epydemic.SIR.I'  #: Compartment for nodes infected.
    REMOVED: Final[str] = 'epydemic.SIR.R'  #: Compartment for nodes recovered/removed.

    # Locus containing the edges at which dynamics can occur
    SI: Final[str] = 'epydemic.SIR.SI'  #: Edge able to transmit infection.

    # Dict of infected individuals
    infections = dict()
    # Dict of new cases by day
    cases = dict()
    # Lockdown thresholds
    thresh_r: int
    thresh_c: int
    # List of removed edges by a lockdown
    removed_edges = []
    # Lockdown Obedience Parameters
    variance: bool
    mean: float
    deviation = 0.1
    cluster: bool

    # State tracking variables
    locked_down_c: bool
    locked_down_r: bool
    isolation: float
    tracing: float

    trace_dict = dict()

    communities: list

    def __init__(self):
        super(SIR_Q, self).__init__()
        self.locked_down_c = False
        self.locked_down_r = False
        self.thresh_r = -1
        self.thresh_c = -1
        self.removed_edges = []
        self.cases = dict()
        self.infections = dict()
        self.trace_dict = dict()
        self.variance = False
        self.cluster = False
        self.deviation = 0.1
        self.communities = []
        self.isolation = -1
        self.tracing = -1

    def set_r(self, thresh_r):
        self.thresh_r = float(thresh_r)

    def set_cases(self, thresh_c):
        self.thresh_c = float(thresh_c)

    def set_isolation(self, iso):
        self.isolation = iso

    def set_tracing(self, tracing):
        self.tracing = tracing

    def check_variance(self, mean):
        self.variance = True
        self.cluster = False
        self.mean = mean

    def check_cluster(self, mean, G):
        self.cluster = True
        self.variance = False
        self.get_communities(G)
        self.mean = mean

    def build(self, params: Dict[str, Any]):
        '''Build the SIR model.

        :param params: the model parameters'''
        super(SIR_Q, self).build(params)

        pInfected = params[self.P_INFECTED]
        pInfect = params[self.P_INFECT]
        pRemove = params[self.P_REMOVE]

        self.addCompartment(self.SUSCEPTIBLE, 1 - pInfected)
        self.addCompartment(self.INFECTED, pInfected)
        self.addCompartment(self.REMOVED, 0.0)

        self.trackEdgesBetweenCompartments(self.SUSCEPTIBLE, self.INFECTED, name=self.SI)
        self.trackNodesInCompartment(self.INFECTED)

        self.addEventPerElement(self.SI, pInfect, self.infect)
        self.addEventPerElement(self.INFECTED, pRemove, self.remove)

        if self.thresh_c >= 0:
            self.postRepeatingEvent(1, 1, self.SI, self.check_cases)
        if self.thresh_r >= 0:
            self.postRepeatingEvent(1, 1, self.SI, self.check_r)

    def infect(self, t: float, e: Any):
        '''Perform an infection event. This changes the compartment of
        the susceptible-end node to :attr:`INFECTED`. It also marks the edge
        traversed as occupied.

        :param t: the simulation time
        :param e: the edge transmitting the infection, susceptible-infected'''
        (n, _) = e
        self.changeCompartment(n, self.INFECTED)
        self.markOccupied(e, t)
        self.update_tracking(t, e)
        if self.isolation > 0:
            self.postEvent(t + 5, e, self.isolate)

    def remove(self, t: float, n: Any):
        '''Perform a removal event. This changes the compartment of
        the node to :attr:`REMOVED`.

        :param t: the simulation time (unused)
        :param n: the node'''
        self.changeCompartment(n, self.REMOVED)
        self.remove_tracking(t, n)

    def isolate(self, t: float, e: any):
        '''Perform an isolation event. This removes all contacts
        from the node and posts a tracing event if appropriate.

        :param t: the simulation time (unused)
        :param e: the edge'''
        (n, _) = e
        r = rnd.random()
        if r < self.isolation:
            neighbors = list(self.network().neighbors(n))
            for each in neighbors:
                self.removeEdge(n, each)
            if self.tracing > 0:
                pass_param = [n, neighbors]
                self.postEvent(t + 1, pass_param, self.trace)

    def trace(self, t: float, e: any):
        '''Perform a Trace event. This traces the percentage
        of contacts input by the user. Traced contacts lose
        all contacts.

        :param t: the simulation time (unused)
        :param e: an object containing the node and its neighbours'''
        n = e[0]
        neighbours = e[1]
        edges = self.network()
        remove_int = round(len(neighbours) * self.tracing)
        to_remove = rnd.sample(neighbours, remove_int)
        for each in to_remove:
            self.quarantine(t, each)
        self.postEvent(t + 14, to_remove, self.undo_trace)

    def quarantine(self, t: float, n: any):
        '''Perform a quarantine. This function removes all the
        neighbours of a single node.

        :param t: the simulation time (unused)
        :param n: the node'''
        neighbors = list(self.network().neighbors(n))
        for each in neighbors:
            self.removeEdge(n, each)
        self.trace_dict[n] = neighbors

    def undo_trace(self, t: float, neighbours: any):
        '''Perform an trace reversal. This mimics the end of a quarantine
        for a node. Add back edges using trace dictionary.

        :param t: the simulation time (unused)
        :param n: the node'''
        for each in neighbours:
            if each in self.trace_dict:
                child_neighbours = self.trace_dict.pop(each)
                for i in child_neighbours:
                    self.addEdge(each, i)

    def update_tracking(self, t: float, e: Any):
        '''Update the tracking of infections

        :param t: the simulation time
        :param e: the edge between infected nodes'''
        (n, x) = e
        self.infections[n] = []
        # Update tracking of a single infection for R
        if x in self.infections:
            val = self.infections[x]
            val.append(n)
            self.infections[x] = val
        time = int(t)
        # Update daily number of infections
        if time in self.cases:
            self.cases[time] += 1
        else:
            self.cases[time] = 1

    def remove_tracking(self, t: float, n: Any):
        '''Perform a removal from tracking. Don't track
            removed nodes for R
        :param t: the simulation time (unused)
        :param n: the node to remove from tracking'''
        if n in self.infections:
            self.infections.pop(n)

    def check_r(self, t: float, n: Any):
        '''Perform an R check. Use r information calculated from tracking to dictate
                lockdowns.

        :param t: the simulation time
        :param n: the node (unused)'''
        # Calculate total infections per infected node
        total = 0
        r = 0
        if len(self.infections) > 0:
            for key, value in self.infections.items():
                total += len(self.infections[key])
            r = total / len(self.infections)
        if len(self.infections) > self.network().number_of_nodes() * 0.01:
            if r > self.thresh_r and not self.locked_down_r and not self.locked_down_c and len(
                    self.removed_edges) == 0:  # Instigate a lockdown
                self.postEvent(t + 7, self.SI, self.lockdown)
                self.locked_down_r = True
            if r < self.thresh_r and self.locked_down_r and 0 < len(self.removed_edges):  # Remove a lockdown
                self.postEvent(t + 7, self.SI, self.ease_lockdown)
                self.locked_down_r = False

    def check_cases(self, t: float, n: Any):
        '''Perform a cases check. Use new case information calculated from tracking to dictate
                lockdowns.

        :param t: the simulation time
        :param n: the node (unused)'''
        if 1 <= len(self.cases):
            if len(self.cases) < t:
                self.fix_cases(t)
            cases = self.cases[int(t) - 1]
            size = self.network().number_of_nodes()
            if (cases / size) > self.thresh_c and not self.locked_down_c and not self.locked_down_r and \
                    len(self.removed_edges) == 0:  # Instigate a lockdown
                self.postEvent(t + 7, self.SI, self.lockdown)
                self.locked_down_c = True
            if (cases / size) < self.thresh_c and self.locked_down_c and 0 < len(self.removed_edges):  # Remove a lockdown
                flag = True
                for each in range(2, 14):  # Check several days as this is a volatile metric
                    if int(t) - each >= 0:
                        if self.cases[int(t) - each] / size > self.thresh_c:
                            flag = False
                if flag:
                    self.postEvent(t + 7, self.SI, self.ease_lockdown)
                    self.locked_down_c = False

    def lockdown(self, t: float, n: Any):
        '''Perform a lockdown. Remove and record edge removals

        :param t: the simulation time (unused)
        :param n: the node (unused)'''
        if self.variance:
            self.lockdown_obedience_variance()
        elif self.cluster:
            self.lockdown_obedience_cluster()
        else:
            edges = self.network().edges
            remove_number = round(len(edges) * 0.9)
            to_remove = rnd.sample(edges, remove_number)
            for edge in to_remove:
                x = edge[0]
                y = edge[1]
                self.removeEdge(x, y)
                self.removed_edges.append(edge)

    def ease_lockdown(self, t: float, n: Any):
        '''Perform a lockdown removal. Add back any removed edges

        :param t: the simulation time (unused)
        :param n: the node (unused)'''
        for edge in self.removed_edges:
            self.network().add_edge(edge[0], edge[1])
        self.removed_edges = []

    def lockdown_obedience_variance(self):
        '''
        This event begins a lockdown. This lockdown is dictated by the individual nodes
        and the variance selected.
        '''
        nodes = self.network().number_of_nodes()
        remove_list = []
        for i in range(0, nodes):
            remove_percentage = np.random.normal(self.mean, self.deviation)
            if remove_percentage > 1:
                remove_percentage = 1
            if remove_percentage < 0:
                remove_percentage = 0
            self.remove_set(i, remove_percentage, remove_list)
        self.remove_edges(remove_list)

    def lockdown_obedience_cluster(self):
        '''
        This event begins a lockdown. This lockdown is dictated by the communities
        and the variance selected.
        '''
        remove_list = []
        for i in self.communities:
            remove_percentage = np.random.normal(self.mean, self.deviation)
            if remove_percentage > 1:
                remove_percentage = 1
            if remove_percentage < 0:
                remove_percentage = 0
            remove_list.extend(self.community_removal(list(i), remove_percentage))
        self.remove_edges(remove_list)

    def community_removal(self, com: list, percentage: float):
        '''
        This function creates a list of edges in a community and then takes a random sample
        that are to later be removed

        :param com: list of nodes in the communites
        :param percentage: percentage of edges to remove
        :return: list of edges to be removed
        '''
        edges = []
        for i in com:
            neighbours = self.network().neighbors(i)
            for j in neighbours:
                if j <= i:
                    if [j, i] not in edges:
                        edges.append([j, i])
                else:
                    if [i, j] not in edges:
                        edges.append([i, j])
        edges = rnd.sample(edges, int(percentage * len(edges)))
        return edges

    def remove_set(self, node: int, to_remove: int, remove_set: list):
        '''
        This method take a number and node and creates a list of edges
        to be removed of size 'to_remove'.

        :param node: the node to remove nodes
        :param to_remove: number of edges to remove
        :param remove_set: the list to be updated
        '''
        edges = list(self.network().neighbors(node))
        remove_int = round(to_remove * len(edges))
        remove_list = rnd.sample(edges, remove_int)
        for i in remove_list:
            remove_set.append([node, i])

    def remove_edges(self, edge_list: list):
        '''
        This is a simple function that removes the edges passed to it

        :param edge_list: the list of edges to remove
        '''
        for edge in edge_list:
            if self.network().has_edge(edge[0], edge[1]):
                self.removeEdge(edge[0], edge[1])
                self.removed_edges.append([edge[0], edge[1]])

    def get_communities(self, g):
        '''
        This function splits the graph into communities. The communities are determined
        by the size of the input graph.

        :param g: the input graph
        '''
        size = g.number_of_nodes()
        if size > 50000:
            k = int(size / 500)
        elif size < 5000:
            k = int(size / 50)
        else:
            k = 100
        self.communities = community.asyn_fluidc(g, k)

    def fix_cases(self, t):
        '''
        Fix cases is a function that fills missing points in the cases array.
        If nobody s infected on a given day then the dictionary entry is not
        created. The function finds these days and inserts them.
        :param t: the point in time
        '''
        for i in range(0, t):
            if i not in self.cases:
                self.cases[i] = 0


import os
import subprocess as sub

def qsub(pbsfile):
    cmd = "qsub " + pbsfile
    sub.call(cmd, shell=True)

# def qsub(ppn,pbsfile):
#     cmd = "qsub -q middle -V -l nodes=1:ppn="+str(ppn)+" "+pbsfile
#     sub.call(cmd, shell=True)

samples = []
for i in os.listdir("/home/rzli/swxxxjz/bam/10"):
    i = i.split("_")
    sample = i[0]
    samples.append(sample)
    samples = list(set(samples))

os.mkdir("/home/rzli/test/")
os.mkdir("/home/rzli/test/pbs")

for i in samples:
    with open("/home/rzli/test/pbs/" + i + ".sh", "w") as w:
        line = "#!/bin/bash" + "\n"
        line = "samtools view -s 0.5 -bo "
        line = line + "/home/rzli/swxxxjz/bam/10/" + i
        line = line + " /home/rzli/test/0.5_" + i
        w.write(line)

for i in samples:
    qsub("/home/rzli/test/pbs/" + i + ".sh")


import pefile as pe
import os
import numpy as np
import pandas as pd

benign_l = os.listdir("malware_dataset/benignware") #label : 0
malware_l = os.listdir("malware_dataset/malware") #label : 1
test_l = os.listdir("malware_dataset/testdata")
benign_path = "malware_dataset/benignware/"
malware_path = "malware_dataset/malware/"
test_path = "malware_dataset/testdata/"

no_of_sections = []
no_of_symbols = []
timestamp = []
characteristics = []
addr_of_entry_pt = []
size_of_init_data = []
size_of_uninit_data = []
checksum = []
dll_char = []
has_relocs_ = []
verify_checksum_ = []
is_dll_ = []
magic_value = []
directory_entry_imports_len = []
imports_len = []
imports_len = []
label = []

def create_df(directory,path):

    print(directory)
    count = 0
    for i in directory:
        print(count)
        count = count + 1

        try:
            filename = path+str(i)
#            print(filename)
            p = pe.PE(name=filename)
            print("Hash Value : ",p.get_imphash())

            print("e_magic_value : ", hex(p.DOS_HEADER.e_magic))
            magic = p.DOS_HEADER.e_magic
            magic_value.append(magic)

            print("signature : ", hex(p.NT_HEADERS.Signature))

            for section in p.sections:
            #    print(section.Name.decode('utf-8'))
#                print("\t Virtual Address :", section.VirtualAddress)
#               print("\t Misc_Physical Address : ", section.Misc_PhysicalAddress)
#               print("\t Pointer to raw data : ", section.PointerToRawData)
#               print("\t Pointer to relocations : ", section.PointerToRelocations)
#               print("\t Pointer to line numbers : ",section.PointerToLinenumbers)
#               print("\t Number of Relocations", section.NumberOfLinenumbers)
                pass

            has_relocs_.append(p.has_relocs())
            verify_checksum_.append(p.verify_checksum())
            is_dll_.append(p.is_dll())

            no_of_sections.append(p.FILE_HEADER.NumberOfSections)

            no_of_symbols.append(p.FILE_HEADER.NumberOfSymbols)

            timestamp.append(p.FILE_HEADER.TimeDateStamp)

            characteristics.append(p.FILE_HEADER.Characteristics)

            addr_of_entry_pt.append(p.OPTIONAL_HEADER.AddressOfEntryPoint)

            size_of_init_data.append(p.OPTIONAL_HEADER.SizeOfInitializedData)

            size_of_uninit_data.append(p.OPTIONAL_HEADER.SizeOfUninitializedData)

            checksum.append(p.OPTIONAL_HEADER.DllCharacteristics)

            #directory_entry_imports_len.append(len(p.DIRECTORY_ENTRY_IMPORT))

            #imports = [x.imports for x in p.DIRECTORY_ENTRY_IMPORT]
            #imports_len.append(len(imports))

#            label.append(1)
#            label.append(0)
        except OSError as e:
            pass
        except pe.PEFormatError as e:
            pass

#create_df(malware_l,malware_path)
#create_df(benign_l,benign_path)
create_df(test_l,test_path)

df = pd.DataFrame()
df['Magic Number'] = magic_value
df['Number of Sections'] = no_of_sections
df['number of symbols'] = no_of_symbols
df['has relocs'] = has_relocs_
df['is dll'] = is_dll_
df['verify checksum'] = verify_checksum_
df['time stamp'] = timestamp
df['characteristics'] = characteristics
df['address of entry point'] = addr_of_entry_pt
df['size of init data'] = size_of_init_data
df['size of uninit data'] = size_of_uninit_data
df['checksum'] = checksum
#df['label'] = label
df.to_csv(r"malware_dataset/benignware_output/data_frame_benign.csv")
#df.to_csv(r"malware_dataset/benignware_output/data_frame_mal.csv")
df.to_csv(r"malware_dataset/data_frame_test.csv")