"""
# Analysis utilities

This script belongs to the following manuscript:

- Mathôt, Berberyan, Büchel, Ruuskanen, Vilotjević, & Kruijne (in prep.)
  *Causal effects of pupil size on visual ERPs*

This module contains various constants and functions that are used in the main
analysis scripts.
"""
import random
import sys
import multiprocessing as mp
import mne; mne.set_log_level(False)
import eeg_eyetracking_parser as eet
from eeg_eyetracking_parser import braindecode_utils as bdu, \
    _eeg_preprocessing as epp
import numpy as np
import time_series_test as tst
from datamatrix import DataMatrix, convert as cnv, operations as ops, \
    functional as fnc, SeriesColumn, io, MultiDimensionalColumn
from mne.time_frequency import tfr_morlet
import matplotlib as mpl
from matplotlib import pyplot as plt
from scipy.stats import mode
import logging; logging.basicConfig(level=logging.INFO, force=True)

FIXATION_TRIGGER = 1
CUE_TRIGGER = 2
INTERVAL_TRIGGER = 3
TARGET_TRIGGER = 4
RESPONSE_TRIGGER = 5
N_CHANNELS = 26
# Occipital
LEFT_OCCIPITAL = 'O1',
RIGHT_OCCIPITAL = 'O2',
MIDLINE_OCCIPITAL = 'Oz',
# Parietal
LEFT_PARIETAL = 'P3', 'P7', 'CP1'
RIGHT_PARIETAL = 'P4', 'P8', 'CP2'
MIDLINE_PARIETAL = 'Pz', 'POz'
# Central
LEFT_CENTRAL = 'T7', 'C3'
RIGHT_CENTRAL = 'T8', 'C4'
MIDLINE_CENTRAL = 'Cz',
# Frontal
LEFT_FRONTAL = 'FC1', 'F3', 'F7', 'FP1'
RIGHT_FRONTAL = 'FC2', 'F4', 'F8', 'FP2'
MIDLINE_FRONTAL = 'Fz', 'FPz'
# Only CP1 and CP2, which were the best channels
LEFT_CP = 'CP1',
RIGHT_CP = 'CP2',
MIDLINE_CP = tuple()
# Pz, POz, Oz
LEFT_OPM = 'Pz',
RIGHT_OPM = 'POz',
MIDLINE_OPM = 'Oz',
# Select a channel group for further processing. The main analyses focus on the
# the parietal group.
CHANNEL_GROUPS = 'parietal', 'occipital', 'frontal', 'central', 'CP', \
    'occipital-parietal-midline'
# Allow the channel group to be specified on the command line
for arg in sys.argv:
    if arg in CHANNEL_GROUPS:
        CHANNEL_GROUP = arg
        break
else:
    CHANNEL_GROUP = 'parietal'
if CHANNEL_GROUP == 'parietal':
    LEFT_CHANNELS = LEFT_PARIETAL
    RIGHT_CHANNELS = RIGHT_PARIETAL
    MIDLINE_CHANNELS = MIDLINE_PARIETAL
elif CHANNEL_GROUP == 'occipital':
    LEFT_CHANNELS = LEFT_OCCIPITAL
    RIGHT_CHANNELS = RIGHT_OCCIPITAL
    MIDLINE_CHANNELS = MIDLINE_OCCIPITAL
elif CHANNEL_GROUP == 'frontal':
    LEFT_CHANNELS = LEFT_FRONTAL
    RIGHT_CHANNELS = RIGHT_FRONTAL
    MIDLINE_CHANNELS = MIDLINE_FRONTAL
elif CHANNEL_GROUP == 'central':
    LEFT_CHANNELS = LEFT_CENTRAL
    RIGHT_CHANNELS = RIGHT_CENTRAL
    MIDLINE_CHANNELS = MIDLINE_CENTRAL
elif CHANNEL_GROUP == 'CP':
    LEFT_CHANNELS = LEFT_CP
    RIGHT_CHANNELS = RIGHT_CP
    MIDLINE_CHANNELS = MIDLINE_CP
elif CHANNEL_GROUP == 'occipital-parietal-midline':
    LEFT_CHANNELS = LEFT_OPM
    RIGHT_CHANNELS = RIGHT_OPM
    MIDLINE_CHANNELS = MIDLINE_OPM
else:
    raise ValueError(f'Invalid channel group: {CHANNEL_GROUP}')
ALL_CHANNELS = LEFT_CHANNELS + RIGHT_CHANNELS + MIDLINE_CHANNELS
FACTORS = ['inducer', 'bin_pupil', 'intensity', 'valid']
LABELS = ['00:blue:0:100:no',
          '01:blue:0:100:yes',
          '02:blue:0:255:no',
          '03:blue:0:255:yes',
          '04:blue:1:100:no',
          '05:blue:1:100:yes',
          '06:blue:1:255:no',
          '07:blue:1:255:yes',
          '08:red:0:100:no',
          '09:red:0:100:yes',
          '10:red:0:255:no',
          '11:red:0:255:yes',
          '12:red:1:100:no',
          '13:red:1:100:yes',
          '14:red:1:255:no',
          '15:red:1:255:yes']
ALPHA = .05
N_CONDITIONS = 16  # 4 factors with 2 levels each
FULL_FREQS = np.arange(4, 30, 1)
NOTCH_FREQS = np.exp(np.linspace(np.log(4), np.log(30), 15))
DELTA_FREQS = np.arange(.5, 4, .5)
THETA_FREQS = np.arange(4, 8, .5)
ALPHA_FREQS = np.arange(8, 12.5, .5)
BETA_FREQS = np.arange(13, 30, .5)
PERTURB_TIMES = [(-.1, .47),
                 (.18, .74)]
SUBJECTS = list(range(1, 34))
SUBJECTS.remove(7)   # technical error
SUBJECTS.remove(5)   # negative inducer effect
SUBJECTS.remove(18)  # negative inducer effect
DATA_FOLDER = 'data'
EPOCHS_KWARGS = dict(tmin=-.1, tmax=.75, picks='eeg',
                     preload=True, reject_by_annotation=False,
                     baseline=None)
# Plotting colors
RED = 'red'
BLUE = 'blue'
FACTOR_COLORS = {
    'inducer': '#B71C1C',
    'bin_pupil': '#4A148C',
    'intensity': '#263238',
    'valid': '#1B5E20'
}
# TFR plotting parameters
Y_FREQS = np.array([0, 4, 9, 25])
VMIN = -.2
VMAX = .2
CMAP = 'coolwarm'
# Plotting style
plt.style.use('default')
mpl.rcParams['font.family'] = 'Roboto Condensed'
# DATA_CHECKPOINT = 'checkpoints/18012023.dm'
DATA_CHECKPOINT = f'checkpoints/19072023-{CHANNEL_GROUP}.dm'

def read_subject(subject_nr):
    """A simple wrapper function that calls eet.read_subject() with the correct
    parameters.

    Parameters
    ----------
    subject_nr: int

    Returns
    -------
    tuple
        A (raw, events, metadata) tuple
    """
    return eet.read_subject(subject_nr=subject_nr,
                            saccade_annotation='BADS_SACCADE',
                            min_sacc_size=128)

def get_tgt_epoch(raw, events, metadata, channels=None, tmin=-.1, tmax=.5,
                  baseline=(None, 0)):
    """A simple wrapper function that uses eet.autoreject_epochs() to get
    an Epochs object around the target onset.

    Parameters
    ----------
    raw: Raw
    events: tuple
    metadata: DataFrame
    channels: list or None, optional
        A list of channel indices or None to select all channels
    tmin: float, optional
    tmax: float, optional
    baseline: tuple, optional

    Returns
    -------
    Epochs
    """
    return eet.autoreject_epochs(
        raw, eet.epoch_trigger(events, TARGET_TRIGGER), tmin=tmin, tmax=tmax,
        metadata=metadata, picks=channels, baseline=baseline,
        ar_kwargs=dict(n_jobs=8))

def get_fix_epoch(raw, events, metadata, channels=None):
    """A simple wrapper function that uses eet.autoreject_epochs() to get
    an Epochs object around the fixation onset.

    Parameters
    ----------
    raw: Raw
    events: tuple
    metadata: DataFrame
    channels: list or None, optional
        A list of channel indices or None to select all channels

    Returns
    -------
    Epochs
    """
    return eet.autoreject_epochs(
        raw, eet.epoch_trigger(events, FIXATION_TRIGGER), tmin=-.5, tmax=2.5,
        metadata=metadata, picks=channels, ar_kwargs=dict(n_jobs=8))

def get_morlet(epochs, freqs, crop=(0, 2), decim=8, n_cycles=2):
    """A simple wrapper function that uses tfr_morlet() to extract
    time-frequency data.

    Parameters
    ----------
    epochs: Epochs
    freqs: array
        An array of frequencies
    crop: tuple, optional
        A time window to crop after extracting the time-frequency data to
        reduce edge artifacts.
    decim: int, optional
        Downsampling factor to reduce memory consumption
    n_cycles: int, optional
        The number of cycles of the morlet wavelet

    Returns
    -------
    EpochsTFR
    """
    morlet = tfr_morlet(epochs, freqs=freqs, n_cycles=n_cycles, n_jobs=-1,
                        return_itc=False, use_fft=True, average=False,
                        decim=decim,
                        picks=np.arange(len(epochs.info['ch_names'])))
    morlet.crop(*crop)
    return morlet

def z_by_freq(col):
    """Performs z-scoring across trials, channels, and time points but
    separately for each frequency.

    Parameters
    ----------
    col: MultiDimensionalColumn

    Returns
    -------
    MultiDimensionalColumn
    """
    zcol = col[:]
    for i in range(zcol.shape[2]):
        zcol._seq[:, :, i] = (
            (zcol._seq[:, :, i] - np.nanmean(zcol._seq[:, :, i]))
            / np.nanstd(zcol._seq[:, :, i])
        )
    return zcol

def subject_data(subject_nr):
    """Performs preprocessing for a single participant. This involves basic
    EEG preprocessing and subsequent epoching and extraction of TFR data. The
    result is a single DataMatrix that contains all information for final
    analysis.

    Parameters
    ----------
    subject_nr: int

    Returns
    -------
    DataMatrix
    """
    print(f'Processing subject {subject_nr}')
    raw, events, metadata = read_subject(subject_nr)
    raw['PupilSize'] = area_to_mm(raw['PupilSize'][0])
    dm = cnv.from_pandas(metadata)
    print('- eeg')
    tgt_epoch = get_tgt_epoch(raw, events, metadata)
    dm.tgt_erp = cnv.from_mne_epochs(tgt_epoch)
    tgt_tfr = get_morlet(
        get_tgt_epoch(raw, events, metadata, baseline=None, tmax=1),
        FULL_FREQS, crop=(0, .5), decim=4)
    dm.tgt_tfr = cnv.from_mne_tfr(tgt_tfr)
    dm.tgt_tfr = z_by_freq(dm.tgt_tfr)
    fix_epoch = get_fix_epoch(raw, events, metadata)
    fix_tfr = get_morlet(fix_epoch, FULL_FREQS)
    dm.fix_erp = cnv.from_mne_epochs(fix_epoch)
    dm.fix_tfr = cnv.from_mne_tfr(fix_tfr)
    dm.fix_tfr = z_by_freq(dm.fix_tfr)
    print('- pupils')
    pupil_fix = eet.PupilEpochs(
        raw, eet.epoch_trigger(events, FIXATION_TRIGGER), tmin=0, tmax=2,
        metadata=metadata, baseline=None)
    pupil_target = eet.PupilEpochs(
        raw, eet.epoch_trigger(events, TARGET_TRIGGER), tmin=-.05, tmax=2,
        metadata=metadata)
    del raw
    dm.pupil_fix = cnv.from_mne_epochs(pupil_fix, ch_avg=True)
    dm.pupil_target = cnv.from_mne_epochs(pupil_target, ch_avg=True)
    return dm

@fnc.memoize(persistent=True, key='merged-data')
def get_merged_data():
    """Merges data for all participants into a single DataMatrix. Uses
    multiprocessing for performance.

    Returns
    -------
    DataMatrix
    """
    return fnc.stack_multiprocess(subject_data, SUBJECTS, processes=10)

def add_bin_pupil(raw, events, metadata):
    """Adds bin pupil to the metadata. This is a patch to allow decoding to
    take bin pupil as a decoding factor into account.

    Parameters
    ----------
    raw: Raw
    events: tuple
    metadata: DataFrame

    Returns
    -------
    tuple
        A (raw, events, metadata) tuple where bin_pupil has been added as a
        column to metadata.
    """
    # This adds the bin_pupil pseudo-factor to the data. This requires that
    # this has been generated already by `analyze.py`.
    dm = io.readtxt('output/bin-pupil.csv')
    dm = dm.subject_nr == metadata.subject_nr[0]
    metadata.loc[16:, 'bin_pupil'] = dm.bin_pupil
    dummy_factor = 192 * [0] + 192 * [1]
    random.shuffle(dummy_factor)
    metadata.loc[16:, 'dummy_factor'] = dummy_factor
    return raw, events, metadata

def decode_subject(subject_nr):
    """A wrapper function around bdu.decode_subject() that performs overall
    decoding for one subject.

    Parameters
    ----------
    subject_nr: int

    Returns
    -------
    DataMatrix
        See bdu.decode_subject()
    """
    read_subject_kwargs = dict(subject_nr=subject_nr,
                               saccade_annotation='BADS_SACCADE',
                               min_sacc_size=128)
    return bdu.decode_subject(
        read_subject_kwargs=read_subject_kwargs, factors=FACTORS,
        epochs_kwargs=EPOCHS_KWARGS, trigger=TARGET_TRIGGER, window_stride=1,
        window_size=200, n_fold=4, epochs=4, patch_data_func=add_bin_pupil)

def crossdecode_subject(subject_nr, from_factor, to_factor):
    """A wrapper function around bdu.decode_subject() that performs
    cross-decoding for one subject.

    Parameters
    ----------
    subject_nr: int
    from_factor: str
        The factor to train on
    to_factor: str
        The factor to test on

    Returns
    -------
    DataMatrix
        See bdu.decode_subject()
    """
    read_subject_kwargs = dict(subject_nr=subject_nr,
                               saccade_annotation='BADS_SACCADE',
                               min_sacc_size=128)
    if 'bin_pupil' in (from_factor, to_factor):
        return bdu.decode_subject(
            read_subject_kwargs=read_subject_kwargs, factors=from_factor,
            crossdecode_factors=to_factor, epochs_kwargs=EPOCHS_KWARGS,
            trigger=TARGET_TRIGGER, window_stride=1, window_size=200, n_fold=4,
            epochs=4, patch_data_func=add_bin_pupil)
    return bdu.decode_subject(
        read_subject_kwargs=read_subject_kwargs, factors=from_factor,
        crossdecode_factors=to_factor, epochs_kwargs=EPOCHS_KWARGS,
        trigger=TARGET_TRIGGER, window_stride=1, window_size=200, n_fold=4,
        epochs=4)

@fnc.memoize(persistent=True)
def blocked_decode_subject(subject_nr, factor, query1, query2):
    """Decodes a factor for a single subject, using two different queries to
    separate the training and testing data.

    Parameters
    ----------
    subject_nr: int
    factor: str
    query1: str
        A pandas-style query to select the training set
    query2: str
        A pandas-style query to select the testing set

    Returns
    -------
    float
        Decoding accuracy
    """
    read_subject_kwargs = dict(subject_nr=subject_nr,
                               saccade_annotation='BADS_SACCADE',
                               min_sacc_size=128)
    train_data, train_labels, train_metadata = bdu.read_decode_dataset(
        read_subject_kwargs, factor, EPOCHS_KWARGS, TARGET_TRIGGER, query1)
    test_data, test_labels, test_metadata = bdu.read_decode_dataset(
        read_subject_kwargs, factor, EPOCHS_KWARGS, TARGET_TRIGGER, query2)
    clf = bdu.train(train_data, test_data)
    y_pred = clf.predict(test_data)
    resized_pred = y_pred.copy()
    resized_pred.resize(
            (len(test_data.datasets), len(test_data.datasets[0])))
    resized_pred = mode(resized_pred, axis=1)[0].flatten()
    y_true = [d.y[0] for d in test_data.datasets]
    return np.mean([p == t for p, t in zip(resized_pred, y_true)])

def statsplot(rm):
    """A simple wrapper function that plots statistical values as a function of
    time and factors.

    Parameters
    ----------
    rm: DataMatrix
        A DataMatrix with statistical results as returned by time_series_test
        functions.
    """
    rm = rm[:]
    rm.sign = SeriesColumn(depth=rm.p.depth)
    colors = ['red', 'green', 'blue', 'orange']
    for y, row in enumerate(rm[1:]):
        for linewidth, alpha in [(1, .05), (2, .01), (4, .005), (8, .001)]:
            row.sign[row.p >= alpha] = np.nan
            row.sign[row.p < alpha] = y
            plt.plot(row.sign, '-', label=f'{row.effect}, p < {alpha}',
                     linewidth=linewidth, color=colors[y])
    plt.legend(bbox_to_anchor=(1, 1), loc="upper left")

def select_ica(raw, events, metadata, exclude_component=0):
    """A helper function that excludes (rather than selects, as the name
    suggests) an independent component from the signal.

    Parameters
    ----------
    raw: Raw
    events: tuple
    metadata: DataFrame
    exclude_component: int, optional
        The index of the excluded independent component

    Returns
    -------
    tuple
        A (raw, events, metadata) tuple where the independent component has
        been excluded from the `raw` object.
    """
    global weights_dict
    raw, events, metadata = add_bin_pupil(raw, events, metadata)
    print(f'running ica to exclude component {exclude_component}')

    @fnc.memoize(persistent=True)
    def run_ica(raw):
        return epp.run_ica(raw)

    # run_ica.clear()
    raw.info['bads'] = []
    ica = run_ica(raw)
    print('applying ica')
    ica.apply(raw, exclude=[exclude_component])
    weights = np.dot(ica.mixing_matrix_[:, exclude_component].T,
                     ica.pca_components_[:ica.n_components_])
    weights_dict = {ch_name: weight
                    for ch_name, weight in zip(ica.ch_names, weights)}
    print(f'weights: {weights_dict} (len={len(weights_dict)})')
    return raw, events, metadata

@fnc.memoize(persistent=True)
def ica_perturbation_decode(subject_nr, factor):
    """Performs the ICA perturbation analysis.

    Parameters
    ----------
    subject_nr: int
    factor: str

    Returns
    -------
    tuple
        The first element of the tuple is a DataMatrix with the regular
        decoding results. The second element is a dict with independent
        component indices as keys and (dm, weights_dict) tuples as values.
        Here, dm is the DataMatrix with the decoding results after excluding
        the independent component, and weights_dict is a mapping with channel
        names as keys and weights (loading of the channel of the independent
        component) as values.
    """
    read_subject_kwargs = dict(subject_nr=subject_nr,
                               saccade_annotation='BADS_SACCADE',
                               min_sacc_size=128)
    fdm = bdu.decode_subject(
        read_subject_kwargs=read_subject_kwargs, factors=factor,
        epochs_kwargs=EPOCHS_KWARGS, trigger=TARGET_TRIGGER, window_stride=1,
        window_size=200, n_fold=4, epochs=4, patch_data_func=add_bin_pupil)
    print(f'full-data accuracy: {fdm.braindecode_correct.mean}')
    perturbation_results = {}
    for exclude_component in range(N_CHANNELS):
        bdu.decode_subject.clear()
        dm = bdu.decode_subject(
            read_subject_kwargs=read_subject_kwargs, factors=factor,
            epochs_kwargs=EPOCHS_KWARGS, trigger=TARGET_TRIGGER,
            window_stride=1, window_size=200, n_fold=4, epochs=4,
            patch_data_func=lambda raw, events, metadata: select_ica(
                    raw, events, metadata, exclude_component))
        perturbation_results[exclude_component] = dm, weights_dict
        print(f'perturbation accuracy({exclude_component}): '
              f'{dm.braindecode_correct.mean}')
    return fdm, perturbation_results

def notch_filter(raw, events, metadata, freq):
    """A helper function that excludes a frequency from the signal using a
    notch filter.

    Parameters
    ----------
    raw: Raw
    events: tuple
    metadata: DataFrame
    freq: float
        The frequency to remove.

    Returns
    -------
    tuple
        A (raw, events, metadata) tuple where the frequency has been removed
        from the `raw` object using a notch filter.
    """
    global weights_dict
    raw, events, metadata = add_bin_pupil(raw, events, metadata)
    width = np.exp(np.log(freq / 4))
    print(f'notch-filtering frequency band: {freq:.2f} / {width:.2f}')
    raw.notch_filter(freq, notch_widths=width, trans_bandwidth=width)
    return raw, events, metadata

@fnc.memoize(persistent=True)
def freq_perturbation_decode(subject_nr, factor):
    """Performs the frequency perturbation analysis.

    Parameters
    ----------
    subject_nr: int
    factor: str

    Returns
    -------
    tuple
        The first element of the tuple is a DataMatrix with the regular
        decoding results. The second element is a dict with frequencies
        as keys and the DataMatrix objects with the decoding results after
        excluding the frequencies as values.
    """
    read_subject_kwargs = dict(subject_nr=subject_nr,
                               saccade_annotation='BADS_SACCADE',
                               min_sacc_size=128)
    fdm = bdu.decode_subject(
        read_subject_kwargs=read_subject_kwargs, factors=factor,
        epochs_kwargs=EPOCHS_KWARGS, trigger=TARGET_TRIGGER, window_stride=1,
        window_size=200, n_fold=4, epochs=4, patch_data_func=add_bin_pupil)
    print(f'full-data accuracy: {fdm.braindecode_correct.mean}')
    perturbation_results = {}
    for freq in NOTCH_FREQS:
        bdu.decode_subject.clear()
        dm = bdu.decode_subject(
            read_subject_kwargs=read_subject_kwargs, factors=factor,
            epochs_kwargs=EPOCHS_KWARGS, trigger=TARGET_TRIGGER,
            window_stride=1, window_size=200, n_fold=4, epochs=4,
            patch_data_func=lambda raw, events, metadata: notch_filter(
                    raw, events, metadata, freq))
        perturbation_results[freq] = dm
        print(f'perturbation accuracy({freq}): {dm.braindecode_correct.mean}')
    return fdm, perturbation_results

def area_to_mm(au):
    """Converts in arbitrary units to millimeters of diameter. This is specific
    to the recording set-up.

    Parameters
    ----------
    au: float

    Returns
    -------
    float
    """
    return -0.9904 + 0.1275 * au ** .5

def pupil_plot(dm, dv='pupil_target', **kwargs):
    """A simple wrapper function that plots pupil size over time.

    Parameters
    ----------
    dm: DataMatrix
    dv: str, optional
    **kwargs: dict, optional
    """
    tst.plot(dm, dv=dv, legend_kwargs={'loc': 'lower left'},
             **kwargs)
    x = np.linspace(12, 262, 6)
    t = [f'{int(s)}' for s in np.linspace(0, 1000, 6)]
    plt.xticks(x, t)
    plt.xlabel('Time (ms)')
    if dv == 'pupil_target':
        plt.axhline(0, linestyle=':', color='black')
        plt.ylim(-.6, .2)
    else:
        plt.ylim(2, 8)
    plt.xlim(0, 250)
    plt.ylabel('Baseline-corrected pupil size (mm)')

def erp_plot(dm, dv='lat_erp', ylim=None, **kwargs):
    """A simple wrapper function that plots ERPs.

    Parameters
    ----------
    dm: DataMatrix
    dv: str, optional
    ylim: float or None, optional
    **kwargs: dict, optional
    """
    tst.plot(dm, dv=dv, **kwargs)
    plt.xticks(np.arange(25, 150, 25), np.arange(0, 500, 100))
    plt.axvline(25, color='black', linestyle=':')
    plt.axhline(0, color='black', linestyle=':')
    plt.xlabel('Time (ms)')
    if ylim:
        plt.ylim(*ylim)

def tfr_plot(dm, dv):
    """A simple wrapper function that creates a multipanel TFR plot.

    Parameters
    ----------
    dm: DataMatrix
    dv: str, optional
    """
    plt.figure(figsize=(12, 4))
    plt.subplots_adjust(wspace=0)
    plt.subplot(141)
    tfr_red = (dm.inducer == 'red')[dv][...]
    tfr_blue = (dm.inducer == 'blue')[dv][...]
    plt.title('a) Induced Pupil Size (Large - Small)')
    plt.imshow(tfr_red - tfr_blue, aspect='auto', vmin=VMIN, vmax=VMAX,
               cmap=CMAP, interpolation='bicubic')
    plt.yticks(Y_FREQS, FULL_FREQS[Y_FREQS])
    plt.xticks(np.arange(0, 31, 6.25), np.arange(0, 499, 100))
    plt.xlabel('Time (ms)')
    plt.ylabel('Frequency (Hz)')

    plt.subplot(142)
    tfr_large = (dm.bin_pupil == 1)[dv][...]
    tfr_small = (dm.bin_pupil == 0)[dv][...]
    plt.title('b) Spontaneous Pupil Size (Large - Small)')
    plt.imshow(tfr_large - tfr_small, aspect='auto', vmin=VMIN, vmax=VMAX,
               cmap=CMAP, interpolation='bicubic')
    plt.gca().get_yaxis().set_visible(False)
    plt.xticks(np.arange(0, 31, 6.25), np.arange(0, 499, 100))
    plt.xlabel('Time (ms)')

    plt.subplot(143)
    tfr_bright = (dm.intensity == 255)[dv].mean
    tfr_dim = (dm.intensity == 100)[dv].mean
    plt.title('c) Stimulus Intensity (Bright - Dim)')
    plt.imshow(tfr_bright - tfr_dim, aspect='auto', vmin=VMIN, vmax=VMAX,
               cmap=CMAP, interpolation='bicubic')
    plt.gca().get_yaxis().set_visible(False)
    plt.xticks(np.arange(0, 31, 6.25), np.arange(0, 499, 100))
    plt.xlabel('Time (ms)')

    plt.subplot(144)
    tfr_attended = (dm.valid == 'yes')[dv].mean
    tfr_unattended = (dm.valid == 'no')[dv].mean
    plt.title('d) Covert Visual Attention (Attended - Unattended)')
    plt.imshow(tfr_attended - tfr_unattended, aspect='auto', vmin=VMIN,
               vmax=VMAX, cmap=CMAP, interpolation='bicubic')
    plt.gca().get_yaxis().set_visible(False)
    plt.xticks(np.arange(0, 31, 6.25), np.arange(0, 499, 100))
    plt.xlabel('Time (ms)')


# -*- coding: utf-8 -*-
# @Time    : 19-1-24 下午9:35
# @Author  : ccs
import json

from django.http import HttpResponse

def calc(request):
    a = request.GET['a']
    b = request.GET['b']
    c = request.GET['c']
    print(a,b,c)
    m = a+b+c
    n = b+a
    rets = {"m":m,"n":n}
    retsj = json.dumps(rets)
    return HttpResponse(retsj)


# -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function, unicode_literals

import logging

from temba_client.v2 import TembaClient

from django.conf import settings
from django.db import migrations

from ureport.utils import datetime_to_json_date, json_date_to_datetime

logger = logging.getLogger(__name__)

class Migration(migrations.Migration):
    def populate_poll_poll_date(apps, schema_editor):
        Poll = apps.get_model("polls", "Poll")
        Org = apps.get_model("orgs", "Org")

        agent = getattr(settings, "SITE_API_USER_AGENT", None)
        host = settings.SITE_API_HOST

        for org in Org.objects.all():
            temba_client = TembaClient(host, org.api_token, user_agent=agent)
            api_flows = temba_client.get_flows()
            flows_date = dict()
            for flow in api_flows:
                flows_date[flow.uuid] = datetime_to_json_date(flow.created_on)

            for poll in Poll.objects.filter(org=org):
                json_date = flows_date.get(poll.flow_uuid, None)
                if json_date:
                    date = json_date_to_datetime(json_date)
                else:
                    logger.info("using created_on for flow_date on poll with id %s" % poll.pk)
                    date = poll.created_on

                poll.poll_date = date
                poll.save()

    dependencies = [("polls", "0022_poll_flow_date")]

    operations = [migrations.RunPython(populate_poll_poll_date)]


from firebase_admin import firestore
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics.pairwise import cosine_similarity

def get_user_skills(userid):
    item = firestore.client().collection('user').document(userid).get().to_dict()['skills']
    user_skill_string = ' '.join(str(e) for e in item)
    return user_skill_string

def get_result(field, user_skill_string, keyword_string):
    content = (user_skill_string, keyword_string)
    cv = CountVectorizer()
    matrix = cv.fit_transform(content)
    singularity_matrix = cosine_similarity(matrix)
    result = (singularity_matrix[1][0]*100)
    result = str(result)
    result = result.split('.', 1)[0]
    # return 'You scored {}% in {} jobs.'.format(result, field)
    return result

def check_result(field, result):
    return 'You scored {}% in {} jobs.'.format(result, field)


import sys
from PySide6.QtWidgets import QApplication, QPushButton
from PySide6.QtCore import Slot

# 这个例子包含Signals and Slots(信号与槽机制)
# 使用@Slot()表明这是一个槽函数
# @Slot() 服了，没使用这个居然也能照常运行
def say_hello():
    print("Button, clicked, hello!")

app = QApplication([])

# QPushButton里面的参数是按钮上会显示的文字
button = QPushButton("点我！")
button.clicked.connect(say_hello)
button.show()

app.exec()


import sys
import form
from PyQt4 import QtCore, QtGui
import letters
import pygame

class ConnectorToMainWindow(QtGui.QMainWindow):
    def __init__(self, parent = None):
        super(ConnectorToMainWindow, self).__init__()
        self.expected_letter = ''
        self.timer = QtCore.QTimer()
        self.ui = form.Ui_MainWindow()
        self.ui.setupUi(self)
        self.my_letters = letters.Letters()
        QtCore.QObject.connect(self.ui.restart_button, QtCore.SIGNAL(form._fromUtf8("clicked()")),
                               self.on_restartbutton_click)
        QtCore.QObject.connect(self.ui.audio_button, QtCore.SIGNAL(form._fromUtf8("clicked()")),
                               self.on_audiobutton_click)
        QtCore.QMetaObject.connectSlotsByName(self)
        self.on_restartbutton_click()

    def on_restartbutton_click(self) :
        self.my_letters.re_init()
        self.ui.answer_window.clear()
        self.ui.img_display.clear()
        self.generate_a_letter()

    def on_audiobutton_click(self) :
        self.play_audio("./audio/letters/{}.mp3".format(self.expected_letter))

    def generate_a_letter(self) :
        self.ui.img_display.clear()
        letter = self.my_letters.pick_a_letter()
        self.play_audio("./audio/letters/{}.mp3".format(letter))
        self.expected_letter = letter
        self.ui.question_window.setText('{0}  {1}'.format(letter.upper(), letter))
        self.ui.answer_window.setFocus()

    def verify_result(self) :
        if str(self.ui.answer_window.toPlainText()).lower() == self.expected_letter :
            self.display_image(QtGui.QPixmap("./imgs/happy.png"))
            self.ui.answer_window.clear()
            QtCore.QTimer.singleShot(2000, self.generate_a_letter)
        else :
            self.display_image(QtGui.QPixmap("./imgs/try_again.png"))
            self.ui.answer_window.clear()
            self.retry()

    def retry(self) :
        self.ui.question_window.setText('{0}  {1}'.format(self.expected_letter.upper(),
                                                        self.expected_letter))
        self.on_audiobutton_click()
        self.ui.answer_window.setFocus()

    def display_image(self, img) :
        self.ui.img_display.setPixmap(img)
        self.ui.img_display.show()

    def play_audio(self, path) :
        pygame.mixer.init(frequency = 8000, channels = 1)
        pygame.mixer.music.load(path)
        pygame.mixer.music.play()

    def keyPressEvent(self, e):
        self.ui.img_display.clear()
        if e.key() != QtCore.Qt.Key_Return and e.key() != QtCore.Qt.Key_Enter:
            if e.key() != QtCore.Qt.Key_Backspace:
                self.ui.answer_window.insertPlainText(str(e.text()).upper())
            else:
                mstr = str(self.ui.answer_window.toPlainText())[:-1]
                self.ui.answer_window.clear()
                self.ui.answer_window.insertPlainText(mstr)
        else:
            self.verify_result()

if __name__ == "__main__":
    app = QtGui.QApplication(sys.argv)
    myapp = ConnectorToMainWindow()
    myapp.show()
    sys.exit(app.exec_())


import gi
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk

ICONSIZE = Gtk.IconSize.SMALL_TOOLBAR

class controlBar(Gtk.HeaderBar):
    def __init__(self):
        Gtk.HeaderBar.__init__(self)

        self.set_show_close_button(True)
        self.props.title = "PyFlowChart"

        self.info_box = Gtk.Box(spacing=10,orientation=Gtk.Orientation.HORIZONTAL)

        self.file_button = Gtk.Button.new_from_icon_name("open-menu-symbolic", ICONSIZE)

        self.settings_button = Gtk.Button.new_from_icon_name("preferences-system", ICONSIZE)
        self.help_button = Gtk.Button.new_from_icon_name("help-about", ICONSIZE)

        self.main_menu = self.init_main_menu()
        self.settings_menu = self.init_settings_menu()
        self.help_menu = self.init_help_menu()
        #self.init_edit_menu()

        self.file_button.connect('clicked', self.file_clicked)
        self.settings_button.connect('clicked', self.settings_clicked)
        self.help_button.connect('clicked', self.help_clicked)

        self.info_box.add(self.settings_button)
        self.info_box.add(self.help_button)

        self.pack_start(self.file_button)
        self.pack_end(self.info_box)

        self.buttons = []
        # self.populate_buttons()

        self.show_all()

    def init_main_menu(self):
        main_menu = Gtk.Menu()

        self.new_button = Gtk.MenuItem.new_with_label('New')
        self.open_button = Gtk.MenuItem.new_with_label('Open')

        self.open_stock_button = Gtk.MenuItem.new_with_label('Import Stock')

        view_button = Gtk.MenuItem.new_with_label('View')

        self.view_menu = Gtk.Menu()
        self.viewer_button = Gtk.MenuItem.new_with_label('Viewer')
        self.builder_button = Gtk.MenuItem.new_with_label('Builder')
        self.view_menu.append(self.viewer_button)
        self.view_menu.append(self.builder_button)
        self.view_menu.show_all()

        view_button.set_submenu(self.view_menu)

        self.save_button = Gtk.MenuItem.new_with_label('Save')
        self.save_as_button = Gtk.MenuItem.new_with_label('Save As...')
        self.quit_button = Gtk.MenuItem.new_with_label('Quit')

        main_menu.append(self.new_button)
        main_menu.append(self.open_button)
        main_menu.append(self.open_stock_button)
        main_menu.append(Gtk.SeparatorMenuItem())
        main_menu.append(view_button)
        main_menu.append(Gtk.SeparatorMenuItem())
        main_menu.append(self.save_button)
        main_menu.append(self.save_as_button)
        main_menu.append(Gtk.SeparatorMenuItem())
        main_menu.append(self.quit_button)
        main_menu.show_all()

        return main_menu

    def init_settings_menu(self):
        settings_menu = Gtk.Menu()

        self.preferences_button = Gtk.MenuItem.new_with_label('Preferences')
        settings_menu.append(self.preferences_button)

        settings_menu.show_all()
        return settings_menu

    def init_help_menu(self):
        help_menu = Gtk.Menu()

        self.app_help_button = Gtk.MenuItem.new_with_label('Help')
        self.about_button = Gtk.MenuItem.new_with_label('About')

        help_menu.append(self.about_button)
        help_menu.append(self.app_help_button)

        help_menu.show_all()
        return help_menu

    def file_clicked(self, widget):
        self.main_menu.popup( None, None, None, None, 0, Gtk.get_current_event_time())

    def settings_clicked(self, widget):
        self.settings_menu.popup( None, None, None, None, 0, Gtk.get_current_event_time())

    def help_clicked(self, widget):
        self.help_menu.popup( None, None, None, None, 0, Gtk.get_current_event_time())


from sklearn.base import BaseEstimator, ClusterMixin, TransformerMixin
from sklearn.cluster import DBSCAN, KMeans, SpectralClustering
from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import LabelBinarizer

class ClusterTransformer(TransformerMixin, BaseEstimator):
    """Turns sklearn clustering algorithm into a transformer component that
    you can use in a pipeline.
    If the clustering method cannot be used for prediction
    (aka. it does not have a predict method), then a nearest neighbour
    vot will be used to infer cluster labels for unseen samples.

    Parameters
    ----------
    model: ClusterMixin
        Sklearn clustering model.
    n_neighbors: int
        Number of neighbours to use for inference.
    metric: str
        Metric to use for determining nearest neighbours.

    Attributes
    ----------
    labeler: LabelBinarizer
        Component that turns cluster labels into one-hot embeddings.
    neighbors: KNeighborsClassifier
        Classifier to use for out of sample prediction.
    """

    def __init__(
        self, model: ClusterMixin, n_neighbors: int = 5, metric: str = "cosine"
    ):
        self.model = model
        self.labeler = LabelBinarizer()
        self.neighbors = KNeighborsClassifier(
            n_neighbors=n_neighbors, metric=metric
        )

    def fit(self, X, y=None):
        """Fits the clustering algorithm and label binarizer.

        Parameters
        ----------
        X: ndarray of shape (n_observations, n_features)
            Observations to cluster.
        y: None
            Ignored, exists for compatiblity.

        Returns
        -------
        self
        """
        labels = self.model.fit_predict(X)
        if not hasattr(self.model, "predict"):
            self.neighbors.fit(X, labels)
        self.labeler.fit(labels)
        return self

    def transform(self, X):
        """Infers cluster labels for given data points.

        Parameters
        ----------
        X: ndarray of shape (n_observations, n_features)
            Observations to cluster.

        Returns
        -------
        ndarray of shape (n_observations, n_clusters)
            One-hot encoding of cluster labels.
        """
        if hasattr(self.model, "predict"):
            labels = self.model.predict(X)
        else:
            labels = self.neighbors.predict(X)
        return self.labeler.transform(labels)

    def get_feature_names_out(self):
        """Returns the cluster classes for each dimension.

        Returns
        -------
        ndarray of shape (n_clusters)
            Cluster names.
        """
        return self.labeler.classes_

def DBSCANTransformer(
    eps: float = 0.5, min_samples: int = 5, metric: str = "cosine"
) -> ClusterTransformer:
    """Convenience function for creating a DBSCAN transformer.

    Parameters
    ----------
    eps : float, default 0.5
        The maximum distance between two samples for one to be considered
        as in the neighborhood of the other. This is not a maximum bound
        on the distances of points within a cluster. This is the most
        important DBSCAN parameter to choose appropriately for your data set
        and distance function.

    min_samples : int, default 5
        The number of samples (or total weight) in a neighborhood for a point
        to be considered as a core point. This includes the point itself.

    metric : str, default 'cosine'
        The metric to use when calculating distance between instances in a
        feature array.

    Returns
    -------
    ClusterTransformer
        Sklearn transformer component that wraps DBSCAN.
    """
    model = DBSCAN(eps=eps, min_samples=min_samples, metric=metric)
    return ClusterTransformer(model, metric="cosine")

def KMeansTransformer(n_clusters: int) -> ClusterTransformer:
    """Convenience function for creating a KMeans transformer.

    Parameters
    ----------
    n_clusters: int
        Number of clusters.

    Returns
    -------
    ClusterTransformer
        Sklearn transformer component that wraps KMeans.
    """
    model = KMeans(n_clusters=n_clusters)
    return ClusterTransformer(model, metric="cosine")

def SpectralClusteringTransformer(n_clusters: int) -> ClusterTransformer:
    """Convenience function for creating a Spectral Clustering transformer.

    Parameters
    ----------
    n_clusters: int
        Number of clusters.

    Returns
    -------
    ClusterTransformer
        Sklearn transformer component that wraps SpectralClustering.
    """
    model = SpectralClustering(n_clusters=n_clusters)
    return ClusterTransformer(model, metric="cosine")


import requests
import re
import json
from nonebot import on_command, CommandSession
@on_command('lol新闻', aliases=('lol新闻'))
async def weather(session: CommandSession):
    url = "http://l.zhangyoubao.com/news/"
    headers = {
            'User-Agent': 'Mozilla/5.0 (Macintosh; U; PPC Mac OS X Mach-O; en-US; rv:1.6) Gecko/20040206 Firefox/0.8'
        }
    reponse = requests.get(url, headers=headers)
    reponse.encoding = "utf-8"
    html = reponse.text
    if reponse.status_code == 200:
        pattern = re.compile('
(.*)
')
        #html = reponse.text
        items = re.findall(pattern, html)
        #print(items)
        LMG = [];
        for item in items:
            Lmg = item[1].strip()
            Lmg1 = item[0]
            Lmg2 = Lmg + " " + Lmg1
            LMG.append(Lmg2)
        Lmg4 = LMG[0:5]
        Lmg3 = '\n'.join(Lmg4)
        #print(Lmg3)
        await session.send(Lmg3)


from collections import *

import numpy as np

from common.session import AdventSession

session = AdventSession(day=20, year=2017)
data = session.data.strip()
data = data.split('\n')

p1, p2 = 0, 0

class Particle:
    def __init__(self, p, v, a, _id):
        self.p = np.array(p)
        self.v = np.array(v)
        self.a = np.array(a)
        self.p_sum = np.abs(p).sum()
        self.v_sum = np.abs(v).sum()
        self.a_sum = np.abs(a).sum()
        self.id = _id

    def update(self):
        self.v += self.a
        self.p += self.v

particles = []
for i, line in enumerate(data):
    line = line.replace('<', '(').replace('>', ')')
    particles.append(eval(f'Particle({line}, _id={i})'))

sorted_particles = sorted(particles, key=lambda p: (
    p.a_sum, p.v_sum, p.p_sum))

p1 = sorted_particles[0].id

for _ in range(1000):
    positions = defaultdict(list)
    for p in particles:
        positions[tuple(p.p)].append(p)
    for parts in positions.values():
        if len(parts) > 1:
            for part in parts:
                particles.remove(part)

    for p in particles:
        p.update()

p2 = len(particles)

session.submit(p1, part=1)
session.submit(p2, part=2)


#Coded by: QyFashae

import os
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import HashingVectorizer, TfidfTransformer
from sklearn.metrics import accuracy_score, confusion_matrix
from sklearn import tree

# Paths to the directories containing the emails
s_ep = os.path.join("training_data_for_ml_emails", "training_data_ML00")
v_ep = os.path.join("training_data_for_ml_emails", "training_data_ML01")

# List of file directories and corresponding labels
l_fd = [(s_ep, 0), (v_ep, 1)]
e_cs = []
l_bs = []

# Function to read email content and labels
def cve_sh_trfm(s_ep, v_ep, l_fd, e_cs, l_bs):
    for cfs, lbs in l_fd:
        files = os.listdir(cfs)
        for file in files:
            file_path = os.path.join(cfs, file)
            try:
                with open(file_path, "r") as current_file:
                    eml_ctt = current_file.read().replace("\n", "")
                    eml_ctt = str(eml_ctt)
                    e_cs.append(eml_ctt)
                    l_bs.append(lbs)
            except:
                pass

# Splitting the dataset and training the model
def vary_trmd(cve_sh_trfm, s_ep, v_ep, l_fd, e_cs, l_bs):
    x_train, x_test, y_train, y_test = train_test_split(
        e_cs, l_bs, test_size=0.4, random_state=17
    )
    nlp_followed_by_dt = Pipeline(
        [
            ("vect", HashingVectorizer(input="eml_ctt", ngram_range=(1, 4))),
            ("tfidf", TfidfTransformer(use_idf=True)),
            ("dt", tree.DecisionTreeClassifier(class_weight="balanced")),
        ]
    )
    nlp_followed_by_dt.fit(x_train, y_train)

    # Predicting and evaluating the model
    y_test_predict = nlp_followed_by_dt.predict(x_test)
    accuracy = accuracy_score(y_test, y_test_predict)
    confusion = confusion_matrix(y_test, y_test_predict)

    # Writing accuracy and confusion matrix to files
    with open("accuracy_score.txt", "w") as acc_file:
        acc_file.write("Accuracy: " + str(accuracy))

    with open("confusion_matrix.txt", "w") as conf_file:
        conf_file.write("Confusion Matrix:\n" + str(confusion))


    # github üzerinden yapılan arama sonuçlarını consola yazdırma

from selenium import webdriver
from selenium.webdriver.common.keys import Keys
import time

driver = webdriver.Chrome()

url = "https://github.com"

driver.get(url)

searchInput = driver.find_element_by_name("q")
time.sleep(1)
print("\n" + driver.title + "\n")

searchInput.send_keys("python")
time.sleep(1)

searchInput.send_keys(Keys.ENTER)
time.sleep(2)

result = driver.find_elements_by_css_selector(".repo-list-item div div a")

for element in result:
    print(element.text)

time.sleep(2)

driver.close()


# finding the right modules/packages to use is not easy
import os
import pyodbc
from numpy import genfromtxt
import pandas as pd
import sqlalchemy as sa # use sqlalchemy for truncating etc
from sqlalchemy import Column, Integer, Float, Date, String, BigInteger
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.dialects import mssql

import urllib
import csv
import fnmatch2
from time import time

## DATABASE STUFF
# server + instance = DESKTOP-A6J6D7Q\LEROYB_INSTANCE
# database = Tracks

server_name='DESKTOP-A6J6D7Q\LEROYB_INSTANCE'
db_name='Tracks'

path = 'F:\export_6642035' # location of all the files related to STRAVA export
pathArchive = 'F:\export_6642035_archive' # this is the archive folder that will contain all the processed files
searchstring = '*' # we could limit the files to a particular string pattern, however in this case we want to go through all the files

# https://docs.sqlalchemy.org/en/13/dialects/mssql.html#connecting-to-pyodbc
# https://docs.sqlalchemy.org/en/13/core/connections.html
# https://www.pythonsheets.com/notes/python-sqlalchemy.html
# https://auth0.com/blog/sqlalchemy-orm-tutorial-for-python-developers/
# https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.truncate.html

params = urllib.parse.quote_plus('Driver={SQL Server};'
'Server='+server_name +';'
'Database='+db_name+';'
'Trusted_Connection=yes;')

engine = sa.create_engine("mssql+pyodbc:///?odbc_connect=%s" % params,fast_executemany=True, echo=True)

# https://chrisalbon.com/python/data_wrangling/pandas_dataframe_importing_csv/
def Load_Data(file_name):
    #data = genfromtxt(file_name, delimiter=',', skip_header=1, converters={0: lambda s: str(s)})
    data = pd.read_csv(file_name, skiprows=0)
    return data.values.tolist()

Base = declarative_base()
BaseGPX = declarative_base()

class Activities(Base):
    #Tell SQLAlchemy what the table name is and if there's any table-specific arguments it should know about
    __tablename__ = 'Activities'
    __table_args__ = {'schema':'STG'}
    #tell SQLAlchemy the name of column and its attributes:
    Activity_ID = Column(BigInteger, primary_key=True, nullable=False) #
    Activity_Date = Column(Date)
    Activity_Name = Column(String)
    Activity_Type = Column(String)
    Activity_Description = Column(String)
    Elapsed_Time = Column(Integer)
    Distance = Column(Float)
    Commute = Column(String)
    Activity_Gear = Column(String)
    Filename = Column(String)

class GPXFiles(BaseGPX):
    #Tell SQLAlchemy what the table name is and if there's any table-specific arguments it should know about
    __tablename__ = 'GPXFiles'
    __table_args__ = {'schema':'STG'}
    #tell SQLAlchemy the name of column and its attributes:
    GPXFile_ID = Column(BigInteger, primary_key=True, autoincrement=True, nullable=False) #
    GPXFile_Name = Column(String, nullable=True)
    GPXFile_Contents = Column(String, nullable=True)
    #GPXFile_XMLContents = Column(String, nullable=True)
    #GPXFile_Route_ID = Column(BigInteger, nullable=True)

# loop through files
# r=root, d=directories, f = files
for r, d, f in os.walk(path):
    for file in f:
        if fnmatch2.fnmatch2(file, searchstring):

            if file.endswith('.gpx'):
                print('gpxfile')
                # go straight into the STG.GPXFiles table
                #Create the session
                #GPXFiles.__table__.drop(bind=engine)
                GPXFiles.__table__.create(bind=engine, checkfirst=True) # if the table exists, it is dropped (not using checkfirts=true)
                session = sessionmaker(bind=engine)
                #print(session)
                session.configure(bind=engine)
                s = session()
                try:

                    data = Load_Data(os.path.join(r, file))
                    #print(data)
                    #for i in data:
                        #print(str(i[0])+'-'+str(i[1]))
                    record = GPXFiles(**{
                        #'GPXFile_ID' : i[0],
                        'GPXFile_Name' : file, #str(i[1]),
                        'GPXFile_Contents' : data
                    })

                    s.add(record) #Add all the records

                    s.commit() #Attempt to commit all the records

                except Exception as e:
                    print(e)
                    s.rollback() #Rollback the changes on error

                s.close() #Close the connection
                continue
            elif file.endswith('.gpx.gz'):
                # print('a GZ gpxfile')
                # unzip file first, then get the file into the STG.GPXFiles table
                continue
            elif file.endswith('.csv'):
                # the only file we really need is the activities.csv file
                # the columns are
                if file == 'activities.csv':
                    # import it

                    ### using pure panda
                    # insert CSV into the table
                    # https://stackoverflow.com/questions/31394998/using-sqlalchemy-to-load-csv-file-into-a-database

                    #Create the session
                    #Activities.__table__.drop(bind=engine)
                    Activities.__table__.create(bind=engine, checkfirst=True) # if the table exists, it is dropped (not using checkfirts=true)
                    session = sessionmaker(bind=engine)
                    print(session)
                    session.configure(bind=engine)
                    s = session()
                    try:

                        data = Load_Data(os.path.join(r, file))
                        #print(data)
                        for i in data:
                            #print(str(i[0])+'-'+str(i[1]))
                            record = Activities(**{
                                'Activity_ID' : i[0],
                                'Activity_Date' : i[1].replace(',', ''),
                                'Activity_Name' : str(i[2]),
                                'Activity_Type' : str(i[3]),
                                'Activity_Description' : str(i[4]),
                                'Elapsed_Time' : int(i[5]),
                                'Distance' : float(i[6].replace(',', '')),
                                'Commute' : str(i[7]),
                                'Activity_Gear' : str(i[8]),
                                'Filename' : str(i[9])
                            })

                            s.add(record) #Add all the records

                            s.commit() #Attempt to commit all the records

                    except Exception as e:
                        print(e)
                        s.rollback() #Rollback the changes on error

                    s.close() #Close the connection

            else:
                continue

            # move file to archive folder at same level as source folder
            # first check folder exist, if not, then create it

        continue

#conn.close

# now that this is done, we need to run the stored procedures to get the data from the staging tbales into the production conformed tables
# first to brng the data from staging into conformed area
# then run a bunch of procs to update some of th emissing attributes

# don't forget to clean up in some way


import os
import time
from collections import Counter

import numpy as np
from sklearn.preprocessing import MinMaxScaler

import core.leading_tree as lt
import core.lmca as lm
from core.delala_select import DeLaLA_select
from utils import common

def load_parameters(param_path):
    dataset = common.load_csv(param_path)
    dc = dataset[:, 0].astype(float)
    lt_num = dataset[:, 1].astype(int)
    length_scale = dataset[:, 2].astype(float)
    return dc, lt_num, length_scale

def removeMinorData(X, Labels, k):
    """
    Remove the points and labels that have no enough neighbors
    :param X:
    :param Labels:
    :param k:
    :return:
    """
    setLabels = set(Labels)
    classNum = len(setLabels)
    cntArray = np.zeros(classNum, dtype=int)
    i = 0
    rmInds = np.zeros(0, dtype=int)
    for l in setLabels:
        clCnt = list(Labels).count(l)
        cntArray[i] = clCnt
        i += 1
        if clCnt <= k:
            lbInds = [i for i in range(len(Labels)) if Labels[i] == l]
            rmInds = np.append(rmInds, lbInds)
            # print("remove class: ", l)
    if len(rmInds) > 0:
        Remove_index = X[rmInds]
        X = np.delete(X, rmInds, axis=0)
        Labels = np.delete(Labels, rmInds)
        return X, Labels, Remove_index
    else:
        return X, Labels, -1

order = 0
param_path = os.path.join("data", "parameters.csv")
dc_all, lt_num_all, length_scale_all = load_parameters(param_path)
SelectedInds_all = np.zeros(0, dtype=int)

def PredictLabel(train_AL, label_num, layer, dimension):
    global order, SelectedInds_all
    if label_num == 0:
        pass
    elif label_num == 1:  # predict that this subtree sample has the same label as its root node
        y_test = np.delete(y[train_AL], 0)  # test set labels
        y_predict = np.zeros(len(y_test), dtype=int) + y[train_AL[0]]
        y_predict_all[train_AL] = y[train_AL[0]]
        SelectedInds_all = np.append(SelectedInds_all, train_AL[0])
        return 0
    elif 2 <= label_num <= 3 or layer == 3:  # Prediction with LMCA
        D_A = D[train_AL]
        D_A = D_A[:, train_AL]
        LT = lt.LeadingTree(X_train=X[train_AL], dc=dc_all[order], lt_num=lt_num_all[order], D=D_A)
        LT.fit()
        LTgammaPara = LT.density * LT.delta
        selectedInds = DeLaLA_select(LTgammaPara, LT.density, LT.layer, y[train_AL], 2, label_num * 2, 0.5)
        selectedInds_universe = train_AL[selectedInds]
        X_train = X[selectedInds_universe]
        y_train = y[selectedInds_universe]
        lmca = lm.LMCA(dimension=dimension, init_method="kpca", verbose=True, max_iter=100, stepsize=1.E-2,
                       nn_active=False, length_scale=length_scale_all[order], k=1)
        lmca.fit(X_train, y_train)
        order += 1
        X_test = np.delete(X[train_AL], selectedInds, axis=0)  # After removing the training set, the test set samples are obtained
        y_test = np.delete(y[train_AL], selectedInds, axis=0)  # test set labels
        y_predict = np.zeros(len(y_test), dtype=int) - 1  # predict labels

        MatDist = common.euclidian_dist_square(X_train, X_test)  # The kernel matrix corresponding to the training set and the test set
        test_bnd_K = np.exp(-1 * lmca.length_scale * MatDist).T
        B = test_bnd_K.dot(lmca.Omega)  # Test set after dimensionality reduction
        A = lmca.K.dot(lmca.Omega)  # The training set after dimensionality reduction

        # Find the training sample with the closest Euclidean distance for each test sample, and predict that both have the same label.
        D_temp = common.euclidian_dist(B, A)
        Pa = np.zeros(len(y_predict), dtype=int)  # Pa[i] represents the index of the training sample closest to the test sample i
        for j in range(len(y_predict)):
            index1 = np.argmin(D_temp[j])
            Pa[j] = index1

        y_predict = y_train[Pa]
        index_predict = np.delete(train_AL, selectedInds, axis=0)
        y_predict_all[index_predict] = y_predict
        SelectedInds_all = np.append(SelectedInds_all, selectedInds_universe)
        return 0
    else:
        return 1

if __name__ == "__main__":
    dataset_path = os.path.join("data", "letter.csv")
    X, y = common.load_data(dataset_path, label_index=0, map_label=False)

    t1 = time.time()
    scalar = MinMaxScaler()
    X = scalar.fit_transform(X)
    D = common.euclidian_dist(X, X)

    remove_index_all = np.zeros(0, dtype=int)
    y_predict_all = np.zeros(len(y), dtype=int) - 1

    dc_lt_num_arr = [[0.12, 45, 1, 2], [0.19, 10, 2, 2], [0.12, 5, 3, 8]]

    def recursive_partitioning(_X, _D, _layer=0, _train_AL=None):
        global remove_index_all
        lt1 = lt.LeadingTree(X_train=_X, dc=dc_lt_num_arr[_layer][0],
                             lt_num=dc_lt_num_arr[_layer][1], D=_D)  # Constructing the lead tree for the entire dataset
        lt1.fit()
        for i in range(dc_lt_num_arr[_layer][1]):
            if _layer == 0:
                _train_AL2, _y_AL, _remove_index = removeMinorData(lt1.AL[i], y[lt1.AL[i]], 2)
            else:
                _train_AL2, _y_AL, _remove_index = removeMinorData(_train_AL[lt1.AL[i]], y[_train_AL[lt1.AL[i]]], 2)

            _label_num = len(np.unique(y[_train_AL2]))
            remove_index_all = np.append(remove_index_all, _remove_index)
            _a = PredictLabel(_train_AL2, _label_num, dc_lt_num_arr[_layer][2], dc_lt_num_arr[_layer][3])

            if _a == 1:
                print(f"The {_layer}th layer case3: The number of subtree {i} categories is {_label_num}, which needs to be divided again.")
                _D_2 = D[_train_AL2]
                _D_2 = _D_2[:, _train_AL2]
                recursive_partitioning(X[_train_AL2], _D_2, _layer + 1, _train_AL=_train_AL2)
        return lt1

    lt0 = recursive_partitioning(X, D)

    index_None = np.zeros(0, dtype=int)  # Returns -1 if there are no indexes that need to be removed, here they are to be removed
    for i in range(len(remove_index_all)):
        if remove_index_all[i] == -1:
            index_None = np.append(index_None, i)

    remove_index_all = np.delete(remove_index_all, index_None)
    y_remove_Select = np.delete(y_predict_all, np.append(SelectedInds_all, remove_index_all))
    D = common.euclidian_dist(X[remove_index_all], X[SelectedInds_all])
    Pa = np.zeros(len(remove_index_all), dtype=int)
    for i in range(len(remove_index_all)):
        index1 = np.argmin(D[i])
        Pa[i] = index1

    y_predict_all[remove_index_all] = y[SelectedInds_all][Pa]
    arr1 = y[remove_index_all] - y_predict_all[remove_index_all]
    count0 = Counter(arr1)[0]
    print(f'The accuracy of the remove sample prediction is {count0 / len(arr1)}, {count0}/{len(arr1)}')

    # Subtree accuracy
    for i in range(45):
        temp = np.setdiff1d(lt0.AL[i], SelectedInds_all)
        arr2 = y[temp] - y_predict_all[temp]
        count1 = Counter(arr2)[0]
        print(f'The accuracy of subtree {i} is {count1 / len(temp)}, {count1}/{len(temp)}')
    # Overall accuracy
    y_predict_all = np.delete(y_predict_all, SelectedInds_all)
    y_test_all = np.delete(y, SelectedInds_all)
    arr = y_test_all - y_predict_all
    count = Counter(arr)[0]

    t2 = time.time()
    print(f'A total of {len(SelectedInds_all)} points are selected, with an accuracy of {count / len(y_test_all)}, {count}/{len(y_test_all)}')
    print(f'Takes {t2 - t1} seconds')


from django.contrib import admin
from django.urls import path, include
from basesite import views

urlpatterns = [
    path('', views.index, name='index'),
    path('academics', views.academics, name='academics'),
    path('labs', views.labs, name='labs'),
    path('committee', views.committee, name='committee'),
    path('gallery', views.gallery, name='gallery'),
    path('hostel', views.hostel, name='hostel'),
    path('placements', views.placements, name='placements'),
    path('alumni', views.alumni, name='alumni'),
    path('library', views.library, name='library'),
    path('about', views.about, name='about'),
    path('contact', views.contact, name='contact'),
    path('coursempe', views.coursempe, name='coursempe'),
    path('coursemae', views.coursemae, name='coursemae'),
    path('coursecse', views.coursecse, name='coursecse'),
    path('insert', views.insert, name="insert")
]


# from django.shortcuts import render
from rest_framework import generics
from review_app.models import FarmersMarket, Vendor
from fm_api.serializers import FarmersMarketSerializer, VendorSerializer
# Create your views here.

class FarmersMarketListAPIView(generics.ListAPIView):
    queryset = FarmersMarket.objects.all()
    serializer_class = FarmersMarketSerializer

class FarmersMarketRetrieveAPIView(generics.RetrieveAPIView):
    queryset = FarmersMarket.objects.all()
    serializer_class = FarmersMarketSerializer

class VendorListAPIView(generics.ListAPIView):
    queryset = Vendor.objects.all()
    serializer_class = VendorSerializer

class VendorRetrieveAPIView(generics.RetrieveAPIView):
    queryset = Vendor.objects.all()
    serializer_class = VendorSerializer


import json
import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)

def hello(event, context):

    logger.info(f"AWS Lambda processing message from GitHub: {event}.")

    body = {
        "message": "Your function executed successfully!",
        "input": event
    }

    response = {
        "statusCode": 200,
        "body": json.dumps(body)
    }

    return response


from django.db import models
from django.utils.translation import ugettext_lazy as _
from django.conf import settings

# Create your models here.

class Country(models.Model):

    name = models.CharField(_("Name"), db_column='name', max_length = 150, null=True, blank=True)
    code2 = models.CharField(_("Code2"), db_column='code2', max_length = 2, unique = True)
    code3 = models.CharField(_("Code3"), db_column='code3', max_length = 3, unique = True)
    number = models.CharField(_("Number"), db_column='number', max_length = 3, unique = True)

class Document(models.Model):

    image = models.ImageField(_("Document image"), upload_to=settings.UPLOAD_TO, null=True, blank=True)


import frida
import sys

package_name = "com.jni.anto.kalip"

def get_messages_from_js(message, data):
            print(message)
            print (message['payload'])

def instrument_debugger_checks():

        hook_code = """
        setTimeout(function(){
        Dalvik.perform(function () {

            var TM = Dalvik.use("android.os.Debug");

            TM.isDebuggerConnected.implementation = function () {

                send("Called - isDebuggerConnected()");

            return false;
            };

        });

    },0);
        """

        return hook_code

process = frida.get_device_manager().enumerate_devices()[-1].attach(package_name)
script = process.create_script(instrument_debugger_checks())
script.on('message',get_messages_from_js)
script.load()
sys.stdin.read()


import tweepy
import pandas as pd
import re
import time
from textblob import TextBlob
from sqlalchemy import create_engine
import yaml
import json

TWITTER_CONFIG_FILE = '../auth.yaml'

with open(TWITTER_CONFIG_FILE, 'r') as config_file:
	config = yaml.load(config_file)

consumer_key = config['twitter']['consumer_key']
consumer_secret = config['twitter']['consumer_secret']
access_token = config['twitter']['access_token']
access_token_secret = config['twitter']['access_token_secret']

auth = tweepy.OAuthHandler(consumer_key, consumer_secret)
auth.set_access_token(access_token, access_token_secret)

# get geo for top 50 US cities by population
json_data=open('cities_trunc.json').read()
data = json.loads(json_data)

# define twitter API calls by geo coordinates to a 10mi radius
def get_tweet_payload(d):
	return {d['city']: '%s,%s,%s' % (d['latitude'], d['longitude'],'10mi')}

geo = {}
i = 0
while i < len(data):
	geo.update(get_tweet_payload(data[i]))
	i += 1

# manually setting queries against set of popular competitive games - leverage twitch_top.py to keep this dynamic
api = tweepy.API(auth, wait_on_rate_limit=True)
query = ['fortnite', 'overwatch', 'starcraft', 'dota', 'league of legends', 'CSGO', 'hearthstone' ,'pubg', 'tekken', 'ssbm']
d = []

# Since we are using Standard APIs we are limited in data volume, adding additional geo regions and/or queries can result in long run-time
for game in query:
	for city,coords in geo.items():
		public_tweets = [status for status in tweepy.Cursor(api.search,q=game, geocode=coords, count=100).items(1000)]
		for tweet in public_tweets:
			analysis = TextBlob(tweet.text)
			TweetText = re.sub('[^A-Za-z0-9]+', ' ', tweet.text)
			polarity = analysis.sentiment.polarity
			subjectivity = analysis.sentiment.subjectivity
			d.append((TweetText,
						polarity,
						subjectivity,
						game,
						city))

# use Pandas to format analyzed tweets into CSV file for appending to a database
timestr = time.strftime("%Y%m%d-%H%M%S")
filename = timestr + "_tweets.csv"
cols=['Tweet','polarity','subjectivity','game','city']
df = pd.DataFrame(d, columns=cols)
df = df[['Tweet','polarity','subjectivity','game','city']]
df.drop_duplicates(['Tweet'], keep='last')
df.to_csv(filename, encoding='utf-8-sig')


import os

from flask import Flask, jsonify, Blueprint
from flask_cors import CORS
from flask_sqlalchemy import SQLAlchemy
from flask_migrate import Migrate
from flask_bcrypt import Bcrypt
import flask_restplus
from werkzeug.contrib import fixers

# instantiate the extensions
db = SQLAlchemy()
migrate = Migrate()
bcrypt = Bcrypt()

def register_api(_app):
    from users.api.users import ns as users_ns
    from users.api.auth import ns as auth_ns

    blueprint = Blueprint('api', __name__)
    api = flask_restplus.Api(
        app=blueprint,
        doc=_app.config['SWAGGER_PATH'],
        version=_app.config['API_VERSION'],
        title='Shows On Demand - Users Service REST API',
        description='Shows on deman Users service API for users access.',
        validate=_app.config['RESTPLUS_VALIDATE']
    )
    api.add_namespace(auth_ns, path='/{}/auth'.format(_app.config['API_VERSION']))
    api.add_namespace(users_ns, path='/{}/users'.format(_app.config['API_VERSION']))
    _app.register_blueprint(blueprint)

def create_app():
    # instantiate the app
    app = Flask(__name__)

    # enable CORS
    CORS(app)

    # set config
    app_settings = os.getenv('APP_SETTINGS')
    app.config.from_object(app_settings)
    app_config = app.config

    # set up extensions
    db.init_app(app)
    bcrypt.init_app(app)
    migrate.init_app(app, db)

    app.wsgi_app = fixers.ProxyFix(app.wsgi_app)

    # register blueprints
    register_api(app)
    return app


from django.db import models

class UserMailmapManager(models.Manager):
    """A queryset manager which defers all :class:`models.DateTimeField` fields, to avoid
    resetting them to an old value involuntarily."""

    @classmethod
    def deferred_fields(cls):
        try:
            return cls._deferred_fields
        except AttributeError:
            cls._deferred_fields = [
                field.name
                for field in UserMailmap._meta.get_fields()
                if isinstance(field, models.DateTimeField) and not field.auto_now
            ]
            return cls._deferred_fields

    def get_queryset(self):
        return super().get_queryset().defer(*self.deferred_fields())

class UserMailmap(models.Model):
    """
    Model storing mailmap settings submitted by users.
    """

    user_id = models.CharField(max_length=50, null=True)
    """Optional user id from Keycloak"""

    from_email = models.TextField(unique=True, null=False)
    """Email address to find author in the archive"""

    from_email_verified = models.BooleanField(default=False)
    """Indicates if the from email has been verified"""

    from_email_verification_request_date = models.DateTimeField(null=True)
    """Last from email verification request date"""

    display_name = models.TextField(null=False)
    """Display name to use for the author instead of the archived one"""

    display_name_activated = models.BooleanField(default=False)
    """Indicates if the new display name should be used"""

    to_email = models.TextField(null=True)
    """Optional new email to use in the display name instead of the archived one"""

    to_email_verified = models.BooleanField(default=False)
    """Indicates if the to email has been verified"""

    to_email_verification_request_date = models.DateTimeField(null=True)
    """Last to email verification request date"""

    mailmap_last_processing_date = models.DateTimeField(null=True)
    """Last mailmap synchronisation date with swh-storage"""

    last_update_date = models.DateTimeField(auto_now=True)
    """Last date that mailmap model was updated"""

    class Meta:
        app_label = "swh_web_mailmap"
        db_table = "user_mailmap"

    # Defer _date fields by default to avoid updating them by mistake
    objects = UserMailmapManager()

    @property
    def full_display_name(self) -> str:
        if self.to_email is not None and self.to_email_verified:
            return "%s <%s>" % (self.display_name, self.to_email)
        else:
            return self.display_name

class UserMailmapEvent(models.Model):
    """
    Represents an update to a mailmap object
    """

    timestamp = models.DateTimeField(auto_now=True, null=False)
    """Timestamp of the moment the event was submitted"""

    user_id = models.CharField(max_length=50, null=False)
    """User id from Keycloak of the user who changed the mailmap.
    (Not necessarily the one who the mail belongs to.)"""

    request_type = models.CharField(max_length=50, null=False)
    """Either ``add`` or ``update``."""

    request = models.TextField(null=False)
    """JSON dump of the request received."""

    successful = models.BooleanField(default=False, null=False)
    """If False, then the request failed or crashed before completing,
    and may or may not have altered the database's state."""

    class Meta:
        indexes = [
            models.Index(fields=["timestamp"]),
        ]
        app_label = "swh_web_mailmap"
        db_table = "user_mailmap_event"


from . import views
from django.urls import path

app_name = 'bankapp'    #namespace

urlpatterns = [
    path('',views.home,name='home'),
    path('login/',views.login,name='login'),
    path('register/',views.register,name='register'),
    path('logout/',views.logout,name='logout'),
    path('user/',views.user,name='user'),
    path('form/',views.form,name="form"),
]


# sourcery skip: do-not-use-staticmethod
"""
A module that contains the AIConfig class object that contains the configuration
"""
from __future__ import annotations

import os
from typing import Type

import yaml

class AIConfig:
    """
    A class object that contains the configuration information for the AI

    Attributes:
        ai_name (str): El nombre de la IA.
        ai_role (str): La descripción de la función de la IA.
        ai_goals (list): La lista de objetivos que la IA debe cumplir.
    """

    def __init__(
        self, ai_name: str = "", ai_role: str = "", ai_goals: list | None = None
    ) -> None:
        """
        Initialize a class instance

        Parameters:
            ai_name (str): El nombre de la IA.
            ai_role (str): La descripción de la función de la IA.
            ai_goals (list): La lista de objetivos que la IA debe cumplir.
        Returns:
            None
        """
        if ai_goals is None:
            ai_goals = []
        self.ai_name = ai_name
        self.ai_role = ai_role
        self.ai_goals = ai_goals

    # Soon this will go in a folder where it remembers more stuff about the run(s)
    SAVE_FILE = os.path.join(os.path.dirname(__file__), "..", "ai_settings.yaml")

    @staticmethod
    def load(config_file: str = SAVE_FILE) -> "AIConfig":
        """
        Returns class object with parameters (ai_name, ai_role, ai_goals) loaded from
          yaml file if yaml file exists,
        else returns class with no parameters.

        Parameters:
           config_file (int): The path to the config yaml file.
             DEFAULT: "../ai_settings.yaml"

        Returns:
            cls (object): An instance of given cls object
        """

        try:
            with open(config_file, encoding="utf-8") as file:
                config_params = yaml.load(file, Loader=yaml.FullLoader)
        except FileNotFoundError:
            config_params = {}

        ai_name = config_params.get("ai_name", "")
        ai_role = config_params.get("ai_role", "")
        ai_goals = config_params.get("ai_goals", [])
        # type: Type[AIConfig]
        return AIConfig(ai_name, ai_role, ai_goals)

    def save(self, config_file: str = SAVE_FILE) -> None:
        """
        Saves the class parameters to the specified file yaml file path as a yaml file.

        Parameters:
            config_file(str): The path to the config yaml file.
              DEFAULT: "../ai_settings.yaml"

        Returns:
            None
        """

        config = {
            "ai_name": self.ai_name,
            "ai_role": self.ai_role,
            "ai_goals": self.ai_goals,
        }
        with open(config_file, "w", encoding="utf-8") as file:
            yaml.dump(config, file, allow_unicode=True)

    def construct_full_prompt(self) -> str:
        """
        Returns a prompt to the user with the class information in an organized fashion.

        Parameters:
            None

        Returns:
            full_prompt (str): A string containing the initial prompt for the user
              including the ai_name, ai_role and ai_goals.
        """

        prompt_start = (
            "Sus decisiones deben tomarse siempre de forma independiente, sin"
            " buscar asistencia al usuario. Aprovecha tus puntos fuertes como LLM y persigue"
            " estrategias sencillas y sin complicaciones jurídicas."
            ""
        )

        from nancy.prompt import get_prompt

        # Construct full prompt
        full_prompt = (
            f"Usted es {self.ai_name}, {self.ai_role}\n{prompt_start}\n\nOBJETIVOS:\n\n"
        )
        for i, goal in enumerate(self.ai_goals):
            full_prompt += f"{i+1}. {goal}\n"

        full_prompt += f"\n\n{get_prompt()}"
        return full_prompt


def _extend_pre_ranges(df, upstream:int=0, downstream:int=0, start:str="Start", end:str="End", strand:str="Strand"):
    strand_rm = False
    if strand not in df.columns:
        strand_rm = True
        df[strand] = "+"
    df.loc[df[strand] == "+", start] -= upstream
    df.loc[df[strand] == "-", start] -= downstream
    df.loc[df[strand] == "+", end] += downstream
    df.loc[df[strand] == "-", end] += upstream
    if strand_rm:
        del df[strand]
    return df

def _extend_ranges(gf,
                   min_upstream:int=1000,
                   max_upstream:int=100000,
                   min_downstream:int=1000,
                   max_downstream:int=100000,
                   gene_upstream:int=5000,
                   gene_downstream:int=0,
                   gene_scale_factor:float=5.):
    import numpy as np
    gf["gene_length"] = gf["End"] - gf["Start"]
    ### In case of duplicated indices, sum lengths across indices
    gb = gf.groupby(level=0).agg(gene_length=("gene_length", sum))
    gf["gene_length"] = gb.loc[gf.index.values, "gene_length"].values
    ### scale by inverse gene length
    gs = 1/gf["gene_length"].values
    ### min max scale, plus epsilon to avoid divide by 0
    gs = (gs - np.min(gs)) / (np.max(gs) - np.min(gs)) + 1e-300
    ### expand to 1 to GSF range. But take log, so log1p((gs-1)*s) = log(1 + (gs - 1) * s)
    gf["log_gene_scale"] = np.log1p((gene_scale_factor - 1) * gs)
    #gf.index = gf["gene_id"].values
    gf = _extend_pre_ranges(gf, upstream=gene_upstream, downstream=gene_downstream)
    gf["MinStart"] = gf["Start"].values
    gf["MinEnd"] = gf["End"].values
    gf = _extend_pre_ranges(gf, upstream=min_upstream, downstream=min_downstream, start="MinStart", end="MinEnd")
    gf["interval"] = ["%s:%d-%d" % (chrom, start, end) for chrom, start, end in zip(gf["Chromosome"], gf["Start"], gf["End"])]
    gf["min_interval"] = ["%s:%d-%d" % (chrom, start, end) for chrom, start, end in zip(gf["Chromosome"], gf["MinStart"], gf["MinEnd"])]
    return gf

def estimate_features_archr(adata, feature_df,
                            min_upstream:int=1000,
                            max_upstream:int=100000,
                            min_downstream:int=1000,
                            max_downstream:int=100000,
                            gene_upstream:int=5000,
                            gene_downstream:int=0,
                            target_sum:int=None,
                            gene_scale_factor:float=5.,
                            peak_column:str=None, ## If not provided, will use peak index
                            feature_column:str=None,
                            var_column_tolerance:float=0.999, ### Tolerance used for determining if a feature should be kept, in case of weird things
                            distal:bool=True, ### Use nearest gene to a peak if unassigned
                            log1p:bool=False,
                            save_raw:bool=False,
                            layer:str=None):
    import numpy as np
    import pandas as pd
    import scipy.sparse
    import pyranges
    import anndata
    import scanpy as sc
    from .timer import template
    sw = template()
    if not isinstance(feature_df, pd.DataFrame):
        raise ValueError("Feature_df is not a dataframe")
    if not np.all(np.isin(["Chromosome", "Start", "End"], feature_df.columns)):
        raise ValueError("Feature_df does not have ranges")
    if feature_column is not None:
        feature_df.index = feature_df[feature_column].values
    with sw("Extending ranges"):
        gf = _extend_ranges(feature_df, min_upstream=min_upstream, min_downstream=min_downstream,
                            max_upstream=max_upstream, max_downstream=max_downstream,
                            gene_upstream=gene_upstream, gene_downstream=gene_downstream,
                            gene_scale_factor=gene_scale_factor)
        gr = pyranges.from_dict({"Chromosome": gf["Chromosome"],
                                 "Start": gf["Start"],
                                 "End": gf["End"],
                                 "feature_index": gf.index.values,
                                 "feature_interval": gf["interval"].values})
        mingr = pyranges.from_dict({"Chromosome": gf["Chromosome"],
                                    "Start": gf["MinStart"],
                                    "End": gf["MinEnd"],
                                    "feature_index": gf.index.values,
                                    "feature_interval": gf["interval"].values})
    ##
    ## Now, get peak ranges (pr) from peak frame (pf)
    ##
    with sw("Extracting peak ranges"):
        if peak_column is None:
            pstr = adata.var_names.values
        else:
            pstr = adata.var[peak_column].values
        pf = pd.DataFrame([x.replace(":", "-", 1).split("-") for x in pstr], columns=["Chromosome", "Start", "End"], index=adata.var_names)
        pr = pyranges.from_dict({"Chromosome": pf["Chromosome"], "Start": pf["Start"], "End": pf["End"], "peak_name": pf.index.values})
    with sw("Calculating overlaps"):
        iif = gf.drop_duplicates("interval")
        iif.index = iif["interval"]
        ## Once peak ranges are gathered, find intersecting gene bodies:
        inter_df = pr.join(gr).df.loc[:, ["peak_name", "feature_index", "feature_interval"]]
        inter_df["Distance"] = 0
        ## Then, find genes with minimum upstream/downstream distance away
        min_df = pr.join(mingr).df.loc[:, ["peak_name", "feature_index", "feature_interval"]]
        ## diff. is accurate, unless overlapping intervals. Do not need to worry, as duplicates from inter_df will take care
        diff = pf.loc[min_df["peak_name"].values, ["Start", "Start", "End", "End"]].values.astype(int) - iif.loc[min_df["feature_interval"].values, ["Start", "End", "Start", "End"]].values.astype(int)
        min_df["Distance"] = np.abs(diff).min(1) + 1
        ## Finally, find distal. Only need nearest gene

        if distal:
            distance_df = pr.nearest(gr).df.loc[:, ["peak_name", "feature_index", "feature_interval", "Distance"]]
            ## Concat such that 1) prioritized intersections, then 2) minimum distance away, then 3) distal
            df = pd.concat([inter_df, min_df, distance_df]).drop_duplicates(["peak_name", "feature_index"])
        else:
            df = pd.concat([inter_df, min_df]).drop_duplicates(["peak_name", "feature_index"])
        df["weight"] = np.exp(-1 - np.abs(df["Distance"]) / 5000. + iif.loc[df["feature_interval"].values, "log_gene_scale"].values)
    with sw("Calculating accessibility"):
        if gf.index.duplicated().sum() > 0:
            ### Get columns that are the same across repeated indices
            nf = gf.groupby(level=0).nunique()
            gf = gf.loc[~gf.index.duplicated(keep="first"), nf.columns[(nf==1).mean() >= var_column_tolerance]]
        S = scipy.sparse.csr_matrix((df["weight"].values,
                                     (pf.index.get_indexer(df["peak_name"].values),
                                      gf.index.get_indexer(df["feature_index"].values))),
                                    shape=(pf.shape[0], gf.shape[0]))
        if layer is not None and layer in adata.layers:
            X = adata.layers[layer]
        else:
            X = adata.X
        gdata = anndata.AnnData(X.dot(S), obs=adata.obs, var=gf, dtype=np.float32, obsm=adata.obsm, obsp=adata.obsp,
                                uns={k: v for k, v in adata.uns.items() if k in ["neighbors", "files", "lsi", "pca", "umap", "leiden"]})
    if save_raw:
        gdata.layers["raw"] = gdata.X.copy()
    if target_sum is not None and target_sum > 0:
        sc.pp.normalize_total(gdata, target_sum=target_sum)
    else:
        print("Using median normalization")
        sc.pp.normalize_total(gdata)
    if log1p:
        sc.pp.log1p(gdata)
    return gdata

def estimate_genes_archr(adata, gtf:str,
                         min_upstream:int=1000,
                         max_upstream:int=100000,
                         min_downstream:int=1000,
                         max_downstream:int=100000,
                         gene_upstream:int=5000,
                         gene_downstream:int=0,
                         target_sum:int=None,
                         gene_scale_factor:float=5.,
                         peak_column:str=None, ## If not provided, will use peak index
                         feature_column:str="gene_id", ### If not provided, will use feature index
                         log1p:bool=True,
                         distal:bool=True,
                         save_raw:bool=False,
                         layer:str=None):
    import numpy as np
    import pandas as pd
    import scipy.sparse
    import pyranges
    import anndata
    import scanpy as sc
    from .timer import template
    sw = template()
    with sw("Reading GTF"):
        gf = pyranges.read_gtf(gtf).df
        gf = gf.loc[gf["Feature"] == "gene", :]
        if feature_column in gf.columns:
            gf.index = gf[feature_column].values.astype(str)
    gdata = estimate_features_archr(adata, feature_df=gf,
                                   min_upstream=min_upstream, min_downstream=min_downstream,
                                   max_upstream=max_upstream, max_downstream=max_downstream,
                                   gene_upstream=gene_upstream, gene_downstream=gene_downstream,
                                   target_sum=target_sum, gene_scale_factor=gene_scale_factor,
                                   peak_column=peak_column,
                                   feature_column=feature_column,
                                    log1p=log1p, layer=layer, save_raw=save_raw)
    gdata.var = gdata.var.loc[:, ["gene_id", "gene_name"]]
    gdata.var_names = gdata.var["gene_name"].values
    gdata.var_names_make_unique()
    del gdata.var["gene_name"]
    gdata.var.columns = ["gene_ids"]
    add_gene_length(gdata.var, gtf)
    return gdata

def get_tss(tss:str):
    import pandas as pd
    tss = pd.read_csv(tss, sep="\t", header=None)
    tss.columns = ["Chromosome", "Start", "End", "gene_id", "score", "strand"]
    df = tss.groupby(["Chromosome", "gene_id", "strand"]).agg(left=("Start", "min"),
                                                              right=("End", "max")).reset_index()
    df["interval"] = df["Chromosome"] + ":" + df["left"].astype(str) + "-" + df["right"].astype(str)
    df.index = df["gene_id"].values
    return df

def add_interval(var, tss:str, inplace=True):
    tf = get_tss(tss)
    interval = [tf["interval"].get(g, "NA") for g in var["gene_ids"]]
    if inplace:
        var["interval"] = interval
    else:
        import pandas as pd
        return pd.Series(interval, index=var.index, name="interval")

def add_gene_length(var, gtf:str=None, inplace=True):
    import pandas as pd
    import pyranges
    from .timer import template
    sw = template()
    with sw("Reading GTF"):
        gf = pyranges.read_gtf(gtf).df
        gf = gf.loc[gf["Feature"] == "gene",:]
        gf["gene_length"] = gf["End"] - gf["Start"]
        gf.index = gf["gene_id"].values
    gl = [gf["gene_length"].get(g, -1) for g in var["gene_ids"]]
    gs = [gf["Strand"].get(g, "*") for g in var["gene_ids"]]
    if not inplace:
        var = var.copy()
    var["gene_length"] = gl
    var["strand"] = gs
    if not inplace:
        return var.loc[:, ["gene_length", "strand"]]

def add_gene_info(var, gene_info:str=None, inplace=True):
    """Add a STAR geneInfo.tab file to .var"""
    import pandas as pd
    df = pd.read_csv(gene_info, sep="\t", skiprows=[0], header=None)
    df.index = df[0].values
    gi = [df[2].get(g, "NA") for g in var["gene_ids"]]
    if inplace:
        var["gene_type"] = gi
    else:
        return pd.Series(gi, index=var.index, name="gene_type")


import json
from sksurv.functions import StepFunction
from sksurv.linear_model import CoxPHSurvivalAnalysis
from sksurv.metrics import concordance_index_censored
from sksurv.nonparametric import nelson_aalen_estimator, kaplan_meier_estimator

from core.cox_wrapper import CoxFairBaseline
from core.drawing import draw_points_tsne
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sksurv.ensemble import RandomSurvivalForest

from exp_config import CONFIG, RES_DIR
from core.cox_generator import CoxGenerator
from survshap import SurvivalModelExplainer, ModelSurvSHAP
from survbex.estimators import BeranModel
from survbex.explainers import SurvBexExplainer

########################################################################################################################
# ------------------------------------------------ PREPARE DATA --------------------------------------------------------
########################################################################################################################
def get_cox_data(coefs: np.ndarray):
    cox_generator = CoxGenerator(coefs=coefs)
    x_cox_train, x_cox_test, y_cox_train, y_cox_test = train_test_split(
        *cox_generator.generate_data(size=CONFIG['TRAIN_SIZE'], censored_part=0.2),
        train_size=0.7
    )

    x_cox_train = pd.DataFrame(x_cox_train, columns=[f'f{i + 1}' for i in range(len(coefs))])
    x_cox_test = pd.DataFrame(x_cox_test, columns=[f'f{i + 1}' for i in range(len(coefs))])

    return [x_cox_train, y_cox_train], [x_cox_test, y_cox_test]

# np.random.seed(42)
# train, test = get_veterans_data()
cox_clusters = [get_cox_data(coefs=cox_coefs) for cox_coefs in CONFIG['COX_COEFS_CLS']]

cox_clusters = [
    (
        [cox_cluster[0][0] + 2.0 / len(cox_clusters) * cl_i, cox_cluster[0][1]],
        [cox_cluster[1][0] + 2.0 / len(cox_clusters) * cl_i, cox_cluster[1][1]]
        # [cox_cluster[0][0] + 1. * cl_i, cox_cluster[0][1]],
        # [cox_cluster[1][0] + 1. * cl_i, cox_cluster[1][1]]
    )
    for cl_i, cox_cluster in enumerate(cox_clusters)
]

all_train = [
    pd.concat([cox_cluster[0][0] for cox_cluster in cox_clusters]),
    np.hstack([cox_cluster[0][1] for cox_cluster in cox_clusters])
]

all_test = [
    pd.concat([cox_cluster[1][0] for cox_cluster in cox_clusters]),
    np.hstack([cox_cluster[1][1] for cox_cluster in cox_clusters])
]

# Use SurvLimeExplainer class to find the feature importance
training_events = np.array([event for event, _ in all_train[1]])
training_times = np.array([time for _, time in all_train[1]])
training_features = all_train[0]

test_events = np.array([event for event, _ in all_test[1]])
test_times = np.array([time for _, time in all_test[1]])
test_features = all_test[0]

with open(f'{RES_DIR}/dataset.json', 'w+') as fp:
    json.dump(fp=fp, obj=dict(
        training_features=training_features.to_dict(orient='raw'),
        training_events=training_events.tolist(),
        training_times=training_times.tolist(),
        test_features=test_features.to_dict(orient='raw'),
        test_events=test_events.tolist(),
        test_times=test_times.tolist()
    ))

########################################################################################################################
# ------------------------------------------------ BUILD BBOX ----------------------------------------------------------
########################################################################################################################
if CONFIG['BBOX'] == 'rf':
    model = RandomSurvivalForest(n_estimators=100, max_samples=min(500, len(all_train[0])), max_depth=8)
    model.fit(all_train[0], all_train[1])
    pred_surv_fn = model.predict_survival_function
    pred_hazard_fn = model.predict_cumulative_hazard_function
    pred_risk_fn = model.predict
elif CONFIG['BBOX'] == 'beran':
    assert len(CONFIG['COX_COEFS_CLS']) == 1
    model = BeranModel(kernel_width=250, kernel_name='gaussian')
    model.fit(X=all_train[0].to_numpy(), b=CONFIG['COX_COEFS_CLS'][0],
              y_events=training_events, y_event_times=training_times)

    def surv_np_to_step_surv(surv_arr: np.ndarray):
        return np.array([StepFunction(x=model.unique_times_, y=sample) for sample in surv_arr])

    pred_surv_fn = lambda X: surv_np_to_step_surv(model.predict_survival_torch_optimized(X))
    pred_hazard_fn = lambda X: -np.log(model.predict_survival_torch_optimized(X))
    pred_risk_fn = lambda X: np.sum(pred_hazard_fn(X), axis=1)

elif 'cox' in CONFIG['BBOX']:
    model = CoxPHSurvivalAnalysis(alpha=1)

    model.fit(all_train[0], all_train[1])
    pred_surv_fn = model.predict_survival_function
    pred_hazard_fn = model.predict_cumulative_hazard_function
    pred_risk_fn = model.predict

    if CONFIG['BBOX'] in ['cox_na', 'cox_km']:
        if CONFIG['BBOX'] == 'cox_na':
            cox_fair_baseline = CoxFairBaseline(
                training_events=training_events,
                training_times=training_times,
                baseline_estimator_f=nelson_aalen_estimator
            )
        elif CONFIG['BBOX'] == 'cox_km':
            cox_fair_baseline = CoxFairBaseline(
                training_events=training_events,
                training_times=training_times,
                baseline_estimator_f=kaplan_meier_estimator
            )
        else:
            raise Exception(f'Undefined cox model = {CONFIG["BBOX"]}')

        model.coef_ /= np.abs(model.coef_).sum()
        pred_surv_fn = lambda X: cox_fair_baseline.predict_survival_function(X, cox_coefs=model.coef_)
        pred_hazard_fn = lambda X: cox_fair_baseline.predict_cum_hazard_from_surv_np(X, cox_coefs=model.coef_)
        pred_risk_fn = lambda X: np.dot(X, model.coef_)
    elif CONFIG['BBOX'] != 'cox':
        raise Exception(f'Undefined cox model = {CONFIG["BBOX"]}')
else:
    raise Exception(f"Undefined bbox = {CONFIG['BBOX']}")

cindex_train = concordance_index_censored(
    event_indicator=training_events, event_time=training_times, estimate=pred_risk_fn(training_features))[0]
print(f'cindex train = {cindex_train}')
cindex_test = concordance_index_censored(
    event_indicator=test_events, event_time=test_times, estimate=pred_risk_fn(test_features))[0]
print(f'cindex test = {cindex_test}')

########################################################################################################################
# ------------------------------------------------ SELECT POINTS TO EXPLAIN --------------------------------------------
########################################################################################################################
# draw_comparison(ex_i=random.randint(0, len(test)))
cluster_centroids = [
    cox_cluster[0][0].mean() + all_test[0].std() * CONFIG['DATA_POINT_DEV']
    for cox_cluster in cox_clusters
]

cl_distances = [
    [sum((cl_centroid - fs) ** 2) for fs in all_test[0].to_numpy()]
    for cl_centroid in cluster_centroids
]

exp_test_ids = [np.argmin(distances) for distances in cl_distances]

draw_points_tsne(
    pt_groups=[
        *[cox_cluster[0][0].to_numpy() for cox_cluster in cox_clusters],
        *list(all_test[0].to_numpy()[exp_test_ids])
    ],
    names=[
        *[f'cl{i}' for i, _ in enumerate(cox_clusters)],
        *[f'ex for cl {i}' for i, _ in enumerate(exp_test_ids)]
    ],
    colors=[None] * len(cox_clusters) * 2,
    path=f'{RES_DIR}/clusters.png'
    # path=f'clusters.png'
)

with open(RES_DIR.joinpath("y_true.json"), 'w+') as fp:
    json.dump(
        fp=fp,
        obj=[
            dict(event=bool(all_test[1][ex_i][0]), event_time=all_test[1][ex_i][1])
            for ex_i in exp_test_ids
        ]
    )

########################################################################################################################
# ------------------------------------------------ SurvSHAP ------------------------------------------------------------
########################################################################################################################
surv_shap = SurvivalModelExplainer(model, all_test[0].iloc[exp_test_ids], all_test[1][exp_test_ids],
                                   predict_survival_function=lambda model, X: pred_surv_fn(X))

exp_survshap = ModelSurvSHAP(random_state=42)
exp_survshap.fit(surv_shap)
shap_explanations = np.array(
    [
        [
            imp[1]
            for imp in pt_exp.simplified_result.values
        ]
        for pt_exp in exp_survshap.individual_explanations

    ]
)

with open(RES_DIR.joinpath("explanation_shap.json"), 'w+') as fp:
    json.dump(fp=fp, obj=shap_explanations.tolist())

########################################################################################################################
# ------------------------------------------------ SurvLIME ------------------------------------------------------------
########################################################################################################################
explainer = SurvBexExplainer(
    training_features=training_features,
    training_events=list(training_events),
    training_times=list(training_times),
    model_output_times=model.event_times_,
    kernel_width=CONFIG['KERNEL_WIDTH']
)

cox_explanations = np.array(
    [
        explainer.explain_instance(
            data_row=all_test[0].iloc[ex_i],
            predict_fn=pred_surv_fn,
            num_samples=CONFIG['NEIGH_SIZE'],
            type_fn='survival',
            optimizer='convex'
        )
        for ex_i in exp_test_ids
    ]
)

with open(RES_DIR.joinpath("explanation_cox.json"), 'w+') as fp:
    json.dump(fp=fp, obj=cox_explanations.tolist())

########################################################################################################################
# ------------------------------------------------ SurvBeX -------------------------------------------------------------
########################################################################################################################
beran_explanations = []
for cl_i, ex_i in enumerate(exp_test_ids):
    beran_explanations.append(
        explainer.explain_instance(
            data_row=all_test[0].iloc[ex_i],
            predict_fn=pred_surv_fn,
            num_samples=CONFIG['NEIGH_SIZE'],
            num_val_samples=CONFIG['NEIGH_VAL_SIZE'],
            type_fn='survival',
            optimizer='gradient',
            grid_info_file=f"{RES_DIR}/optimization_cl={cl_i}.csv",
            max_iter=CONFIG['MAX_ITER']

        )
    )

with open(RES_DIR.joinpath("explanation_beran.json"), 'w+') as fp:
    json.dump(
        fp=fp,
        obj=np.array(beran_explanations).tolist()
    )