import os

import pytest
import numpy as np

from pynpoint.core.pypeline import Pypeline
from pynpoint.readwrite.fitsreading import FitsReadingModule
from pynpoint.processing.frameselection import RemoveFramesModule, FrameSelectionModule, \
                                               RemoveLastFrameModule, RemoveStartFramesModule, \
                                               ImageStatisticsModule, FrameSimilarityModule, \
                                               SelectByAttributeModule, ResidualSelectionModule
from pynpoint.util.tests import create_config, remove_test_data, create_star_data

class TestFrameSelection:

    def setup_class(self) -> None:

        self.limit = 1e-10
        self.test_dir = os.path.dirname(__file__) + '/'

        create_star_data(self.test_dir+'images')
        create_config(self.test_dir+'PynPoint_config.ini')

        self.pipeline = Pypeline(self.test_dir, self.test_dir, self.test_dir)

    def teardown_class(self) -> None:

        remove_test_data(self.test_dir, folders=['images'])

    def test_read_data(self) -> None:

        module = FitsReadingModule(name_in='read',
                                   image_tag='read',
                                   input_dir=self.test_dir+'images',
                                   overwrite=True,
                                   check=True)

        self.pipeline.add_module(module)
        self.pipeline.run_module('read')

        data = self.pipeline.get_data('read')
        assert np.sum(data) == pytest.approx(105.54278879805277, rel=self.limit, abs=0.)
        assert data.shape == (10, 11, 11)

        attr = self.pipeline.get_attribute('read', 'NDIT', static=False)
        assert np.sum(attr) == pytest.approx(10, rel=self.limit, abs=0.)
        assert attr.shape == (2, )

        attr = self.pipeline.get_attribute('read', 'NFRAMES', static=False)
        assert np.sum(attr) == pytest.approx(10, rel=self.limit, abs=0.)
        assert attr.shape == (2, )

        self.pipeline.set_attribute('read', 'NDIT', [4, 4], static=False)

    def test_remove_last_frame(self) -> None:

        module = RemoveLastFrameModule(name_in='last',
                                       image_in_tag='read',
                                       image_out_tag='last')

        self.pipeline.add_module(module)
        self.pipeline.run_module('last')

        data = self.pipeline.get_data('last')
        assert np.sum(data) == pytest.approx(84.68885503527224, rel=self.limit, abs=0.)
        assert data.shape == (8, 11, 11)

        self.pipeline.set_attribute('last', 'PARANG', np.arange(8.), static=False)
        self.pipeline.set_attribute('last', 'STAR_POSITION', np.full((8, 2), 5.), static=False)

        attr = self.pipeline.get_attribute('last', 'PARANG', static=False)
        assert np.sum(attr) == pytest.approx(28., rel=self.limit, abs=0.)
        assert attr.shape == (8, )

        attr = self.pipeline.get_attribute('last', 'STAR_POSITION', static=False)
        assert np.sum(attr) == pytest.approx(80., rel=self.limit, abs=0.)
        assert attr.shape == (8, 2)

    def test_remove_start_frame(self) -> None:

        module = RemoveStartFramesModule(frames=1,
                                         name_in='start',
                                         image_in_tag='last',
                                         image_out_tag='start')

        self.pipeline.add_module(module)
        self.pipeline.run_module('start')

        data = self.pipeline.get_data('start')
        assert np.sum(data) == pytest.approx(64.44307047549808, rel=self.limit, abs=0.)
        assert data.shape == (6, 11, 11)

        attr = self.pipeline.get_attribute('start', 'PARANG', static=False)
        assert np.sum(attr) == pytest.approx(24., rel=self.limit, abs=0.)
        assert attr.shape == (6, )

        attr = self.pipeline.get_attribute('start', 'STAR_POSITION', static=False)
        assert np.sum(attr) == pytest.approx(60., rel=self.limit, abs=0.)
        assert attr.shape == (6, 2)

    def test_remove_frames(self) -> None:

        module = RemoveFramesModule(name_in='remove',
                                    image_in_tag='start',
                                    selected_out_tag='selected',
                                    removed_out_tag='removed',
                                    frames=[2, 5])

        self.pipeline.add_module(module)
        self.pipeline.run_module('remove')

        data = self.pipeline.get_data('selected')
        assert np.sum(data) == pytest.approx(43.68337741822863, rel=self.limit, abs=0.)
        assert data.shape == (4, 11, 11)

        data = self.pipeline.get_data('removed')
        assert np.sum(data) == pytest.approx(20.759693057269445, rel=self.limit, abs=0.)
        assert data.shape == (2, 11, 11)

        attr = self.pipeline.get_attribute('selected', 'PARANG', static=False)
        assert np.sum(attr) == pytest.approx(14., rel=self.limit, abs=0.)
        assert attr.shape == (4, )

        attr = self.pipeline.get_attribute('selected', 'STAR_POSITION', static=False)
        assert np.sum(attr) == pytest.approx(40., rel=self.limit, abs=0.)
        assert attr.shape == (4, 2)

        attr = self.pipeline.get_attribute('removed', 'PARANG', static=False)
        assert np.sum(attr) == pytest.approx(10., rel=self.limit, abs=0.)
        assert attr.shape == (2, )

        attr = self.pipeline.get_attribute('removed', 'STAR_POSITION', static=False)
        assert np.sum(attr) == pytest.approx(20., rel=self.limit, abs=0.)
        assert attr.shape == (2, 2)

    def test_frame_selection(self) -> None:

        module = FrameSelectionModule(name_in='select1',
                                      image_in_tag='start',
                                      selected_out_tag='selected1',
                                      removed_out_tag='removed1',
                                      index_out_tag='index1',
                                      method='median',
                                      threshold=2.,
                                      fwhm=0.1,
                                      aperture=('circular', 0.1),
                                      position=(None, None, 0.2))

        self.pipeline.add_module(module)
        self.pipeline.run_module('select1')

        data = self.pipeline.get_data('selected1')
        assert np.sum(data) == pytest.approx(54.58514780071149, rel=self.limit, abs=0.)
        assert data.shape == (5, 11, 11)

        data = self.pipeline.get_data('removed1')
        assert np.sum(data) == pytest.approx(9.857922674786586, rel=self.limit, abs=0.)
        assert data.shape == (1, 11, 11)

        data = self.pipeline.get_data('index1')
        assert np.sum(data) == pytest.approx(5, rel=self.limit, abs=0.)
        assert data.shape == (1, )

        attr = self.pipeline.get_attribute('selected1', 'PARANG', static=False)
        assert np.sum(attr) == pytest.approx(17., rel=self.limit, abs=0.)
        assert attr.shape == (5, )

        attr = self.pipeline.get_attribute('selected1', 'STAR_POSITION', static=False)
        assert np.sum(attr) == pytest.approx(50, rel=self.limit, abs=0.)
        assert attr.shape == (5, 2)

        attr = self.pipeline.get_attribute('removed1', 'PARANG', static=False)
        assert np.sum(attr) == pytest.approx(7., rel=self.limit, abs=0.)
        assert attr.shape == (1, )

        attr = self.pipeline.get_attribute('removed1', 'STAR_POSITION', static=False)
        assert np.sum(attr) == pytest.approx(10, rel=self.limit, abs=0.)
        assert attr.shape == (1, 2)

        module = FrameSelectionModule(name_in='select2',
                                      image_in_tag='start',
                                      selected_out_tag='selected2',
                                      removed_out_tag='removed2',
                                      index_out_tag='index2',
                                      method='max',
                                      threshold=1.,
                                      fwhm=0.1,
                                      aperture=('annulus', 0.05, 0.1),
                                      position=None)

        self.pipeline.add_module(module)
        self.pipeline.run_module('select2')

        data = self.pipeline.get_data('selected2')
        assert np.sum(data) == pytest.approx(21.42652724866543, rel=self.limit, abs=0.)
        assert data.shape == (2, 11, 11)

        data = self.pipeline.get_data('removed2')
        assert np.sum(data) == pytest.approx(43.016543226832646, rel=self.limit, abs=0.)
        assert data.shape == (4, 11, 11)

        data = self.pipeline.get_data('index2')
        assert np.sum(data) == pytest.approx(10, rel=self.limit, abs=0.)
        assert data.shape == (4, )

        attr = self.pipeline.get_attribute('selected2', 'PARANG', static=False)
        assert np.sum(attr) == pytest.approx(8., rel=self.limit, abs=0.)
        assert attr.shape == (2, )

        attr = self.pipeline.get_attribute('selected2', 'STAR_POSITION', static=False)
        assert np.sum(attr) == pytest.approx(20, rel=self.limit, abs=0.)
        assert attr.shape == (2, 2)

        attr = self.pipeline.get_attribute('removed2', 'PARANG', static=False)
        assert np.sum(attr) == pytest.approx(16., rel=self.limit, abs=0.)
        assert attr.shape == (4, )

        attr = self.pipeline.get_attribute('removed2', 'STAR_POSITION', static=False)
        assert np.sum(attr) == pytest.approx(40, rel=self.limit, abs=0.)
        assert attr.shape == (4, 2)

        module = FrameSelectionModule(name_in='select3',
                                      image_in_tag='start',
                                      selected_out_tag='selected3',
                                      removed_out_tag='removed3',
                                      index_out_tag='index3',
                                      method='range',
                                      threshold=(10., 10.7),
                                      fwhm=0.1,
                                      aperture=('circular', 0.1),
                                      position=None)

        self.pipeline.add_module(module)
        self.pipeline.run_module('select3')

        data = self.pipeline.get_data('selected3')
        assert np.sum(data) == pytest.approx(22.2568501695632, rel=self.limit, abs=0.)
        assert data.shape == (2, 11, 11)

        data = self.pipeline.get_data('removed3')
        assert np.sum(data) == pytest.approx(42.18622030593487, rel=self.limit, abs=0.)
        assert data.shape == (4, 11, 11)

        data = self.pipeline.get_data('index3')
        assert np.sum(data) == pytest.approx(12, rel=self.limit, abs=0.)
        assert data.shape == (4, )

    def test_image_statistics_full(self) -> None:

        module = ImageStatisticsModule(name_in='stat1',
                                       image_in_tag='read',
                                       stat_out_tag='stat1',
                                       position=None)

        self.pipeline.add_module(module)
        self.pipeline.run_module('stat1')

        data = self.pipeline.get_data('stat1')
        assert np.sum(data) == pytest.approx(115.68591492205017, rel=self.limit, abs=0.)
        assert data.shape == (10, 6)

    def test_image_statistics_position(self) -> None:

        module = ImageStatisticsModule(name_in='stat2',
                                       image_in_tag='read',
                                       stat_out_tag='stat2',
                                       position=(5, 5, 0.1))

        self.pipeline.add_module(module)
        self.pipeline.run_module('stat2')

        data = self.pipeline.get_data('stat2')
        assert np.sum(data) == pytest.approx(118.7138708968444, rel=self.limit, abs=0.)
        assert data.shape == (10, 6)

    def test_frame_similarity_mse(self) -> None:

        module = FrameSimilarityModule(name_in='simi1',
                                       image_tag='read',
                                       method='MSE',
                                       mask_radius=(0., 0.1))

        self.pipeline.add_module(module)
        self.pipeline.run_module('simi1')

        attr = self.pipeline.get_attribute('read', 'MSE', static=False)
        assert np.min(attr) > 0.
        assert np.sum(attr) == pytest.approx(0.11739141370277852, rel=self.limit, abs=0.)
        assert attr.shape == (10, )

    def test_frame_similarity_pcc(self) -> None:

        module = FrameSimilarityModule(name_in='simi2',
                                       image_tag='read',
                                       method='PCC',
                                       mask_radius=(0., 0.1))

        self.pipeline.add_module(module)
        self.pipeline.run_module('simi2')

        attr = self.pipeline.get_attribute('read', 'PCC', static=False)
        assert np.min(attr) > 0.
        assert np.sum(attr) == pytest.approx(9.134820985662829, rel=self.limit, abs=0.)
        assert attr.shape == (10, )

    def test_frame_similarity_ssim(self) -> None:

        module = FrameSimilarityModule(name_in='simi3',
                                       image_tag='read',
                                       method='SSIM',
                                       mask_radius=(0., 0.1),
                                       temporal_median='constant')

        self.pipeline.add_module(module)
        self.pipeline.run_module('simi3')

        attr = self.pipeline.get_attribute('read', 'SSIM', static=False)
        assert np.min(attr) > 0.
        assert np.sum(attr) == pytest.approx(9.074290801266256, rel=self.limit, abs=0.)
        assert attr.shape == (10, )

    def test_select_by_attribute(self) -> None:

        self.pipeline.set_attribute('read', 'INDEX', np.arange(44), static=False)

        module = SelectByAttributeModule(name_in='frame_removal_1',
                                         image_in_tag='read',
                                         attribute_tag='SSIM',
                                         number_frames=6,
                                         order='descending',
                                         selected_out_tag='select_sim',
                                         removed_out_tag='remove_sim')

        self.pipeline.add_module(module)
        self.pipeline.run_module('frame_removal_1')

        attr = self.pipeline.get_attribute('select_sim', 'INDEX', static=False)
        assert np.sum(attr) == pytest.approx(946, rel=self.limit, abs=0.)
        assert attr.shape == (44, )

        attr = self.pipeline.get_attribute('select_sim', 'SSIM', static=False)
        assert np.sum(attr) == pytest.approx(5.545578246610884, rel=self.limit, abs=0.)
        assert attr.shape == (6, )

        attr = self.pipeline.get_attribute('remove_sim', 'SSIM', static=False)
        assert np.sum(attr) == pytest.approx(3.528712554655373, rel=self.limit, abs=0.)
        assert attr.shape == (4, )

    def test_residual_selection(self) -> None:

        module = ResidualSelectionModule(name_in='residual_select',
                                         image_in_tag='start',
                                         selected_out_tag='res_selected',
                                         removed_out_tag='res_removed',
                                         percentage=80.,
                                         annulus_radii=(0.1, 0.2))

        self.pipeline.add_module(module)
        self.pipeline.run_module('residual_select')

        data = self.pipeline.get_data('res_selected')
        assert np.sum(data) == pytest.approx(41.77295229983322, rel=self.limit, abs=0.)
        assert data.shape == (4, 11, 11)

        data = self.pipeline.get_data('res_removed')
        assert np.sum(data) == pytest.approx(22.670118175664847, rel=self.limit, abs=0.)
        assert data.shape == (2, 11, 11)


from torch import nn

class Convolution(nn.Module):
    def __init__(self, in_channels, out_channels, k_size=3, rate=1):
        super().__init__()
        self.convlayer = nn.Sequential(
            nn.Conv2d(in_channels=int(in_channels),
                      out_channels=int(out_channels),
                      kernel_size=(k_size, k_size),
                      dilation=(rate, rate),
                      padding='same'),
            nn.BatchNorm2d(int(out_channels)),
            nn.ReLU(inplace=True))

    def forward(self, x):
        return self.convlayer(x)


#!/usr/bin/env python3
import os
import pickle
import warnings

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy.stats import norm
import ROOT
import uproot
import yaml

SPEED_OF_LIGHT = 2.99792458
SPLIT = True

centrality_colors = [ROOT.kOrange+7, ROOT.kAzure+4, ROOT.kTeal+4, ROOT.kBlack, ROOT.kBlack, ROOT.kBlack]

# avoid pandas warning
warnings.simplefilter(action='ignore', category=FutureWarning)
ROOT.gROOT.SetBatch()

##################################################################
# read configuration file
##################################################################
config = 'config.yaml'
with open(os.path.expandvars(config), 'r') as stream:
    try:
        params = yaml.full_load(stream)
    except yaml.YAMLError as exc:
        print(exc)

ANALYSIS_RESULTS_PATH = params['ANALYSIS_RESULTS_PATH']
CT_BINS_CENT = params['CT_BINS_CENT']
CENTRALITY_LIST = params['CENTRALITY_LIST']
RANDOM_STATE = params['RANDOM_STATE']
##################################################################

# split matter/antimatter
SPLIT_LIST = ['']
if SPLIT:
    SPLIT_LIST = ['antimatter', 'matter']

if not os.path.isdir("plots/cpt_and_lifetime"):
    os.mkdir("plots/cpt_and_lifetime")

ROOT.gStyle.SetOptStat(0)
ROOT.gStyle.SetOptFit(0)

eff_cut_dict = pickle.load(open("file_eff_cut_dict", "rb"))
presel_eff_file = uproot.open('PreselEff.root')
signal_extraction_file = ROOT.TFile.Open('SignalExtraction.root')
signal_extraction_keys = uproot.open('SignalExtraction.root').keys()
analysis_results_file = uproot.open(os.path.expandvars(ANALYSIS_RESULTS_PATH))

# get centrality selected histogram
cent_counts, cent_edges = analysis_results_file['Centrality_selected;1'].to_numpy()
cent_bin_centers = (cent_edges[:-1]+cent_edges[1:])/2

abs_correction_file = ROOT.TFile.Open('He3_abs.root')
eff_correction_file = ROOT.TFile.Open('EffAbsCorrection_1.root')
ratio_file = ROOT.TFile.Open('Ratio.root', 'recreate')

for i_cent_bins in range(len(CENTRALITY_LIST)):
    cent_bins = CENTRALITY_LIST[i_cent_bins]

    # get number of events
    cent_range_map = np.logical_and(cent_bin_centers > cent_bins[0], cent_bin_centers < cent_bins[1])
    counts_cent_range = cent_counts[cent_range_map]
    evts = np.sum(counts_cent_range)
    print(f'Number of events: {evts}')

    h_corrected_yields = [ROOT.TH1D(), ROOT.TH1D()]
    for i_split, split in enumerate(SPLIT_LIST):
        print(f'{i_split} -> {split}')
        # get preselection efficiency and abs correction histograms
        presel_eff_counts, presel_eff_edges = presel_eff_file[
            f'fPreselEff_vs_ct_{split}_{cent_bins[0]}_{cent_bins[1]};1'].to_numpy()
        presel_eff_bin_centers = (presel_eff_edges[1:]+presel_eff_edges[:-1])/2

        func_name = 'BGBW'
        if (cent_bins[1] == 90) or (cent_bins[0] == 0 and cent_bins[1] == 10):
            func_name = 'BlastWave'
        g_abs_correction = ROOT.TGraphAsymmErrors()
        g_abs_correction = abs_correction_file.Get(f"{cent_bins[0]}_{cent_bins[1]}/{func_name}/fEffCt_{split}_{cent_bins[0]}_{cent_bins[1]}_{func_name}")
        eff_abs_correction = ROOT.TH1D()
        eff_abs_correction = eff_correction_file.Get(f"{cent_bins[0]}_{cent_bins[1]}/{func_name}/fCorrection_{split}_{cent_bins[0]}_{cent_bins[1]}_{func_name}")

        # list of corrected yields
        ct_bins_tmp = [0]
        ct_bins_tmp += CT_BINS_CENT[i_cent_bins]
        if cent_bins[0]==30 or cent_bins[0]==50:
            ct_bins_tmp = [0, 2, 4, 7, 14, 35]
        if cent_bins[0]==10:
            ct_bins_tmp = [0, 2, 7, 14, 35]
        #if cent_bins[1] == 90:
        #    ct_bin_tmp = CT_BINS_CENT[i_cent_bins]s
        bins = np.array(ct_bins_tmp, dtype=float)
        # print(bins)
        h_corrected_yields[i_split] = ROOT.TH1D(

            f'fYields_{split}_{cent_bins[0]}_{cent_bins[1]}', f'{split}, {cent_bins[0]}-{cent_bins[1]}%', len(bins)-1, bins)

        for ct_bins in zip(CT_BINS_CENT[i_cent_bins][:-1], CT_BINS_CENT[i_cent_bins][1:]):

            bin = f'{split}_{cent_bins[0]}_{cent_bins[1]}_{ct_bins[0]}_{ct_bins[1]}'
            formatted_eff_cut = "{:.2f}".format(eff_cut_dict[bin])

            # look for plot with eff = eff_cut (or the nearest one)
            bkg_shape = 'pol1'
            eff_cut_increment = 0
            eff_cut_sign = -1
            while signal_extraction_keys.count(f'{bin}_{bkg_shape}/fInvMass_{formatted_eff_cut};1') == 0:
                if eff_cut_sign == -1:
                    eff_cut_increment += 0.01
                eff_cut_sign *= -1
                formatted_eff_cut = "{:.2f}".format(eff_cut_dict[bin]+eff_cut_increment*eff_cut_sign)

            # get signal
            h_raw_yield = signal_extraction_file.Get(f'{bin}_{bkg_shape}/fRawYields;1')
            eff_index = h_raw_yield.FindBin(float(formatted_eff_cut))
            raw_yield = h_raw_yield.GetBinContent(eff_index)
            raw_yield_error = h_raw_yield.GetBinError(eff_index)
            print(f'eff_cut = {formatted_eff_cut}, raw_yield = {raw_yield}+{raw_yield_error}')

            # apply corrections
            # 1. efficiency (presel x BDT)
            presel_eff_map = np.logical_and(
                presel_eff_bin_centers > ct_bins[0],
                presel_eff_bin_centers < ct_bins[1])
            presel_eff = presel_eff_counts[presel_eff_map]
            bdt_eff = float(formatted_eff_cut)
            print(f'bin: {presel_eff_map}, presel_eff: {presel_eff}')
            eff = presel_eff * eff_cut_dict[bin]
            # 2. absorption correction
            abs = g_abs_correction.GetPointY(CT_BINS_CENT[i_cent_bins].index(ct_bins[0]))
            print(f"absorption correction for point {CT_BINS_CENT[i_cent_bins].index(ct_bins[0])}: {abs}")
            # 3. efficiency correction
            eff_correct = eff_abs_correction.GetBinContent(eff_abs_correction.FindBin(ct_bins[0]))

            ct_bin_index = h_corrected_yields[i_split].FindBin(ct_bins[0]+0.5)
            h_corrected_yields[i_split].SetBinContent(ct_bin_index, raw_yield/eff[0]/abs/eff_correct)
            h_corrected_yields[i_split].SetBinError(ct_bin_index, raw_yield_error/eff[0]/abs/eff_correct)

        # set labels
        h_corrected_yields[i_split].GetXaxis().SetTitle("#it{c}t (cm)")
        h_corrected_yields[i_split].GetYaxis().SetTitle("d#it{N}/d(#it{c}t) (cm^{-1})")
        #h_corrected_yields[i_split].Scale(1, "width")
        for i_bin in range(len(bins))[2:]:
            bin_width = h_corrected_yields[i_split].GetBinWidth(i_bin)
            print(f"bin: {h_corrected_yields[i_split].GetBinLowEdge(i_bin)}; bin width: {bin_width}")
            bin_content = h_corrected_yields[i_split].GetBinContent(i_bin)
            bin_error = h_corrected_yields[i_split].GetBinError(i_bin)
            h_corrected_yields[i_split].SetBinContent(i_bin, bin_content/bin_width)
            h_corrected_yields[i_split].SetBinError(i_bin, bin_error/bin_width)
        h_corrected_yields[i_split].GetYaxis().SetRangeUser(1., 450.)
        h_corrected_yields[i_split].SetMarkerStyle(20)
        h_corrected_yields[i_split].SetMarkerSize(0.8)
        h_corrected_yields[i_split].Write()

        # fit with exponential pdf
        fit_function_expo = ROOT.TF1("expo", "expo", 2, 35)
        if cent_bins[0] == 30 or cent_bins[0]==50:
            fit_function_expo = ROOT.TF1("expo", "expo", 2, 14)
        elif cent_bins[1] == 90:
            fit_function_expo = ROOT.TF1("expo", "expo", 0, 35)
        res_expo = h_corrected_yields[i_split].Fit(fit_function_expo, "RMIS+")

        # compute and display lifetime
        tau = -1/fit_function_expo.GetParameter(1)*100/SPEED_OF_LIGHT # ps
        tau_error = fit_function_expo.GetParError(1)*100/SPEED_OF_LIGHT/fit_function_expo.GetParameter(1)/fit_function_expo.GetParameter(1) # ps
        tau_text = ROOT.TLatex(20, 100, '#tau = ' + "{:.0f}".format(tau) + '#pm' + "{:.0f}".format(tau_error) + ' ps')
        tau_text.SetTextSize(0.05)

        # display chi2
        formatted_chi2_lifetime = "{:.2f}".format(fit_function_expo.GetChisquare())
        chi2_lifetime_text = ROOT.TLatex(20, 70, "#chi^{2}/NDF = "+formatted_chi2_lifetime+"/"+str(fit_function_expo.GetNDF()))
        chi2_lifetime_text.SetTextSize(0.05)

        # compute yield
        integral = float()
        integral_error = float()
        integral = fit_function_expo.Integral(0, 1.e9)
        integral_error = fit_function_expo.IntegralError(0, 1.e9, res_expo.GetParams(), res_expo.GetCovarianceMatrix().GetMatrixArray())
        formatted_integral = "{:.10f}".format(integral/2./evts)
        formatted_integral_error = "{:.10f}".format(integral_error/evts/2.)
        print(f"N_ev = {evts}")
        integral_yield = ROOT.TLatex(20, 40, "1/#it{N_{ev}} #times BR #times d#it{N}/d#it{y} = "+formatted_integral+" #pm "+formatted_integral_error)
        integral_yield.SetTextSize(0.05)

        # draw on canvas
        canv = ROOT.TCanvas()
        ROOT.gStyle.SetOptStat(0)
        canv.SetTicks(1, 1)
        canv.SetName(h_corrected_yields[i_split].GetName())
        canv.SetTitle(h_corrected_yields[i_split].GetTitle())
        canv.cd()
        h_corrected_yields[i_split].Draw("")
        tau_text.Draw("same")
        chi2_lifetime_text.Draw("same")
        integral_yield.Draw("same")
        canv.SetLogy()
        canv.Write() # write to file
        canv.Print(f'plots/cpt_and_lifetime/{h_corrected_yields[i_split].GetName()}.pdf')

    # ratios
    ROOT.gStyle.SetOptStat(0)
    h_ratio = ROOT.TH1D(h_corrected_yields[0])
    h_ratio.SetName(f'fRatio_{cent_bins[0]}_{cent_bins[1]}')
    h_ratio.SetTitle(f'{cent_bins[0]}-{cent_bins[1]}%')
    h_ratio.Divide(h_corrected_yields[0], h_corrected_yields[1], 1, 1)
    h_ratio.SetLineColor(centrality_colors[i_cent_bins])
    h_ratio.SetMarkerColor(centrality_colors[i_cent_bins])
    h_ratio.GetYaxis().SetTitle("Ratio ^{3}_{#bar{#Lambda}}#bar{H} / ^{3}_{#Lambda}H")
    h_ratio.GetYaxis().SetRangeUser(0., 1.8)
    h_ratio.GetXaxis().SetRangeUser(0., 35.)
    h_ratio.SetMarkerStyle(20)
    h_ratio.SetMarkerSize(0.8)
    h_ratio.Fit("pol0","R","",2,35)
    h_ratio.Write()

    # plot ratios
    c = ROOT.TCanvas("c", "c")
    c.SetTicks(1, 1)
    c.cd()
    h_ratio.Draw()
    formatted_ratio = "{:.2f}".format(h_ratio.GetFunction("pol0").GetParameter(0))
    formatted_ratio_error = "{:.2f}".format(h_ratio.GetFunction("pol0").GetParError(0))
    text_x_position = 20
    ratio_text = ROOT.TLatex(text_x_position, 1.6, f"R = {formatted_ratio} #pm {formatted_ratio_error}")
    ratio_text.SetTextFont(44)
    ratio_text.SetTextSize(28)
    ratio_text.Draw("same")
    formatted_chi2 = "{:.2f}".format(h_ratio.GetFunction("pol0").GetChisquare())
    chi2_text = ROOT.TLatex(text_x_position, 1.45, "#chi^{2}/NDF = "+formatted_chi2+"/"+str(h_ratio.GetFunction("pol0").GetNDF()))
    chi2_text.SetTextFont(44)
    chi2_text.SetTextSize(28)
    chi2_text.Draw("same")
    c.Print(f"plots/{h_ratio.GetName()}.pdf")

    del h_corrected_yields
    del h_ratio

ratio_file.Close()


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

options = webdriver.ChromeOptions()
options.add_argument(r"user-data-dir=C:\Github\python-whatsapp-messages\whatsapp-web\data")
options.add_experimental_option("excludeSwitches", ["enable-automation"])
options.add_experimental_option('useAutomationExtension', False)
driver = webdriver.Chrome(options=options)
driver.maximize_window()
driver.get('https://web.whatsapp.com')

time.sleep(20)

driver.find_element(By.XPATH, "//*[@title='Eu']").click()

time.sleep(5)

message_field = driver.find_element(By.XPATH, "//div[@class='to2l77zo gfz4du6o ag5g9lrv bze30y65 kao4egtt']//p[@class='selectable-text copyable-text iq0m558w g0rxnol2']")
message_field.send_keys("Message to be sent!")

time.sleep(5)

message_field.send_keys(Keys.ENTER)

time.sleep(5)


__author__ = 'diegopinheiro'

from genetic_algorithms.genetic_algorithm import GeneticAlgorithm
from stop_criterions.iteration import IterationStopCriterion
from common.set_reader import SetReader
from selection_strategies.fitness_proportionate import FitnessProportionate
from selection_strategies.tournament_selection import TournamentSelection
from selection_strategies.rank_selection import RankSelection
from crossover_strategies.single_point import SinglePoint
from crossover_strategies.two_point import TwoPoint
from fitness_fuctions.accuracy_set import AccuracySet
from initialization_strategies.random_initialization import RandomInitialization
from initialization_strategies.data_initialization import DataInitialization
from mutation_strategies.one_gene_mutation import OneGeneMutation
import random
from common.writer import  Writer
from common.attribute import Attribute
from common.attribute_converter import AttributeConverter
from common.float_converter_2 import FloatConverter2

import sys

def main(output_path):

    print("\r\n\r\n\r\n")
    print("######## Experiment: testIris ########")

    data_set_name = "iris"
    data_set_reader = SetReader(data_set_name)
    training_set = data_set_reader.read_train_set()
    test_set = data_set_reader.read_test_set()
    data_set_attributes = data_set_reader.attributes

    population_size = 100

    replacement_rate = 0.3

    max_number_iteration = 3000

    use_elitism = True
    use_default = True

    initial_number_rules = 4
    min_number_rules = 4
    max_number_rules = 4

    stop_criterion = IterationStopCriterion(max_number_iteration=max_number_iteration)
    selection_strategy = FitnessProportionate()
    # selection_strategy = TournamentSelection(probability=0.7)
    # selection_strategy = RankSelection()
    # crossover_strategy = TwoPoint()
    crossover_strategy = SinglePoint()
    fitness_function = AccuracySet(min_number_rules=min_number_rules, max_number_rules=max_number_rules, dataset=training_set)
    initialization_strategy = DataInitialization(data_set=training_set)
    mutation_strategy = OneGeneMutation()

    # TODO remove fixed seed
    # random.seed(0)

    for rule in [2,3,4]:
        for mutation_rate in [.01, .05, .1, .2, .5]:
            print("rule:%d mutation:%.2f" % (rule, mutation_rate) )
            for run in range(0,30):
                initial_number_rules = rule

                genetic_algorithm = GeneticAlgorithm(population_size=population_size,
                                                     replacement_rate=replacement_rate,
                                                     mutation_rate=mutation_rate,
                                                     use_elitism=use_elitism,
                                                     use_default=use_default,
                                                     initial_number_rules=initial_number_rules,
                                                     stop_criterion=stop_criterion,
                                                     selection_strategy=selection_strategy,
                                                     crossover_strategy=crossover_strategy,
                                                     fitness_function=fitness_function,
                                                     initialization_strategy=initialization_strategy,
                                                     mutation_strategy=mutation_strategy,
                                                     attributes=data_set_attributes,

                                                     test_set=test_set)

                dimensions = genetic_algorithm.calculate_rule_size() * initial_number_rules
                writer = Writer(output_path+'ga_iris_%d_%d_%d_%s_%s_%.2f_%.2f_%d.with_positions' % (population_size, dimensions, initial_number_rules,
                                                                           str.lower(selection_strategy.__class__.__name__),
                                                                           str.lower(crossover_strategy.__class__.__name__),
                                                                                mutation_rate,
                                                                                replacement_rate,
                                                                                run))

                genetic_algorithm.writer = writer

                genetic_algorithm.training_set = training_set
                genetic_algorithm.learn()
                accuracy_training_set = genetic_algorithm.calculate_accuracy(training_set)
                accuracy_test_set = genetic_algorithm.calculate_accuracy(test_set)

                print(genetic_algorithm.get_best_individual().print_individual())
                print("Accuracy Training Set: " + str(accuracy_training_set))
                print("Accuracy Test Set: " + str(accuracy_test_set))

if __name__ == "__main__":
    output_path = sys.argv[1]

    main(output_path)


import struct
from newsroom import jsonl
import os
from utils import preprocess_text
import parameters as p
import pandas as pd
import pickle as pkl
from tensorflow.core.example import example_pb2
import argparse

def trim_and_transform(example_generator, new_filename, transformation, constraint):
    oldcount, newcount = 0, 0
    if os.path.isfile(new_filename):
        os.remove(new_filename)
    with jsonl.open(new_filename, gzip=True) as newfile:
        for line in example_generator:
            oldcount += 1
            line = transformation(line)
            if constraint(line):
                newcount += 1
                newfile.appendline(line)
            if oldcount % 1000 == 0:
                print(oldcount)
    print('# of old lines: %i, # of new lines: %i' % (oldcount, newcount))

def newsroom_constraint(line):
    return line['text'] is not None and line['summary'] is not None

def newsroom_preprocess(line):
    text = preprocess_text(line['text']) if line['text'] is not None else None
    summary = preprocess_text(line['summary']) if line['summary'] is not None else None
    return dict(text=text, summary=summary)

# # TODO: clean this up
# def cnn_preprocess(example_str):
#     abstract = []
#     article = []
#     in_abstract = False
#     in_article = False
#     example_str = example_str.decode('utf-8', 'replace')
#     example_str = example_str.replace('<s>', ' ')
#     example_str = example_str.replace('</s>', ' . ')
#     prev_c = None
#     for c in example_str.split():
#         if c == '.' and prev_c == c:
#             continue
#         if 'abstract' in c and c != 'abstract':
#             in_abstract = True
#             in_article = False
#             continue
#         if 'article' in c and c != 'article':
#             in_abstract = False
#             in_article = True
#             continue
#         c = c.replace('</s>', '.')
#         if '<s>' in c: continue
#         if '�' in c: continue
#         if in_abstract:
#             abstract.append(c)
#         if in_article:
#             article.append(c)
#         prev_c = c
#     pdb.set_trace()
#     return dict(text=article, summary=abstract)
def cnn_preprocess(example_str):
    # convert to tensorflow example e
    e = example_pb2.Example.FromString(example_str)
    # extract text and summary
    try:
        # the article text was saved under the key 'article' in the data files
        article_text = e.features.feature['article'].bytes_list.value[0].decode().split(' ')
        # the abstract text was saved under the key 'abstract' in the data files
        abstract_text = e.features.feature['abstract'].bytes_list.value[0].decode().split(' ')
    except ValueError:
        article_text = abstract_text = None
    return dict(text=article_text, summary=abstract_text)

def cnn_constraint(line):
    return line['text'] is not None and line['summary'] is not None

def pico_preprocess(line):
    line = dict(text=line.abstract, P=line.population, I=line.intervention, O=line.outcome)
    if pico_constraint(line):
        return {k:preprocess_text(v) for k,v in line.items()}
    else:
        return line

def pico_constraint(line):
	return line['text'] == line['text'] and \
	       line['P'] == line['P'] and \
	       line['I'] == line['I'] and \
	       line['O'] == line['O']

def preprocess_newsroom_datafile(filename, new_filename):
    with jsonl.open(filename, gzip=True) as oldfile:
        trim_and_transform(oldfile, new_filename, newsroom_preprocess, newsroom_constraint)

def preprocess_cnn_datafile(filename, new_filename):
    def cnn_dataset_generator():
        with open(filename, "rb") as file:
            while True:
                len_bytes = file.read(8)
                if not len_bytes: break # finished reading this file
                str_len = struct.unpack('q', len_bytes)[0]
                example_str = struct.unpack('%ds' % str_len, file.read(str_len))[0]
                yield example_str
    trim_and_transform(cnn_dataset_generator(), new_filename, cnn_preprocess, cnn_constraint)

def preprocess_pico_dataset(filename, new_filename_train, new_filename_dev, new_filename_test, aspect_file):
    df = pd.read_csv(filename)
    def train_generator():
        for i,row in df[:30000].iterrows():
            yield row
    def dev_generator():
        for i,row in df[30000:40000].iterrows():
            yield row
    def test_generator():
        for i,row in df[40000:].iterrows():
            yield row
    trim_and_transform(train_generator(), new_filename_train, pico_preprocess, pico_constraint)
    trim_and_transform(dev_generator(), new_filename_dev, pico_preprocess, pico_constraint)
    trim_and_transform(test_generator(), new_filename_test, pico_preprocess, pico_constraint)
    with open(aspect_file, 'w') as aspectfile:
        aspectfile.write(str(['P','I','O']))

def preprocess_all_newsroom_dataset_files(folder):
    preprocess_newsroom_datafile(
        os.path.join(folder, 'train.data'),
        os.path.join(folder, 'train_processed.data'))
    preprocess_newsroom_datafile(
        os.path.join(folder, 'val.data'),
        os.path.join(folder, 'val_processed.data'))
    preprocess_newsroom_datafile(
        os.path.join(folder, 'test.data'),
        os.path.join(folder, 'test_processed.data'))

def preprocess_all_cnn_dataset_files(folder):
    preprocess_cnn_datafile(
        os.path.join(folder, 'train.bin'),
        os.path.join(folder, 'train_processed.data'))
    preprocess_cnn_datafile(
        os.path.join(folder, 'val.bin'),
        os.path.join(folder, 'val_processed.data'))
    preprocess_cnn_datafile(
        os.path.join(folder, 'test.bin'),
        os.path.join(folder, 'test_processed.data'))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('dataset_type')
    parser.add_argument('data_folder')
    args = parser.parse_args()
    if args.dataset_type == 'newsroom':
        # for newsroom dataset
        preprocess_all_newsroom_dataset_files(args.data_folder)
    elif args.dataset_type == 'cnn':
        # for cnn dataset
        preprocess_all_cnn_dataset_files(args.data_folder)
    else:
        raise Exception

    # IGNORE FOR NOW
    # for pico dataset
    # aspect_file = '/Volumes/JEREDUSB/aspects.txt'
    # filename = '/Volumes/JEREDUSB/pico_cdsr.csv'
    # new_filename_train = '/Volumes/JEREDUSB/train_processed.data'
    # new_filename_dev = '/Volumes/JEREDUSB/dev_processed.data'
    # new_filename_test = '/Volumes/JEREDUSB/test_processed.data'
    # preprocess_pico_dataset(filename, new_filename_train, new_filename_dev, new_filename_test, aspect_file)
    # with open(aspect_file, 'w') as aspectfile:
    #     aspectfile.write(str(['P','I','O']))


from django import forms
from myapp1.models import OrderItem

class OrderItemForm(forms.ModelForm):
    class Meta:
        model = OrderItem
        fields = ['item', 'client', 'items_ordered']
        widgets = {
            'client': forms.RadioSelect(choices=[
                ('Client A', 'Client A'),
                ('Client B', 'Client B'),
                ('Client C', 'Client C'),
            ]),
        }
        labels = {
            'items_ordered': 'Quantity',
            'client': 'Client Name',
        }

    class InterestForm(forms.Form):
        interested = forms.TypedChoiceField(
            choices=((1, 'Yes'), (0, 'No')),
            coerce=lambda x: bool(int(x)),
            widget=forms.RadioSelect,
            label='Are you interested?'
        )
        quantity = forms.IntegerField(min_value=1, initial=1, label='Quantity')
        comments = forms.CharField(widget=forms.Textarea, required=False, label='Additional Comments')


import requests
from bs4 import BeautifulSoup

# セッションの生成
session = requests.session()

# ログイン
url = "http://somewhere/?controller=login&action=login"
data = { # from Form Data
    "login.x": "39", # <someValue>
    "login.y": "12", # <someValue>
    "accountName": "<ID>", # <ID>
    "password": "<PW>" # <PW>
}

response = session.post(url, data=data) # post方式
response.raise_for_status()
# print(response.text)

soupHome = BeautifulSoup(response.text, "html.parser")
name = soupHome.select_one("#profile-nav > div.menuBody > div.data > div.item.name > a")
print("name: ", name.string)

# 登録時間を取得
url = "http://somewhere/?controller=blog&action=top"
response = session.get(url)
response.raise_for_status()
# print(response.text)

soup = BeautifulSoup(response.text, "html.parser")
regTime = soup.select_one("div.regTime")
print("regTime: ", regTime.string)

# ログインが可能になれば、アクセスできる情報が増える
# ログインする方法を簡単にまとめると、
#  1. 開発者ツールを用いて、どこでログインの処理が行われているのかを確認する
#  2. その際、どこに（;php）何で（;方式）リクエストをするのかを確認する
#  3. pythonで上記(;1., 2.)を実装する
#    3-1. requestモヂュールでsessionを生成する
#    3-2. sessionにリクエストの方式を選択する
#    3-3. raise_for_status関数でログイン処理を行う


# coding=utf-8
import webapp2
import logging

from controllers.restore import ExportHandler, ExportHttpWorker, ExportWorker
from controllers.event import SaveEventsHandler, LastEventsHandler, UserEventsHandler, EventsCountHandler

class BlankPageHandler(webapp2.RequestHandler):
    def get(self):
        self.response.write('')

app = webapp2.WSGIApplication([(r'/', BlankPageHandler),
                               (r'/api/user/(\w+)', SaveEventsHandler),
                               (r'/count', EventsCountHandler),
                               (r'/get', LastEventsHandler),
                               (r'/get/(\w+)', UserEventsHandler),
                               (r'/export/worker/http', ExportHttpWorker),
                               (r'/export/worker', ExportWorker),
                               (r'/export', ExportHandler)], debug=True)

def handle_404(request, response, exception):
    logging.exception(exception)
    response.write('Oops! I could swear this page was here!')
    response.set_status(404)

def handle_500(request, response, exception):
    logging.exception(exception)
    response.write('success')
    response.set_status(200)

app.error_handlers[404] = handle_404
app.error_handlers[500] = handle_500


#!/usr/bin/env python
'''
A helper script to extract regions from LRW with dlib.
Mouth ROIs are fixed based on median mouth region across 29 frames.

@author Peratham Wiriyathammabhum
@date Jan 10, 2020
'''
import argparse
import os, os.path
import sys
import glob
import errno
import pickle
import math
import time
import copy
from multiprocessing import Pool
from time import time as timer

import numpy as np
import cv2
import yaml

cwd = os.getcwd()
from os.path import expanduser
hp = expanduser("~")
sys.path.insert(0, '/cfarhomes/peratham/lrcodebase')
from lr_config import *

from collections import ChainMap
import re

def get_stats(filename):
	stat_dict = {}
	vidname = filename.replace('.txt', '.mp4')

	# .... ex. 'Duration: 0.53 seconds' -> 0.53 float
	# stat_dict['duration'] = ''
	lastline = ''
	with open(filename,'r') as fp:
		lastline = list(fp)[-1]
	x = re.match('\w+: (\d+\.\d+) \w+', lastline)
	duration = float(x.group(1))
	stat_dict['duration'] = duration

	# ....
	# stat_dict['fps'] = ''
	cap = cv2.VideoCapture(vidname)
	fps = cap.get(cv2.CAP_PROP_FPS)
	stat_dict['fps'] = fps

	# ....
	# stat_dict['num_frames'] = ''
	stat_dict['num_frames'] = int(round(fps*duration))

	return {filename:stat_dict}

def process_boundary_stats(sample_paths, pool):
	try:
		batch_stats = pool.map(get_stats, sample_paths)
	except:
		print('[Error] {}'.format(file_paths[i]))
	return dict(ChainMap(*batch_stats))

def main(args):
	image_dir = args.dataset
	nthreads = int(args.nthreads)
	split = args.split
	filenames = glob.glob(os.path.join(image_dir, '*', '{}'.format(split), '*.txt'))
	filenames = sorted(filenames)
	total_size = len(filenames)
	pickle.dump( filenames, open( os.path.join(args.outdir, "lrw.{}.filenames.p".format(split)), "wb" ) )

	# ....
	res_dict = {} # result dict {filename:{duration:float.sec, fps:int1, num_frames:int2}}
	current_iter = 0
	chunk = 4*nthreads
	while current_iter < total_size:
		curr_batch_size = chunk if current_iter + chunk <= total_size else total_size - current_iter
		with Pool(nthreads) as pool:
			sample_paths = filenames[current_iter:current_iter+curr_batch_size]
			bdict = process_boundary_stats(sample_paths, pool)
			res_dict = {**res_dict, **bdict}
		current_iter += curr_batch_size
		if current_iter // chunk % 20 == 0:
			print('[Info] Operating...{}'.format(current_iter))
	# ....

	with open(args.outpickle,'wb') as fp:
		pickle.dump(res_dict, fp)

	with open(args.outfile,'w') as fp:
		yaml.dump(res_dict, fp)

	return

if __name__ == '__main__':
	parser = argparse.ArgumentParser(description='Pytorch Video-only BBC-LRW Example')
	parser.add_argument('--dataset', default='/cfarhomes/peratham/datapath/lrw/lipread_mp4',
		help='path to dataset')
	parser.add_argument('--split', default='train',
		help='train, val, test')
	parser.add_argument('--outdir', default='/cfarhomes/peratham/datapath/lrw/boundary_stats/',
		help='path to output files')
	parser.add_argument('--outfile', default='/cfarhomes/peratham/datapath/lrw/boundary_stats/boundary_stats.yaml',
		help='path to output yaml')
	parser.add_argument('--outpickle', default='/cfarhomes/peratham/datapath/lrw/boundary_stats/boundary_stats.p',
		help='path to output pickle')
	parser.add_argument('--nthreads', required=False, type=int,
		default=64, help='num threads')
	args = parser.parse_args()
	main(args)


from tqdm import tqdm
import torch
import torchvision
import torchvision.transforms as transforms
import pandas as pd
import numpy as np
from torch.utils.data import Dataset, DataLoader, Subset
from model import dCNN
from dataset import NumpySBDDataset
import torch.optim as optim
import os
import yaml
import argparse

device = "cuda" if torch.cuda.is_available() else "cpu"

def getdatasetstate(args={}):
    return {k: k for k in range(args["train_size"])}

def pre_process_fold(root, fold_name, class_to_label, meta, viz=False):
    path = os.path.join(root, fold_name)
    files = os.listdir(path)
    print("Processing fold:", fold_name)

    fold_data = []
    fold_labels = []

    for file in tqdm(files):
        try:
            if file.split(".")[1] != "wav":
                continue
            class_label = meta[meta["slice_file_name"] == file]["classID"].values[0]
            label_name = class_to_label[class_label]

            data, sr = librosa.load(os.path.join(path, file))

            S = librosa.feature.melspectrogram(data, sr, n_mels=128, fmax=8000)
            S_dB = librosa.power_to_db(S, ref=np.max)
            # pad S_dB with zeros to be (128, 173), pad the 1st dimension

            length = S_dB.shape[
                1
            ]  # the 4 second clip has length of 173 (trim anything with length greater than 173)

            if length > 173:
                # trim the length of the spectrogram to 173
                S_dB = S_dB[:173, :]
            else:
                # print(length)
                padding_to_add = 173 - length
                S_dB = np.pad(S_dB, [(0, 0), (0, padding_to_add)], mode="constant")

            #            print(S_dB.shape)
            if S_dB.shape == (128, 173):
                fold_data.append(S_dB)
                fold_labels.append(class_label)
            else:
                print(f"Size mismatch! {S_dB.shape}")

        except IndexError:
            print("Index error while processing file", file)
    return fold_data, fold_labels

def create_10_fold_data(ROOT="./data/"):
    train_file = "train_x.npy"
    full_path = os.path.join(ROOT, train_file)

    if os.path.isfile(full_path) and os.access(full_path, os.R_OK):
        train_x = np.load(ROOT + "train_x.npy")
        train_y = np.load(ROOT + "train_y.npy")
        test_x = np.load(ROOT + "test_x.npy")
        test_y = np.load(ROOT + "test_y.npy")

        # print(f"{train_x.shape} {train_y.shape} {test_x.shape} {test_y.shape}")

        return train_x, train_y, test_x, test_y

    train_x = []
    train_y = []

    test_x = []
    test_y = []

    print("Generating 10-fold data...")
    for fold in tqdm(folds):
        train = True

        if str(TEST_FOLD) in fold:
            train = False

        fold_data, fold_labels = pre_process_fold(
            "./data/audio/", fold, class_to_label, meta, viz=False
        )

        if train:
            train_x.extend(fold_data)
            train_y.extend(fold_labels)
        else:
            test_x.extend(fold_data)
            test_y.extend(fold_labels)

    ROOT_SAVE = "./data/"
    train_x = np.array(train_x)
    train_y = np.array(train_y)
    test_x = np.array(test_x)
    test_y = np.array(test_y)
    np.save(ROOT_SAVE + "train_x.npy", train_x)
    np.save(ROOT_SAVE + "train_y.npy", train_y)
    np.save(ROOT_SAVE + "test_x.npy", test_x)
    np.save(ROOT_SAVE + "test_y.npy", test_y)
    return train_x, train_y, test_x, test_y

def processData(args, stageFor="train", indices=None):
    class_to_label = {
        0: "air_conditioner",
        1: "car_horn",
        2: "children_playing",
        3: "dog_bark",
        4: "drilling",
        5: "engine_idling",
        6: "gun_shot",
        7: "jackhammer",
        8: "siren",
        9: "street_music",
    }

    meta = pd.read_csv("./data/UrbanSound8K.csv")

    folds = list(map(lambda x: "fold" + str(x), list(range(1, 11))))

    TEST_FOLD = 10  # TODO: make this an arg
    ROOT_DATA_DIR = "./data/audio/"  # TODO: make this an arg

    train_x, train_y, test_x, test_y = create_10_fold_data()

    if stageFor == "train":
        dataset = NumpySBDDataset(train_x, train_y)
        dataset = Subset(dataset, indices)
    elif stageFor == "infer":
        dataset = NumpySBDDataset(train_x, train_y)
        dataset = Subset(dataset, indices)
    else:
        dataset = NumpySBDDataset(train_x, train_y)

    loader = DataLoader(
        dataset, batch_size=64, shuffle=True
    )  # TODO: add arg for batch size

    return loader

def train(args, labeled, resume_from, ckpt_file):
    print("========== In the train step ==========")
    batch_size = args["batch_size"]
    lr = args["learning_rate"]
    momentum = args["momentum"]
    epochs = args["train_epochs"]
    train_split = args["split_train"]

    loader = processData(args, stageFor="train", indices=labeled)

    net = dCNN()
    net = net.to(device=device)

    criterion = torch.nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters())

    if resume_from is not None:
        ckpt = torch.load(os.path.join(args["EXPT_DIR"], resume_from + ".pth"))
        net.load_state_dict(ckpt["model"])
        optimizer.load_state_dict(ckpt["optimizer"])
    else:
        getdatasetstate(args)

    net.train()

    for epoch in tqdm(range(args["train_epochs"]), desc="Training"):

        running_loss = 0

        for i, batch in enumerate(loader, start=0):
            data, labels = batch

            data = data.to(device)
            labels = labels.to(device)

            optimizer.zero_grad()
            output = net(data)
            loss = criterion(output, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

            if i % 1000:
                print(
                    "epoch: {} batch: {} running-loss: {}".format(
                        epoch + 1, i + 1, running_loss / 1000
                    ),
                    end="\r",
                )
                running_loss = 0

    print("Finished Training. Saving the model as {}".format(ckpt_file))

    ckpt = {"model": net.state_dict(), "optimizer": optimizer.state_dict()}
    torch.save(ckpt, os.path.join(args["EXPT_DIR"], ckpt_file + ".pth"))

    return

def test(args, ckpt_file):
    print("========== In the test step ==========")
    batch_size = args["batch_size"]
    lr = args["learning_rate"]
    momentum = args["momentum"]
    epochs = args["train_epochs"]
    train_split = args["split_train"]

    loader = processData(args, stageFor="test")

    net = dCNN()
    net = net.to(device=device)

    net.load_state_dict(
        torch.load(os.path.join(args["EXPT_DIR"], ckpt_file + ".pth"))["model"]
    )

    net.eval()
    predix = 0
    predictions = {}
    truelabels = {}

    n_val = args["test_size"]
    with tqdm(total=n_val, desc="Testing round", unit="batch", leave=False) as pbar:
        for step, (batch_x, batch_y) in enumerate(loader):
            with torch.no_grad():
                batch_x = batch_x.to(device)
                batch_y = batch_y.to(device)

                prediction = net(batch_x)

            for logit, label in zip(prediction, batch_y):
                # predictions[predix] = logit.cpu().numpy().tolist()
                truelabels[predix] = label.cpu().numpy().tolist()

                class_probabilities = logit.cpu().numpy().tolist()
                index_max = np.argmax(class_probabilities)
                predictions[predix] = index_max

                predix += 1

            pbar.update()

    # unpack predictions
    predictions = [val for key, val in predictions.items()]
    truelabels = [val for key, val in truelabels.items()]

    return {"predictions": predictions, "labels": truelabels}

def infer(args, unlabeled, ckpt_file):
    print("========== In the inference step ==========")
    batch_size = args["batch_size"]
    lr = args["learning_rate"]
    momentum = args["momentum"]
    epochs = args["train_epochs"]
    train_split = args["split_train"]

    loader = processData(args, stageFor="infer", indices=unlabeled)

    net = dCNN()
    net = net.to(device=device)

    net.load_state_dict(
        torch.load(os.path.join(args["EXPT_DIR"], ckpt_file + ".pth"))["model"]
    )

    net.eval()

    n_val = len(unlabeled)
    predictions = {}
    predix = 0

    with tqdm(total=n_val, desc="Inference round", unit="batch", leave=False) as pbar:
        for step, (batch_x, batch_y) in enumerate(loader):
            with torch.no_grad():
                batch_x = batch_x.to(device)
                batch_y = batch_y.to(device)
                prediction = net(batch_x)

            for logit in prediction:
                predictions[unlabeled[predix]] = {}

                class_probabilities = logit.cpu().numpy().tolist()
                predictions[unlabeled[predix]]["pre_softmax"] = class_probabilities
                index_max = np.argmax(class_probabilities)
                predictions[unlabeled[predix]]["prediction"] = index_max
                predix += 1

            pbar.update()

    return {"outputs": predictions}

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config",
        default=os.path.join(os.getcwd(), "config.yaml"),
        type=str,
        help="Path to config.yaml",
    )
    args = parser.parse_args()

    with open(args.config, "r") as stream:
        args = yaml.safe_load(stream)

    labeled = list(range(5000))
    resume_from = None
    ckpt_file = "ckpt_0"

    print("Testing getdatasetstate")
    getdatasetstate(args=args)
    train(args=args, labeled=labeled, resume_from=resume_from, ckpt_file=ckpt_file)
    test(args=args, ckpt_file=ckpt_file)
    infer(args=args, unlabeled=[10, 20, 30], ckpt_file=ckpt_file)


import numpy as np
from collections import defaultdict
import rl.env
import rl.player
import rl.display

class Env(rl.env.DiscreteEnv):

    def __init__(self, n = 10, h = 1, p = 2, k = 3) -> None:
        super().__init__()

        self.n = n
        self.h = h
        self.p = p
        self.k = k

        # List of players
        self.players = list()
        self.state_size = 0

    # Return the next state and reward
    def step(self, state : tuple, actions : list[tuple], d = None) -> tuple[list[tuple],list[int]]:
        # Random uniformly distributed demand between 0 and n
        if d is None:
            demand = np.random.randint(0, self.n+1)
        else:
            demand = d

        # Holding cost
        cost = max(0,self.h*(state + actions - demand))

        # Lost sales penalty
        cost += max(0,self.p*(demand - state - actions[0]))

        # Fixed ordering cost
        cost += self.k * (actions[0] > 0)

        next_state = state + actions[0] - min(demand, state + actions[0])

        reward = -cost

        return int(next_state), int(reward)

    # Function that returns the available actions
    def A(self, state):
        return tuple([tuple([i]) for i in range(self.n - state + 1)])

    def goal_test(self, state) -> bool:
        raise NotImplementedError("ContinuousEnv.goal_test() is not implemented")

    def reset(self) -> None:
        return 0

    def init_render(self, grid_scale = 50):
        self.grid_scale = grid_scale

        self.width = (self.n+1)*self.grid_scale + self.grid_scale
        self.height = (self.n+1)*self.grid_scale

        self.display = rl.display.displayHandler(self.width, self.height)

    def render(self, agent):
        gs = self.grid_scale
        min_q = 1e10
        max_q = 0

        color_r = [(i, 0, 0) for i in np.linspace(255, 0, int(self.n/2 + 1))]
        color_g = [(0, i, 0) for i in np.linspace(0, 255, int(self.n/2 + 1))]
        color = color_r[:-1] + color_g

        min_q = min(agent.Q.values())
        max_q = max(agent.Q.values())

        for a in range(self.n+1):
            for s in range(self.n+1):
                if(a <= self.n-s):
                    q = agent.Q[s,(a,)]
                    color_index = int((q-min_q) * len(color)/(max_q - min_q + 1e-05))
                    color_index = min(color_index, self.n)
                    self.display.draw_square((a*gs + gs/2, s*gs + gs/2), (gs, gs), color[color_index])
                    self.display.draw_text(str(round(q, 2)), (a*gs + gs/2, s*gs + gs/2), (255,255,255), align="center")
                else:
                    self.display.draw_square((a*gs + gs/2, s*gs + gs/2), (gs, gs), (0,0,0))

        q_range = np.linspace(min_q, max_q, self.n+1)

        if self.n%2 == 0:
            n_ = self.n + 1
        else: n_ = self.n
        for s in range(n_):
            self.display.draw_square(((self.n+1)*gs + gs/2, s*gs + gs/2), (gs, gs), color[s], width = 10)
            self.display.draw_text(str(round(q_range[s], 2)), ((self.n+1)*gs + gs/2, s*gs + gs/2), (255,255,255), align="center")

        self.display.update()

        return self.display.eventHandler()

    def is_terminal(self, state) -> bool:
        return False

    def get_transistion_probabilities(self, state, action):

        p = defaultdict(lambda : 0)

        num_states = len(self.get_states())
        for d in range(num_states):

            next_state, reward = self.step(state, action, d)

            p[(int(next_state), int(reward))] += 1/num_states

        return p

    def get_states(self) -> list:
        return tuple(np.arange(self.n + 1))

class EnvNonUniform(Env):
    def __init__(self, n = 10, h = 1, p = 2, k = 3, b = 1, m = 0) -> None:
        super().__init__(n, h, p, k)

        self.b = b
        self.m = m

        self.p = p

    def step(self, state : tuple, actions : list[tuple], d = None) -> tuple[list[tuple],list[int]]:
        # Random uniformly distributed demand between 0 and n
        if d is None:
            demand_ratio = np.array([(self.b+1)/(self.n+1) \
                                    if x == self.m or x == self.m + 1 \
                                    else (self.n - 1 - 2*self.b)/(self.n**2 - 1) \
                                    for x in range(self.n+1)])

            demand = int(np.random.choice(np.arange(self.n+1),
                                    size = 1,
                                    p = demand_ratio)[0])
        else:
            demand = d

        # Holding cost
        cost = max(0,(self.h*(state + actions[0] - demand)))

        # Lost sales penalty
        cost += max(0,(self.p*(demand - state - actions[0])))

        # Fixed ordering cost
        if(actions[0] > 0): cost += self.k

        next_state = state + actions[0] - min(demand, state + actions[0])

        reward = cost

        return next_state, reward

    def get_transistion_probabilities(self, state, action):

        p = defaultdict(lambda : 0)

        num_states = len(self.get_states())
        for d in range(num_states):

            demand_ratio = np.array([(self.b+1)/(self.n+1) \
                                    if x == self.m or x == self.m + 1 \
                                    else (self.n - 1 - 2*self.b)/(self.n**2 - 1) \
                                    for x in range(self.n+1)])

            next_state, reward = self.step(state, action, d)

            p[(int(next_state), int(reward))] += demand_ratio[d]

        return p


import requests
from bs4 import BeautifulSoup
from gql import gql, Client
from gql.transport.requests import RequestsHTTPTransport

url = 'https://events.uchicago.edu/cal/main/showMain.rdo'
base_url = 'https://events.uchicago.edu'
gql_server = 'https://ucevents.herokuapp.com/'

client = Client(
  transport=RequestsHTTPTransport(url=gql_server, use_json=True),
  fetch_schema_from_transport=True,
)
mutation = gql('''
mutation {
  createEvent(
    name: "jeff test",
    description:"test",
    timeobj: {start: 1521930600, end: 1521940600 },
    tags: ["concert", "billie", "eilish"],
    event_photo_url: "test") {
    name
    description
    time{start, end}
    tags
  }
}
''')

query = gql("""
{
  events {
    name
  }
}
""")

#######################
#     DAY EVENTS
#######################

page = requests.get(url)
soup = BeautifulSoup(page.content, 'html.parser')

events = soup.find_all('table')[1].find_all('tbody')[0].find_all('tr')
parsed_events = []

for event in events:
  parsed = {}
  data = event.find_all('td')
  if len(data) > 1:
    parsed['time'] = " ".join(data[0].get_text().split())
    info = data[1].get_text('|')
    info = " ".join(info.split())
    parsed['name'] = info.split('|')[1]
    parsed['location'] = info.split('|')[3]
    parsed['link'] = data[1].a['href']
    parsed_events.append(parsed)

##############################################
#     GET INFO FROM SINGLE EVENT
##############################################

#   Returns a JSON of form { time, location, description, tags, name, link}

def get_event_info(event_link):
  event_page = requests.get(event_link)
  soup = BeautifulSoup(event_page.content, 'html.parser')
  data = soup.find('table').get_text('|')
  parsed = ' '.join(data.split())
  all_items = [i for i in parsed.split('|') if i.strip()]

  meta = ['When:','Where:','Description:','Cost:','Contact:','Tag:','Notes:']
  our_meta = ['time','location','description','','','tags','']
  parsed_event_info = {}

  for i, m in enumerate(meta):
    if not i == len(meta)-1:
      if meta[i] in all_items and meta[i+1] in all_items:
        start = all_items.index(meta[i])
        end = all_items.index(meta[i+1])
        repieced = ' '.join(all_items[start+1:end])
        if our_meta[i]:
          parsed_event_info[our_meta[i]] = repieced

  title_part = soup.find(id='maincontent')
  parsed_event_info['name'] = title_part.parent.find('h1').get_text()

  if 'tags' in parsed_event_info:
    parsed_event_info['tags'] = str(parsed_event_info['tags'].split(',')).replace("'", '"')

  parsed_event_info['link'] = event_link

  for key in parsed_event_info:
    if not key == 'tags':
      parsed_event_info[key] = parsed_event_info[key].replace('"', '')

  return parsed_event_info

for event in parsed_events[1:]:
  info = get_event_info( base_url + event['link'] )
  if not 'All Events' in info['name']:
    mutation = gql('''
    mutation {
      createEvent(
        name: "%s",
        description: "%s",
        location: "%s",
        timeobj: {start: 1521930600, end: 1521940600 },
        tags: %s,
        event_photo_url: "test") {
        name
        description
        time{start, end}
        tags
      }
    }
    '''  % ( info.get('name', ''), info.get('description', ''), info.get('location', ''), info.get('tags', '[\"test\"]') ) )
    print(client.execute(mutation))


"""OpenAI Functions Service"""
import json

import openai

from .schema import *

async def parse_openai_function(
    response: dict,
    functions: List[Type[F]] = OpenAIFunction.Metadata.subclasses,
    **kwargs: Any,
) -> FunctionCall:
    """Parses the response from OpenAI and returns a FunctionCall object."""
    choice = response["choices"][0]["message"]
    if "function_call" in choice:
        function_call_ = choice["function_call"]
        name = function_call_["name"]
        arguments = function_call_["arguments"]
        for i in functions:
            if i.__name__ == name:
                result = await i(**json.loads(arguments))(**kwargs)
                break
        else:
            raise ValueError(f"Function {name} not found")
        return result
    return FunctionCall(name="chat", data=choice["content"])

async def function_call(
    text: str,
    model: str = "gpt-3.5-turbo-16k-0613",
    functions: List[Type[F]] = OpenAIFunction.Metadata.subclasses,
    **kwargs,
) -> FunctionCall:
    """Calls the function orchestrator."""
    messages = [
        {"role": "user", "content": text},
        {"role": "system", "content": "You are a function Orchestrator"},
    ]
    response = await openai.ChatCompletion.acreate(
        model=model,
        messages=messages,
        functions=[i.openaischema for i in functions],
    )
    return await parse_openai_function(response, functions=functions, **kwargs)


from rest_framework import serializers
from .models import Profile, User
from articles.serializers import ArticleSerializer

class UserSerializer(serializers.ModelSerializer):
    class Meta:
        model = User
        fields = ('id', 'first_name', 'last_name')

class ProfileSerializer(serializers.ModelSerializer):
    user = UserSerializer()

    class Meta:
        model = Profile
        fields = '__all__'

    def update(self, instance, validated_data):
        user_data = validated_data.pop('user')
        instance.user.first_name = user_data.get('first_name')
        instance.user.last_name = user_data.get('last_name')
        instance.user.save()

        instance.city = validated_data.get('city')
        instance.bday = validated_data.get('bday')
        instance.gender = validated_data.get('gender')
        instance.save()
        return instance


import logging
import numpy as np

from ComperfModelling import predict_from_model

def run_kernel_density_estimation(values, normalised=False, normalisation_statistic=0):

	values = values.reshape(-1,1)

	# What should bandwidth be? Add as part of experiment configuration?
	bandwidth = 100000
	evaluation_gap = 50000

	if normalised:
		bandwidth = float(bandwidth)/normalisation_statistic
		evaluation_gap = float(evaluation_gap)/normalisation_statistic

	kde = KernelDensity(kernel='gaussian', bandwidth=bandwidth).fit(values)

	return kde, bandwidth, evaluation_gap

def compute_sample_weights_from_kde(kde, values, normalised=False, normalisation_statistic=0.0):

	# It is very expensive to evaluate every single example, so bin the examples into bins of 100k cycles, and apply the same weight to all examples within a given bin

	range_max = np.amax(values)
	range_min = np.amin(values)

	# Making the bin size equal to the bandwidth of the KDE
	bin_size = 100000
	if normalised:
		bin_size = bin_size/normalisation_statistic

	bins = int((float(range_max) - range_min)/float(bin_size))

	bin_positions = np.linspace(range_min,range_max,bins) + (bin_size/2.0) # i.e. evaluate the middle of each bin

	density_estimations = np.exp(kde.score_samples(bin_positions.reshape(-1,1)))

	# clip density estimations for bins that very likely don't contain any examples, as otherwise sample weights explode
	density_estimations[density_estimations < 0.1] = np.nan
	sample_weights = np.reciprocal(density_estimations)
	sample_weights[np.isnan(sample_weights)] = 10

	# now, apply the correct weight to each original estimation_position
	estimation_position_corresponding_bins = ((values - range_min) / bin_size).astype(np.int32)
	estimation_position_corresponding_bins[estimation_position_corresponding_bins >= bins] = bins-1

	per_example_sample_weights = sample_weights[estimation_position_corresponding_bins]

	logging.debug("Clipped %d examples to have a maximum sample weight of 10", len(per_example_sample_weights[per_example_sample_weights == 10]))

	return per_example_sample_weights

def kde_clustering_get_indexes(values, normalised=False, normalisation_statistic=0):

	counts = values.reshape(-1,1)

	kde, bandwidth, evaluation_gap = run_kernel_density_estimation(counts, normalised, normalisation_statistic)

	num_divisions = int(np.amax(counts)-np.amin(counts))/evaluation_gap
	if num_divisions < 20:
		# i.e. if we evaluate the density at less than this number of divisions, then we can't possibly cluster...
		return None, [], []

	estimation_positions = np.linspace(np.amin(counts),np.amax(counts),num_divisions)
	estimation_positions = estimation_positions.reshape(-1,1)

	estimations = np.exp(kde.score_samples(estimation_positions))

	# form a minimum by starting at 1.0 then descending to 0.0, then continuing with the histogram
	estimations = np.insert(estimations,0,0.0,axis=0)
	estimation_positions = np.insert(estimation_positions,0,np.amin(counts)-evaluation_gap,axis=0)
	estimations = np.insert(estimations,0,1.0,axis=0)
	estimation_positions = np.insert(estimation_positions,0,np.amin(counts)-(2*evaluation_gap),axis=0)

	mi, ma = argrelextrema(estimations, np.less_equal)[0], argrelextrema(estimations, np.greater)[0]

	bounds = []

	populated_cluster_id = 0

	cluster_indexes = []
	cluster_bounds = []

	for idx in np.arange(1,len(mi)):

		lower_bound = estimation_positions[mi[idx-1]]

		if idx == (len(mi)-1):
			upper_bound = np.amax(counts)+1
		else:
			upper_bound = estimation_positions[mi[idx]]

		clustered_task_indexes = np.where(np.logical_and(counts >= lower_bound, counts < upper_bound))[0]

		if(len(clustered_task_indexes) > 0):
			cluster_indexes.append(clustered_task_indexes)
			cluster_bounds.append([lower_bound,upper_bound])

	return kde, cluster_indexes, cluster_bounds

"""
	This function aims to identify significant differences between the responses, and saving the examples to later investigate as a comparative study
	A comparison may be between individual tasks or two task sets, where we necessarily average the latter
"""
def get_comparisons(responses):

	# should also save the cluster details so we can later use them to repeat the equivalent analysis on separate testing sets

	individual_comparisons = []
	cluster_comparisons = []
	cluster_comparison_bounds = []

	central_cluster_indexes = []
	central_cluster_bounds = []
	individual_comparisons = []
	cluster_comparisons = []
	kde = None

	# return a list of pairs, where each element is a list of task indexes to compare against eachother
	# return two types of comparisons: closest-to-mean task, tasks within 1*std-deviation of mean

	kde, clustered_task_indexes, cluster_bounds = kde_clustering_get_indexes(responses)
	if kde is None:
		log.info("Could not cluster instances as there is no significant variation.")
		return central_cluster_indexes, central_cluster_bounds, individual_comparisons, cluster_comparisons, kde

	# filter to only the significant clusters?
	large_clusters = []
	minimum_cluster_size = max([int(0.01*len(responses)),10])
	for cluster_idx, cluster_indexes in enumerate(clustered_task_indexes):
		if len(cluster_indexes) >= minimum_cluster_size:
			large_clusters.append(cluster_idx)

	logging.info("Found %d significant clusters (minimum cluster size was: %d)", len(large_clusters), minimum_cluster_size)

	clustered_task_indexes = [cluster_indexes for cluster_idx, cluster_indexes in enumerate(clustered_task_indexes) if cluster_idx in large_clusters]
	cluster_bounds = [cluster_indexes for cluster_idx, cluster_indexes in enumerate(cluster_bounds) if cluster_idx in large_clusters]

	if len(clustered_task_indexes) < 2:
		logging.info("There were fewer than two clusters, so cannot do a comparative analysis")
		return central_cluster_indexes, central_cluster_bounds, individual_comparisons, cluster_comparisons, kde

	fast_cluster = clustered_task_indexes[0]
	slow_clusters = clustered_task_indexes[1:]

	fast_durations = responses[fast_cluster]
	mean_fast_task_idx = (np.abs(fast_durations - np.mean(fast_durations))).argmin()
	mean_fast_task_idx_original_set = fast_cluster[mean_fast_task_idx]

	constraint_coefficient = 0.5
	fast_lower_central_bound = np.mean(fast_durations) - constraint_coefficient * np.std(fast_durations)
	fast_upper_central_bound = np.mean(fast_durations) + constraint_coefficient * np.std(fast_durations)
	central_fast_task_idxs = np.where(np.logical_and(fast_durations>=fast_lower_central_bound, fast_durations<=fast_upper_central_bound))[0]

	fast_durations_within_constraint = fast_durations[central_fast_task_idxs].flatten()
	distances_from_mean = np.absolute(fast_durations_within_constraint - np.mean(fast_durations))
	distances_sorted_indexes = np.argsort(distances_from_mean)

	# now filter the central indexes to only those that are within the bounds and are the closest to mean
	central_fast_task_idxs = central_fast_task_idxs[distances_sorted_indexes[:50]]

	# update the bounds
	fast_lower_central_bound = np.amin(fast_durations[central_fast_task_idxs])
	fast_upper_central_bound = np.amax(fast_durations[central_fast_task_idxs])

	central_fast_task_idxs_original_set = fast_cluster[central_fast_task_idxs]

	central_cluster_bounds = []
	central_cluster_bounds.append([fast_lower_central_bound, fast_upper_central_bound])

	central_cluster_indexes = []
	central_cluster_indexes.append(central_fast_task_idxs_original_set)

	for slow_cluster_idx, slow_cluster in enumerate(slow_clusters):

		slow_durations = responses[slow_cluster]
		mean_slow_task_idx = (np.abs(slow_durations - np.mean(slow_durations))).argmin()
		mean_slow_task_idx_original_set = slow_cluster[mean_slow_task_idx]
		lower_central_bound = np.mean(slow_durations) - constraint_coefficient * np.std(slow_durations)
		upper_central_bound = np.mean(slow_durations) + constraint_coefficient * np.std(slow_durations)
		central_slow_task_idxs = np.where(np.logical_and(slow_durations>=lower_central_bound, slow_durations<=upper_central_bound))[0]
		central_slow_task_idxs_original_set = slow_cluster[central_slow_task_idxs]

		slow_durations_within_constraint = slow_durations[central_slow_task_idxs].flatten()
		slow_distances_from_mean = np.absolute(slow_durations_within_constraint - np.mean(slow_durations))
		slow_distances_sorted_indexes = np.argsort(slow_distances_from_mean)

		# now filter the central indexes to only those that are within the bounds and are the closest to mean
		central_slow_task_idxs = central_slow_task_idxs[slow_distances_sorted_indexes[:50]]
		central_slow_task_idxs_original_set = slow_cluster[central_slow_task_idxs]

		# update the bounds
		lower_central_bound = np.amin(slow_durations[central_slow_task_idxs])
		upper_central_bound = np.amax(slow_durations[central_slow_task_idxs])

		individual_comparisons.append([mean_fast_task_idx_original_set, mean_slow_task_idx_original_set])
		central_cluster_bounds.append([lower_central_bound, upper_central_bound])
		central_cluster_indexes.append(central_slow_task_idxs_original_set)
		cluster_comparisons.append([0,slow_cluster_idx+1])

	return central_cluster_indexes, central_cluster_bounds, individual_comparisons, cluster_comparisons, kde

def run_pca_transformation(taskset, pca):

	taskset = list(taskset)
	taskset[1] = np.array([sum((taskset[1][i]*np.array(pca.components_[:])).transpose()) for i in range(len(taskset[1]))])
	taskset = tuple(taskset)

def run_sample_weighting_and_pca_transformation(
		training_taskset,
		validation_taskset,
		testing_taskset,
		should_standardise,
		pca_transform,
		training_event_stats):

	training_responses = training_taskset[2]
	validation_responses = validation_taskset[2]
	testing_responses = testing_taskset[2]

	kde_params = {}
	kde_params["normalised"] = should_standardise
	kde_params["normalisation_statistic"] = training_event_stats[-1][1]

	kde_params["values"] = training_responses
	kde, _, _ = run_kernel_density_estimation(**kde_params)

	kde_params["kde"] = kde
	training_sample_weights = compute_sample_weights_from_kde(**kde_params)

	kde_params["values"] = validation_responses
	validation_sample_weights = compute_sample_weights_from_kde(**kde_params)

	kde_params["values"] = testing_responses
	testing_sample_weights = compute_sample_weights_from_kde(**kde_params)

	num_input_events = len(training_taskset[1][0])

	if pca_transform == True:
		from sklearn.decomposition import PCA
		pca = PCA(n_components=num_input_events)
		pca.fit_transform(training_taskset[1])
	else:
		pca = pca_struct()
		pca.components_ = np.identity(num_input_events)

	run_pca_transformation(training_taskset, pca)
	run_pca_transformation(validation_taskset, pca)
	run_pca_transformation(testing_taskset, pca)

	return training_taskset, validation_taskset, testing_taskset, training_sample_weights, validation_sample_weights, testing_sample_weights, pca

def calculate_mae(models, testing_taskset):

	testing_examples = testing_taskset[1]
	testing_responses_true = testing_taskset[2][:,0]

	maes = []
	for model in models:
		testing_responses_pred = predict_from_model(model.model_file, testing_examples)
		mae = np.sum(np.absolute(testing_responses_pred - testing_responses_true)) / float(len(testing_responses_true))
		maes.append(mae)

	mean_mae = np.mean(maes)
	return mean_mae

def get_analysis_tasksets_for_configuration(
		dataset,
		all_events,
		pca_component_struct,
		taskset_kwargs
		):

	# Which configuration to get from the dataset is within the "benchmark_id" kwarg

	taskset, _, _ = dataset.get_taskset(**taskset_kwargs)
	run_pca_transformation(taskset, pca_component_struct)

	taskset_fs, _, _ = dataset.get_taskset(**taskset_kwargs) # No pca transformation here

	taskset_kwargs["input_events"] = all_events
	taskset_kwargs["standardised"] = False

	taskset_all_events_fs_destandardised, _, _ = dataset.get_taskset(**taskset_kwargs) # No pca transformation here

	return taskset, taskset_fs, taskset_all_events_fs_destandardised

def analyse_tasks_across_benchmarks(self,
		models,
		testing_tasksets, # array of length two, for two sets of tasks to compare to eachother
		scores,
		predicted_responses,
		std_devs,
		reference_responses,
		true_responses,
		reference_benchmark_idx, # i.e. what index of testing_tasksets is the reference
		reference_strategy="minimum_predicted_response",
		reference_indexes=[],
		just_get_values=False # If True, then deeplift won't be invoked at all, and all produced contributions/coefficients will be 0
		):

	target_benchmark_idx = [benchmark_idx for benchmark_idx in range(len(testing_tasksets)) if benchmark_idx != reference_benchmark_idx][0]

	if just_get_values:
		models = [0] # We don't need the models if we just want to get the value variations, but we do need an element in the list
	else:
		backend.set_floatx('float32')
		backend.set_learning_phase(False) # this is some keras funkery when using models with backend: https://github.com/fchollet/keras/issues/2310
		config = tf.ConfigProto()
		config.gpu_options.allow_growth = True
		backend.set_session(tf.Session(config=config))

	reference_example_idx = -1

	for model in models:

		if just_get_values:

			reference_example = np.zeros(len(testing_tasksets[reference_benchmark_idx][1][0])).tolist()
			reference_response = 0.0
			deeplift_results = np.array([[0.0 for event in reference_example] for task in testing_tasksets[target_benchmark_idx][1]])

			model_reference_predictions = np.array([0.0 for task in testing_tasksets[target_benchmark_idx][1]])
			model_target_predictions = np.array([0.0 for task in testing_tasksets[target_benchmark_idx][1]])

			if reference_responses == None:
				reference_responses = [reference_response]
			else:
				reference_responses.append(reference_response)

			reference_example_idx = 0
			reference_indexes.append(reference_example_idx)

		else:

			model_reference_predictions = predict_from_model(model.model_file, testing_tasksets[reference_benchmark_idx][1])
			model_target_predictions = predict_from_model(model.model_file, testing_tasksets[target_benchmark_idx][1])

			# If the comparison is between two groups of examples, which ones do we compare?
			# After this switch, we will have one reference example and one target example
			if reference_strategy == "minimum_predicted_response":
				reference_example_idx = np.argmin(model_reference_predictions)
				reference_example = testing_tasksets[reference_benchmark_idx][1][reference_example_idx]
				reference_response = model_reference_predictions[reference_example_idx]

			elif reference_strategy == "closest_to_mean_predicted_response":
				difference_from_mean = 9999999
				reference_example_idx = 0
				mean = np.mean(model_reference_predictions)
				for task_idx, duration in enumerate(model_reference_predictions):
					if abs(mean - duration) < difference_from_mean:
						difference_from_mean = abs(mean - duration)
						reference_example_idx = task_idx

				reference_example = testing_tasksets[reference_benchmark_idx][1][reference_example_idx]
				reference_response = model_reference_predictions[reference_example_idx]

			elif reference_strategy == "minimum_true_response":
				reference_example_idx = np.argmin(testing_tasksets[reference_benchmark_idx][2])
				reference_example = testing_tasksets[reference_benchmark_idx][1][reference_example_idx]
				reference_response = keras_model.predict(np.array([reference_example]))[:,0]

			elif reference_strategy == "minimum_example_inputs":
				reference_example = [np.amin(event_counts) for event_idx, event_counts in enumerate(testing_tasksets[reference_benchmark_idx][1].transpose())]
				reference_response = keras_model.predict(np.array([reference_example]))[:,0]

			elif reference_strategy == "zeros":
				reference_example = np.zeros(len(testing_tasksets[reference_benchmark_idx][1][0])).tolist()
				reference_response = keras_model.predict(np.array([reference_example]))[:,0]

			else:
				logging.error("Invalid reference strategy '%s'", reference_strategy)
				raise ValueError()

			if reference_responses == None:
				reference_responses = [reference_response]
			else:
				reference_responses.append(reference_response)

			del model_reference_predictions

			reference_indexes.append(reference_example_idx)

			target_example = testing_tasksets[target_benchmark_idx][1][0]

			# Now apply deeplift
			deeplift_results = run_deeplift_comparison(model.model_file,
				target_example,
				reference_example)

		# deeplift results is of format [task_idx][event_idx] = score
		# Aim to return a scores format of: [label_rank][input_event_idx] = list of scores, one per model

		if scores is None:
			# a list of scores for every task label and event
			scores = [[[] for event_idx in range(len(reference_example))] for i in list(range(int(np.amax(testing_tasksets[target_benchmark_idx][3]))))]
			# a list of predicted responses for every task label
			predicted_responses = [[] for i in list(range(int(np.amax(testing_tasksets[target_benchmark_idx][3]))))]
			# a list of input values for every task label and event
			std_devs = [[[] for event_idx in range(len(reference_example))] for i in list(range(int(np.amax(testing_tasksets[target_benchmark_idx][3]))))]

		add_true_responses = False
		if true_responses is None:
			add_true_responses = True
			true_responses = [[] for i in list(range(int(np.amax(testing_tasksets[target_benchmark_idx][3]))))]

		for task_idx, rank in enumerate(testing_tasksets[target_benchmark_idx][3]): # the true number of analysed tasks (we have scores for all)

			label_rank = int(rank)
			if label_rank >= len(scores):
				difference = (len(scores) - label_rank) + 1
				for i in range(difference + 1):
					scores.append([[] for event_idx in range(len(reference_example))])
					predicted_responses.append([])
					std_devs.append([[] for event_idx in range(len(reference_example))])
					if add_true_responses:
						true_responses.append([])

			for event_idx, score in enumerate(deeplift_results[task_idx]):
				scores[label_rank][event_idx].append(score)
				std_devs[label_rank][event_idx].append(testing_tasksets[target_benchmark_idx][1][task_idx][event_idx])

			predicted_responses[label_rank].append(model_target_predictions[task_idx])

			if add_true_responses:
				true_responses[label_rank].append(testing_tasksets[target_benchmark_idx][2][task_idx])

	return scores, predicted_responses, std_devs, reference_responses, true_responses

def get_delta_values(
		tasksets,
		reference_benchmark_idx,
		reference_index):

	target_benchmark_idx = [benchmark_idx for benchmark_idx in range(len(tasksets)) if benchmark_idx != reference_benchmark_idx][0]

	reference_example = tasksets[reference_benchmark_idx][1][reference_index]

	values = [[[] for event_idx in range(len(reference_example))] for i in list(range(int(np.amax(tasksets[target_benchmark_idx][3]))))]

	for task_idx, rank in enumerate(tasksets[target_benchmark_idx][3]):

		label_rank = int(rank)
		if label_rank >= len(values):
			difference = (len(values) - label_rank) + 1
			for i in range(difference + 1):
				values.append([[] for event_idx in range(len(reference_example))])

		for event_idx in range(len(tasksets[target_benchmark_idx][1][task_idx])):
			values[label_rank][event_idx].append(tasksets[target_benchmark_idx][1][task_idx][event_idx] - reference_example[event_idx])

	return values

def performance_comparison(
		reference_taskset,
		reference_taskset_fs,
		reference_taskset_all_events_fs_destandardised,
		target_taskset,
		target_taskset_fs,
		target_taskset_all_events_fs_destandardised,
		models,
		training_event_stats,
		training_pca_matrix,
		linreg_coefficients_fs_des
		):

	# The performance comparison may operate as a 1-1 comparison, or a many-many comparison
	# For the latter case, we average across multiple individual comparisons (all-all)

	linreg_contributions_across_comparisons = []
	contributions_across_comparisons = []
	coefficients_across_comparisons = []
	event_variations_across_comparisons = []
	all_event_variations_across_comparisons = []
	pred_duration_variations_across_comparisons = []
	true_duration_variations_across_comparisons = []

	logging.debug("For the performance comparison, there are %d reference examples and %d target examples.", len(reference_taskset[1]))

	reference_instance_indexes_to_use_for_comparison = list(range(len(reference_taskset[1])))[:]

	# At the moment, all reference instances are compared to all target instances
	"""
	num_reference_instances_to_use = 4 # TODO parameterise this arbitrary number!
	if len(reference_instance_indexes_to_use_for_comparison) > num_reference_instances_to_use:
		selection = np.random.choice(np.arange(len(reference_instance_indexes_to_use_for_comparison)),num_reference_instances_to_use,replace=False)
		reference_instance_indexes_to_use_for_comparison = np.array(reference_instance_indexes_to_use_for_comparison)[selection].tolist()
	"""

	# For each reference example, compare it with the corresponding target example
	for reference_instance_idx in reference_instance_indexes_to_use_for_comparison:

		# Build a pseudo taskset containing the one reference example that we are comparing
		reference_taskset_temp = (reference_taskset[0], reference_taskset[1][[reference_instance_idx]],reference_taskset[2][[reference_instance_idx]],reference_taskset[3][[reference_instance_idx]])

		# Build a pseudo taskset containing the one target example that we are comparing
		target_instance_idx = np.where(target_taskset[3] == reference_taskset_temp[3][0])[0] # i.e. find the matching label
		target_taskset_temp = (target_taskset[0], target_taskset[1][[target_instance_idx]],target_taskset[2][[target_instance_idx]],target_taskset[3][[target_instance_idx]])

		# reference and target tasksets in FS
		reference_instance_idx_fs = np.where(reference_taskset_fs[3] == reference_taskset_temp[3][0])[0]
		target_instance_idx_fs = np.where(target_taskset_fs[3] == reference_taskset_temp[3][0])[0]
		reference_instance_idx_fs_all_events = np.where(reference_taskset_all_events_fs_destandardised[3] == reference_taskset_temp[3][0])[0]
		target_instance_idx_fs_all_events = np.where(target_taskset_all_events_fs_destandardised[3] == reference_taskset_temp[3][0])[0]
		target_taskset_fs_temp = (target_taskset_fs[0], target_taskset_fs[1][[target_instance_idx_fs]],target_taskset_fs[2][[target_instance_idx_fs]],target_taskset_fs[3][[target_instance_idx_fs]])
		reference_taskset_fs_temp = (reference_taskset_fs[0], reference_taskset_fs[1][[reference_instance_idx_fs]],reference_taskset_fs[2][[reference_instance_idx_fs]],reference_taskset_fs[3][[reference_instance_idx_fs]])

		comparison_tasksets = [reference_taskset_temp,target_taskset_temp]
		comparison_tasksets_fs = [reference_taskset_fs_temp,target_taskset_fs_temp]

		reference_indexes = []
		task_scores = None
		task_std_devs = None
		predicted_durations = None
		true_durations = None
		reference_response = None

		# TODO this function (and thus the surrounding code) is in need of significant refactoring!
		scores, predicted_durations_tmp, task_std_devs_tmp, reference_responses_tmp, true_durations_tmp = analyse_tasks_across_benchmarks(models,
			comparison_tasksets,
			task_scores,
			predicted_durations,
			task_std_devs,
			reference_response,
			true_durations,
			reference_benchmark_idx=0,
			reference_strategy="closest_to_mean_predicted_response",
			reference_indexes=reference_indexes)

		event_variations_fs = get_delta_values(comparison_tasksets_fs,
			reference_benchmark_idx=0,
			reference_index=reference_indexes[0])

		reference_example = reference_taskset_temp[1][reference_indexes[0]]
		reference_duration_true = reference_taskset_temp[2][reference_indexes[0]]
		reference_duration_pred = reference_responses_tmp[0]

		reference_example_all_events = reference_taskset_all_events_fs_destandardised[1][reference_instance_idx_fs_all_events]
		target_example_all_events = target_taskset_all_events_fs_destandardised[1][target_instance_idx_fs_all_events]

		duration_variation_true = np.array([np.mean(the_durations) for task_idx, the_durations in enumerate(true_durations_tmp) if len(the_durations) > 0]) - reference_duration_true
		duration_variation_destandardised_true = (duration_variation_true * training_event_stats[-1][1])

		duration_variation_pred = np.array([np.mean(the_durations) for task_idx, the_durations in enumerate(predicted_durations_tmp) if len(the_durations) > 0]) - reference_duration_pred
		duration_variation_destandardised_pred = (duration_variation_pred * training_event_stats[-1][1])

		event_variations_fs = np.array([np.mean(delta_values,axis=1) for task_idx, delta_values in enumerate(event_variations_fs) if any(len(s) > 0 for s in delta_values)])
		event_variations_fs_destandardized = event_variations_fs * np.array(training_event_stats).transpose()[1].transpose()[:-1]
		event_variations = np.array([np.mean(the_inputs,axis=1) for task_idx, the_inputs in enumerate(task_std_devs_tmp) if any(len(s) > 0 for s in the_inputs)]) - reference_example

		all_event_variations = np.array([target_value - reference_value for reference_value,target_value in zip(reference_example_all_events,target_example_all_events)])

		scores_pc = np.array([np.mean(the_scores,axis=1) for task_idx, the_scores in enumerate(scores) if any(len(s) > 0 for s in the_scores)])
		coefficients_pc = np.divide(scores_pc, event_variations, out=np.zeros_like(scores_pc), where=event_variations!=0)
		coefficients_fs = np.matmul(coefficients_pc, training_pca_matrix.components_)
		coefficients_fs_destandardized = ((coefficients_fs * np.array(training_event_stats)[-1][1])) / np.array(training_event_stats).transpose()[1].transpose()[:-1]
		scores_fs = coefficients_fs * event_variations_fs
		scores_fs_destandardized = scores_fs * training_event_stats[-1][1]

		# Get the average over the compared examples (though there should only be one comparison here)
		average_scores = scores_fs_destandardized[0]
		average_coefficients = coefficients_fs_destandardized[0]
		average_values = event_variations_fs_destandardized[0]

		linreg_contributions = np.array(average_values) * linreg_coefficients_fs_des
		linreg_contributions_across_comparisons.append(linreg_contributions)

		if configuration_idx == reference_benchmark_id:
			# Record the actual reference values so we can later reconstruct absolute counts from deltas if we want
			average_values = np.matmul(reference_example, training_pca_matrix.components_) * np.array(training_event_stats).transpose()[1].transpose()[:-1] + np.array(training_event_stats).transpose()[0][:-1]
			predicted_durations_destandardised = reference_duration_pred * np.array(training_event_stats[-1][1]) + training_event_stats[-1][0]
			duration_variation_destandardised_pred = predicted_durations_destandardised
			durations_destandardised = reference_duration_true * np.array(training_event_stats[-1][1]) + training_event_stats[-1][0]
			duration_variation_destandardised_true = durations_destandardised

		all_event_variations_across_comparisons.append(all_event_variations)
		contributions_across_comparisons.append(average_scores)
		coefficients_across_comparisons.append(average_coefficients)
		event_variations_across_comparisons.append(average_values)
		pred_duration_variations_across_comparisons.append(np.mean(duration_variation_destandardised_pred))
		true_duration_variations_across_comparisons.append(np.mean(duration_variation_destandardised_true))

	# And I need to get value variations for *all* events
	average_all_event_variations = np.mean(np.array(all_event_variations_across_comparisons),axis=0)[0]

	# Now I have the computed data for all the comparisons between the reference set and target set
	average_linreg_contributions = np.mean(np.array(linreg_contributions_across_comparisons),axis=1)

	average_contributions = np.mean(np.array(contributions_across_comparisons),axis=0)
	average_coefficients = np.mean(np.array(coefficients_across_comparisons),axis=0)
	average_event_variations = np.mean(np.array(event_variations_across_comparisons),axis=0)
	average_predicted_duration_variations = np.mean(np.array(pred_duration_variations_across_comparisons),axis=0)
	average_true_duration_variations = np.mean(np.array(true_duration_variations_across_comparisons),axis=0)

	return average_contributions, average_linreg_contributions, average_coefficients, average_event_variations, average_all_event_variations, np.mean(average_predicted_duration_variations), np.mean(average_true_duration_variations)