# Simple quantification of features

import os
import click
from subprocess import PIPE, Popen, call
from pathlib import Path
from os.path import isfile, isdir
from pprint import pprint
from multiprocessing import Lock, Process, Queue, current_process, Pool
# from tqdm.contrib.concurrent import process_map
from tqdm import tqdm

from .generics import *
from .cli import cli

def call_count(job):

	name, locus = job

	start = int(locus.split(":")[-1].split("-")[0])
	stop  = int(locus.split(":")[-1].split("-")[1])

	c = Counter()

	sam_iter = samtools_view(alignment_file, locus=locus)

	for read in sam_iter:
		sam_strand, sam_length, _, sam_pos, _, sam_rg, sam_seq = read

		if sam_pos >= start and sam_pos + sam_length <= stop:

			c.update([sam_rg])

	line = [name, locus, sum(c.values())]
	for rg in rgs:
		line.append(c[rg])

	# pprint(c)

	lock.acquire()
	# print(".", end='', flush=True)

	with open(output_file, 'a') as outf:
		print("\t".join(map(str, line)), file=outf)

	# pbar.update(1)
	lock.release()

def init(l, r, a, o, ):
	global lock
	global rgs
	global alignment_file
	global output_file
	lock = l
	rgs = r
	alignment_file = a
	output_file = o

@cli.command(group='Calculation', help_priority=3)

@click.option("-a", "--alignment_file",
	required=False,
	type=click.UNPROCESSED, callback=validate_path,
	help='Alignment file input (bam or cram).')

@click.option("-o", "--output_directory",
	required=False,
	type=click.Path(),
	help="Directory name for annotation output")

@click.option("-l", "--locus_files",
	required=False,
	type=click.Path(exists=True),
	multiple=True,
	help='list of loci in 2-column format (name, chrom:start-stop). Multiples allowed. Defaults to the results file.')

@click.option("-g", "--gff_files",
	required=False,
	type=click.Path(exists=True),
	multiple=True,
	help='gff of genes to be evaluated. Looks for "mRNA" features as would be described in an NCBI annotation. Multiples allowed. Defaults to None.')

@click.option("--ignore_zeroes",
	is_flag=True,
	help="Don't save zero depth entries to deepcounts.txt. Saves file space, but can make it harder to analyze.")

def count(** params):
	"""Gets counts for all readgroups, loci, strand, and sizes."""

	rc = requirementClass()
	rc.add_samtools()
	rc.check()

	ic = inputClass(params)
	ic.check(['alignment_file'])

	output_directory     = str(ic.output_directory)
	alignment_file       = ic.inputs['alignment_file']

	locus_files          = params['locus_files']
	gff_files            = params['gff_files']
	ignore_zeroes        = params['ignore_zeroes']

	Path(output_directory+ "/counts/").mkdir(parents=True, exist_ok=True)

	counts_file = f"{output_directory}/counts/counts.txt"
	deep_counts_file = f"{output_directory}/counts/deepcounts.txt"

	if len(locus_files) == 0:
		locus_files = [f"{output_directory}/peak/loci.txt"]

	c = Counter()

	print('annotations:')
	loci = []
	for file in locus_files:
		with open(file, 'r') as f:
			header = f.readline()
			for line in f:
				c[file] += 1
				line = line.strip().split("\t")[:2]
				line[1] = line[1].replace("..", "-")
				loci.append((line[1], line[0]))

	for file in gff_files:
		with open(file, 'r') as f:
			for line in f:
				if not line.startswith("#"):
					line = line.strip().split("\t")
					if line[2] == 'gene':
						c[file] += 1
						coords = f"{line[0]}:{line[3]}-{line[4]}"
						name = line[8].split(";")[0].split("=")[1]

						loci.append((name, coords))

	for key, val in c.items():
		print(" ", key, "->", val, 'loci')

	print('')
	print('processing annotations...')

	chroms, rgs = get_chromosomes(alignment_file)

	# print('  -> making empty genome structure')
	# lookup = {}
	# for c,l in chroms:
	# 	lookup[c] = [None] * l
	# 	# lookup[c] = {}

	# print('  -> populating')
	# pbar = tqdm(total = len(loci), ncols=90)
	# for name, locus in loci:
	# 	pbar.update()
	# 	# print(locus)

	# 	chrom, start, stop = parse_locus(locus)
	# 	# print(chrom, start, stop)

	# 	for r in range(start, stop + 1):
	# 		try:
	# 			lookup[chrom][r].add(name)
	# 		except:
	# 			lookup[chrom][r] = set([name])
	# pbar.close()

	# depth_c = get_rg_depth(output_directory, alignment_file)

	# print(sum(depth_c.values()))
	# pprint(depth_c)
	# sys.exit()

	with open(counts_file, 'w') as outf:
		print('name','locus', "\t".join(rgs), sep='\t', file=outf)

	with open(deep_counts_file, 'w') as outf:
		print('name','rg','length','strand','count', sep='\t', file=outf)

	perc = percentageClass(1,len(loci))

	for i, locus in enumerate(loci):
		name, locus = locus

		c = Counter()
		deep_c = Counter()

		p_count = perc.get_percent(i)
		if p_count:
			print(f"   quantifying... {p_count}%", end='\r')

		chrom, start, stop = parse_locus(locus)

		for read in samtools_view(alignment_file, locus=locus):
			sam_strand, sam_length, sam_size, sam_pos, sam_chrom, sam_rg, sam_seq, sam_read_id = read

			c[sam_rg] += 1
			deep_c[(sam_rg, sam_length, sam_strand)] += 1

		line = [name, locus]

		line += [c[rg] for rg in rgs]

		with open(counts_file, 'a') as outf:

			print("\t".join(map(str, line)), file=outf)

		with open(deep_counts_file, 'a') as outf:
			for rg in rgs:
				for length in range(15,31):
					for strand in {"+", "-"}:
						count = deep_c[(rg, length, strand)]

						if count == 0 and ignore_zeroes:
							pass
						else:
							print(name, rg, length, strand, count, sep='\t', file=outf)

	print()

	# print()
	# print("processing alignments...")
	# print(f"  {sum(depth_c.values()):,} total")

	# pbar = tqdm(total = sum(depth_c.values()), ncols=90)

	# sam_iter = samtools_view(alignment_file)

	# for i,read in enumerate(sam_iter):

	# 	# if i % 1000000 == 0:
	# 	# 	print(read)
	# 	# 	break

	# 	sam_strand, sam_length, _, sam_pos, sam_chrom, sam_rg, sam_seq = read

	# 	if 15 <= sam_length <= 30:
	# 		pbar.update()

	# 		try:
	# 			start_loc = lookup[sam_chrom][sam_pos]
	# 		except IndexError:
	# 			start_loc = set([None])

	# 		try:
	# 			stop_loc  = lookup[sam_chrom][sam_pos + sam_length + 1]
	# 		except IndexError:
	# 			stop_loc = set([None])

	# 		# print(start_loc, stop_loc)

	# 		if start_loc and stop_loc:
	# 			names = list(start_loc & stop_loc)
	# 			names = [n for n in names if n]

	# 			# if len(names) == 0:
	# 			# 	names = ['NoLocus']

	# 		else:
	# 			# names = ['NoLocus']
	# 			names = []

	# 		for name in names:
	# 			c.update([(name, sam_rg)])
	# 			deep_c.update([(name, sam_rg, sam_length, sam_strand)])

		# if sum(c.values()) > 1000000:
		# 	pprint(c)
		# 	sys.exit()

	# pbar.close()

	# pprint(deep_c)
	# sys.exit()

	# print()
	# print('writing...')

	# with open(counts_file, 'w') as outf:
	# 	print('name','locus', "\t".join(rgs), sep='\t', file=outf)

	# with open(deep_counts_file, 'w') as outf:
	# 	print('name','rg','length','strand','count', sep='\t', file=outf)

	# loci.append(("NoLocus", "NA"))

	# pbar = tqdm(total = len(loci), ncols=90)
	# for name, locus in loci:
	# 	pbar.update()

	# 	line = [name, locus]

	# 	line += [c[(name, rg)] for rg in rgs]

	# 	with open(counts_file, 'a') as outf:
	# 		print("\t".join(map(str, line)), file=outf)

	# 	for rg in rgs:

	# 		for length in range(15,31):
	# 			for strand in {"+", "-"}:
	# 				count = deep_c[(name, rg, length, strand)]

	# 				with open(deep_counts_file, 'a') as outf:
	# 					print(name, rg, length, strand, count, sep='\t', file=outf)
	# 				# if count > 0:
	# 				# 	sys.exit()

	# pbar.close()
		# sys.exit()

	# with open(output_file, 'w') as outf:
	# 	header = "name\tlocus\ttotal"
	# 	for r in rgs:
	# 		header += f"\t{r}"
	# 	print(header, file=outf)

	# lock = Lock()

	# with Pool(initializer=init,
	# 	initargs=(lock, rgs, alignment_file, output_file, ),
	# 	processes=threads) as p:

	# 	r = list(tqdm(p.imap(call_count, loci), total=len(loci)))


import time
import logging

import os
import signal

import argparse

class SingingTwice:
    SONG_LIST = ["CHEER UP", "FANCY", "FEEL SPECIAL"]

    def __init__(self, log_file=None):
        logging.basicConfig(level=logging.INFO, format='%(message)s')
        self.logger = logging.getLogger("SingingTwice")
        self.log_file = log_file

        if log_file:
            self.log_handler = logging.FileHandler(self.log_file)
            self.logger.addHandler(self.log_handler)

        self.__stop = False

        signal.signal(signal.SIGINT, self.stop)
        signal.signal(signal.SIGTERM, self.stop)

    def main(self):
        i = 0
        self.logger.info("Start Singing, PID {0}".format(os.getpid()))
        while not self.__stop:
            self.logger.info(self.SONG_LIST[i % len(self.SONG_LIST)])
            i += 1
            time.sleep(1)

    def stop(self, signum, frame):
        self.__stop = True
        self.logger.info("Receive Signal {0}".format(signal))
        self.logger.info("Stop Singing")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--log", help="log filename", default=None)
    args = parser.parse_args()

singing_Twice = SingingTwice()
singing_Twice.main()


from email import header
import json
import pandas as pd
import gspread
from dotenv import load_dotenv
from gspread_dataframe import set_with_dataframe
import os
import json

load_dotenv()

def download_gspread():
    credentials = json.loads(os.environ.get("GCLOUD_CREDENTIALS"))
    gc = gspread.service_account_from_dict(credentials)
    sh = gc.open_by_key(os.environ.get("SPREADSHEET_KEY"))
    worksheet = sh.get_worksheet(0)
    df = pd.DataFrame(worksheet.get_all_values())
    df = df.drop(df.index[0])
    df.columns = ["idOp", "rate", "email"]
    df["idOp"] = pd.to_numeric(df["idOp"])
    df["rate"] = df["rate"].apply(lambda x: float(x.replace(",", ".")))
    return df.to_dict('records')

def upload_datafrae_to_gspread(df):
    credentials = json.loads(os.environ.get("GCLOUD_CREDENTIALS"))
    gc = gspread.service_account_from_dict(credentials)
    sh = gc.open_by_key(os.environ.get("SPREADSHEET_KEY"))
    worksheet = sh.get_worksheet(0)
    worksheet.clear()
    set_with_dataframe(worksheet, df)


#比起JS中习惯的NPM 跟 YARN，py的包管理工具是pip/
# 安装一个Pillow pip install Pillow

from PIL import Image
im = Image.open('liangzi.jpg')
print(im.format, im.size, im.mode)
# JPEG (1920, 1080) RGB
im.thumbnail((200, 100))
im.save('liangzi1.jpg', 'JPEG')
# PYTHON真好玩啊，哈哈哈
# python模块的搜索路径是：默认情况下，Python解释器会搜索当前目录、所有已安装的内置模块和第三方模块，搜索路径存放在sys模块的path变量中

# 如果我们要添加自己的搜索目录，有两种方法：
#
# 一是直接修改sys.path，添加要搜索的目录：
#
# >>> import sys
# >>> sys.path.append('/Users/michael/my_py_scripts')


import pandas as pd
import numpy as np
import datetime

from wtpred.base_predictors import randomforest

import statistics
from keras.models import Sequential
from keras.layers.core import Dense, Activation
from keras.layers.recurrent import SimpleRNN
from keras.optimizers import Adam
from keras.callbacks import EarlyStopping
import tensorflow as tf
import keras.backend as K

N_lookbucks = 6

def make_model(datasetdf):

    ### make error_seq ###
    df = datasetdf.copy()

    # Multi Index
    df.set_index([df.index.year, df.index], inplace=True)
    df.index.names = ['year', 'date']

    error_seq = []
    for name, _ in df.groupby(level='year'):

        # split dataset to test and train
        objective_year = str(name)
        testdf = datasetdf[objective_year].copy()
        traindf = datasetdf.drop(testdf.index)

        # make model
        model = randomforest.make_model(traindf)

        # predict
        testdf = testdf.dropna(how='any')
        X_test = testdf.iloc[:, 1:]
        y_test = testdf.iloc[:, 0]
        pred = model.predict(X_test)

        # calculate error
        error = y_test.values - pred
        error_seq += list(error)

    ### learn error_seq ###

    ## normalize
    mean = statistics.mean(error_seq)
    std = statistics.stdev(error_seq)
    error_seq = [(x - mean) / std for x in error_seq]

    ## make_dataset
    explanatory, response = [], []
    for i in range(len(error_seq)-N_lookbucks):

        # if there are NaN, skip
        if np.nan in error_seq[i:i+N_lookbucks+1]:
            continue

        explanatory.append(error_seq[i:i+N_lookbucks])
        response.append(error_seq[i+N_lookbucks])

    # reshape
    X = np.array(explanatory).reshape(len(explanatory), N_lookbucks, 1)
    y = np.array(response).reshape(len(response), 1)

    ## define model
    N_hidden = 1000
    N_inout_neurons = 1

    model = Sequential()
    model.add(SimpleRNN(N_hidden, batch_input_shape=(None, N_lookbucks, N_inout_neurons), return_sequences=False))
    model.add(Dense(N_hidden))
    model.add(Dense(N_hidden))
    model.add(Dense(N_inout_neurons))
    model.add(Activation("linear"))
    optimizer = Adam(lr=0.001)
    model.compile(loss="mean_squared_error", optimizer=optimizer)

    ## learn
    epochs = 100
    batch_size = 64

    early_stopping = EarlyStopping(monitor='val_loss', patience=10, verbose=1)
    fit = model.fit(X, y,
            batch_size=batch_size,
            epochs=epochs,
            validation_split=0.1,
            callbacks=[early_stopping]
    )

    return model

def predict_nextday_error(model, realdf, preddf, nextdate_str):
    realdf.columns = ['real']
    preddf.columns = ['pred']

    date_index = pd.date_range(end=nextdate_str, periods=N_lookbucks+1, closed='left')

    emptydf = pd.DataFrame([0]*N_lookbucks, index=date_index)
    df = emptydf.join(realdf, how='inner')
    df = df.join(preddf, how='inner')
    df = df[['real', 'pred']]

    df['error'] = df['real'] - df['pred']

    if df['error'].isnull().any():
        return 0

    # reshape
    X = np.array(df['error']).reshape(1, N_lookbucks, 1)

    ## predict
    predicted_error = model.predict(X).flatten()

    return predicted_error


"""Finding the time differences between two timestamps or datetime objects."""
# ------------------------ Import libraries and functions ---------------------
import datetime
import re

import dateutil.parser

from cleaning_utils.constants import DATE_SET, ISO_PATTERN, TIME_SET

# ---------------------------- function definition ----------------------------

def datetime_diff(dt1, dt2):
    """Calculates the differences between two timestamps or datetime objects.

    Args:
        dt1 (dict or ISO 8601 str):
            The first datetime input
        dt2 (dict or ISO 8601 str):
            The second datetime input

    Returns:
        A <class 'datetime.timedelta'> from dt2 - dt1.

    Raises:
        Exception: If the input is not a known ISO 8601.

    Examples:
        >>> datetime1 = {
        ...         'month': 1,
        ...         'year': 1399,
        ...         'day': 4,
        ...         'hour': 23,
        ...         'minute': 5,
        ...         'second': 25
        ...     }
        >>> datetime2 = {
        ...         'month': 1,
        ...         'year': 1399,
        ...         'day': 8,
        ...         'hour': 0,
        ...         'minute': 1,
        ...         'second': 0
        ...     }
        >>> datetime_diff(datetime1, datetime2)
        datetime.timedelta(days=3, seconds=3335)
        >>> datetime3 = '2020-03-23T23:05:25+03:26'
        >>> datetime4 = '2020-03-27T00:01:00+03:26'
        >>> datetime_diff(datetime3, datetime4)
        datetime.timedelta(days=3, seconds=3335)

    """
    # -------------------------- string input branch --------------------------
    # string pattern for checking whether the input is a known ISO 8601
    if isinstance(dt2, str) and isinstance(dt1, str):
        if bool(re.match(re.compile(ISO_PATTERN), dt1)) and bool(
            re.match(re.compile(ISO_PATTERN), dt2),
        ):
            result = dateutil.parser.parse(dt2) - dateutil.parser.parse(dt1)
        else:
            raise Exception("unknown datetime formats")
    # --------------------------- dict input branch ---------------------------
    elif isinstance(dt2, dict) and isinstance(dt1, dict):
        if len(dt1.keys()) == len(dt2.keys()) == 6:
            if (TIME_SET <= dt1.keys()) and (TIME_SET <= dt2.keys()):
                result = datetime.datetime(
                    dt2["year"],
                    dt2["month"],
                    dt2["day"],
                    dt2["hour"],
                    dt2["minute"],
                    dt2["second"],
                ) - datetime.datetime(
                    dt1["year"],
                    dt1["month"],
                    dt1["day"],
                    dt1["hour"],
                    dt1["minute"],
                    dt1["second"],
                )
            else:
                raise Exception("unknown datetime formats")
        elif len(dt1.keys()) == len(dt2.keys()) == 3:
            if (DATE_SET <= dt1.keys()) and (DATE_SET <= dt2.keys()):
                result = (
                    datetime.datetime(
                        dt2["year"],
                        dt2["month"],
                        dt2["day"],
                    )
                    - datetime.datetime(dt1["year"], dt1["month"], dt1["day"])
                )
            else:
                raise Exception("unknown datetime formats")
        else:
            raise Exception("unknown datetime formats")
    # ------------------------- incorrect input branch ------------------------
    else:
        raise Exception("unknown datetime formats")
    return result


from util.data import select_intersection, get_train_label_orig_from_hot, get_assigment_node, get_assignment_index_each_node, get_lists_indices_per_label, get_label_missing, write_dictionary_data, write_dictionary_validation, get_n_sample_to_keep, get_training_testing_partition,get_index_from_one_hot_label
import copy
import numpy as np
import math
from datasets.dataset import get_data
import operator

train_image_all = []
train_label_all = []
test_image_all = []
test_label_all = []
train_label_orig_all = []
keep_track_indices_dataset_train = []
keep_track_indices_dataset_test = []

def multiple_to_single_dataset(cfg):
    cfg.case = 1
    cfg.maxCase = 2

    counter_dataset = 0

    indices_each_node_case = []
    for i in range(0, cfg.maxCase):
        indices_each_node_case.append([])

    for i in range(0, cfg.n_nodes):
        for j in range(0, cfg.maxCase):
            indices_each_node_case[j].append([])

    for dataset in cfg.dataset_list:

        indices_init_train = len(train_label_all)
        indices_init_test = len(test_label_all)

        train_image, train_label, test_image, test_label, train_label_orig = get_data(
            dataset,
            cfg.total_data,
            cfg.dataset_file_path,
            use_test_client_to_data_dict=False)

        # make sure every dataset has train_label_orig
        if len(train_label_orig) == 0:
            train_label_orig = get_train_label_orig_from_hot(train_label)
        train_image_all.extend(train_image)
        train_label_all.extend(train_label)
        test_image_all.extend(test_image)
        test_label_all.extend(test_label)

        number_element_per_node_this_dataset = get_assigment_node(
            train_label_orig, cfg.n_nodes)

        assignment_index_each_node, d = get_assignment_index_each_node(cfg,
            train_label_orig,
            np.arange(indices_init_train, len(train_label_all)),
            number_element_per_node_this_dataset, cfg.n_nodes,
            counter_dataset)
        counter_dataset += 1
        np.save(
            cfg.results_file_path + '/d-' + str(counter_dataset) + '.npy',
            d)

        for c in range(0, cfg.maxCase):
            for i in range(0, cfg.n_nodes):
                indices_each_node_case[c][i].extend(
                    assignment_index_each_node[c][i])

        # indices_each_node.extend()

        ############################################################################################

        keep_track_indices_dataset_train.append(
            np.arange(indices_init_train, len(train_label_all)))
        keep_track_indices_dataset_test.append(
            np.arange(indices_init_test, len(test_label_all)))

    return train_image_all, train_label_all, keep_track_indices_dataset_train, keep_track_indices_dataset_test, indices_each_node_case

def split_focus_noise(keep_track_indices_dataset_train,
                      keep_track_indices_dataset_test):
    # useful for the end to compute loss and accurady
    train_image_focus = [
        train_image_all[i] for i in keep_track_indices_dataset_train[0]
    ]
    train_label_focus = [
        train_label_all[i] for i in keep_track_indices_dataset_train[0]
    ]
    test_image_focus = [
        test_image_all[i] for i in keep_track_indices_dataset_test[0]
    ]
    test_label_focus = [
        test_label_all[i] for i in keep_track_indices_dataset_test[0]
    ]

    train_image_noise = [
        train_image_all[i] for i in range(0, len(train_label_all))
        if i not in keep_track_indices_dataset_train[0]
    ]
    train_label_noise = [
        train_label_all[i] for i in range(0, len(train_label_all))
        if i not in keep_track_indices_dataset_train[0]
    ]
    test_image_noise = [
        test_image_all[i] for i in range(0, len(test_label_all))
        if i not in keep_track_indices_dataset_test[0]
    ]
    test_label_noise = [
        test_label_all[i] for i in range(0, len(test_label_all))
        if i not in keep_track_indices_dataset_test[0]
    ]

    return train_image_focus, train_label_focus, test_image_focus, test_label_focus, train_image_noise, train_label_noise, test_image_noise, test_label_noise

# only on the focus dataset

def train_validation(indices_each_node_case_copy_init, model, dim_w, sim, cfg,
                     lists_indices_per_label_validation, validation_image,
                     validation_label):
    # split benchmark dataset into training and testing
    validation_training_indices, validation_testing_indices, validation_training_image, validation_training_label, validation_testing_image, validation_testing_label = get_training_testing_partition(
        cfg.n_category, lists_indices_per_label_validation, validation_image,
        validation_label)

    dic_ind_grad = {}
    dic_ind_grad_validation = {}

    #print('-----Start of the validation training----------------')
    model.start_consecutive_training(
        model.get_init_weight(dim_w, rand_seed=sim))

    max_val = math.inf
    counter_convergence_reached = 0
    i = 0
    while counter_convergence_reached < 200:

        model.run_one_step_consecutive_training(validation_image,
                                                validation_label,
                                                validation_training_indices)
        w_validation = model.end_consecutive_training_and_get_weights()

        loss_train = model.loss(validation_training_image,
                                validation_training_label, w_validation)
        loss_test = model.loss(validation_testing_image,
                               validation_testing_label, w_validation)
        accu_train = model.accuracy(validation_training_image,
                                    validation_training_label, w_validation)
        accu_test = model.accuracy(validation_testing_image,
                                   validation_testing_label, w_validation)

        if loss_test < max_val:
            max_val = loss_test
            counter_convergence_reached = 0
        else:
            if i > 10:
                counter_convergence_reached += 1
                #print('counter', counter_convergence_reached)

        i = i + 1
        if i > 10000:
            break

        #print('accu train', model.accuracy(validation_training_image, validation_training_label, w_validation))
        #print('accu test', model.accuracy(validation_testing_image, validation_testing_label, w_validation))

        with open(cfg.results_file_path + '/loss-accu-validation.csv',
                  'a') as f:
            f.write(
                str(i) + ',' + str(loss_train) + ',' + str(loss_test) + ',' +
                str(accu_train) + ',' + str(accu_test) + '\n')
            f.close()

    for s in validation_testing_indices:
        norm_diff = model.loss(validation_image[s:s + 1],
                               validation_label[s:s + 1], w_validation)
        dic_ind_grad_validation[s] = (norm_diff, None,
                                      get_index_from_one_hot_label(
                                          validation_label[s]))

    write_dictionary_validation(cfg.results_file_path, '/dico-validation.csv',
                                dic_ind_grad_validation,
                                validation_testing_indices)

    # RANK ALL DATA POINT IN EACH NODE
    for n in range(0, cfg.n_nodes):

        norm_diff_vec = model.loss_vector(train_image_all,train_label_all,w_validation,indices_each_node_case_copy_init[cfg.case][n])
        #print(len(norm_diff_vec))
        #print('check',len(indices_each_node_case_copy_init[cfg.case][n]))
        count = 0

        for s in indices_each_node_case_copy_init[cfg.case][n]:

            #norm_diff = model.loss(train_image_all[s:s + 1],
                                  # train_label_all[s:s + 1], w_validation)

            dic_ind_grad[s] = (norm_diff_vec[count], n,
                               get_index_from_one_hot_label(
                                   train_label_all[s]))
            count+=1

    sorted_d = sorted(dic_ind_grad.items(), key=operator.itemgetter(1))
    #print('here',sorted_d)

    write_dictionary_data(cfg.results_file_path, '/dico-order.csv', sorted_d)

    n_sample_keep, indices_to_keep = get_n_sample_to_keep(
        cfg.results_file_path, len(train_label_all), sorted_d)
    # print('n sample to keep',n_sample_keep)

    return indices_to_keep, w_validation

def get_indices_each_node_opt(cfg,
                              indices_each_node_case,
                              indices_each_node_case_copy_init,
                              indices_to_keep=None):

    option = cfg.option

    if option == 1:
        n_sample_keep = len(keep_track_indices_dataset_train[0])
        indices_to_keep = np.arange(0, n_sample_keep)
        list_nodes_not_enough_data = []
        for n in range(0, cfg.n_nodes):
            indices_each_node_case[cfg.case][n] = select_intersection(
                indices_to_keep, indices_each_node_case_copy_init[cfg.case][n])
            # check if the number of point in a node is smaller than batch size
            if len(indices_each_node_case[cfg.case][n]) == 0:
                list_nodes_not_enough_data.append(n)

    elif option == 2:

        list_nodes_not_enough_data = []
        for n in range(0, cfg.n_nodes):
            indices_each_node_case[
                cfg.case][n] = indices_each_node_case_copy_init[cfg.case][n]
            # check if the number of point in a node is smaller than batch size
            if len(indices_each_node_case[cfg.case][n]) == 0:
                list_nodes_not_enough_data.append(n)

    elif option == 0:

        sampler_list = []
        train_indices_list = []
        list_nodes_not_enough_data = []
        for n in range(0, cfg.n_nodes):

            indices_each_node_case[cfg.case][n] = select_intersection(
                indices_to_keep, indices_each_node_case_copy_init[cfg.case][n])
            # check wehter the number of point in a node is smaller than batch size
            if len(indices_each_node_case[cfg.case][n]) == 0:
                list_nodes_not_enough_data.append(n)

    return indices_each_node_case, list_nodes_not_enough_data

def get_benchmark_dataset(sim, size_validation_set, cfg):

    #init
    label_missing_in_validation = [1000]
    while (len(label_missing_in_validation) != 0):

        train_indices_focus = keep_track_indices_dataset_train[0]
        train_indices_focus_copy = copy.deepcopy(train_indices_focus)
        np.random.seed(sim)
        np.random.shuffle(train_indices_focus_copy)
        list_indices_validation = train_indices_focus_copy[
            0:size_validation_set]
        validation_image = [
            train_image_all[i] for i in list_indices_validation
        ]
        validation_label = [
            train_label_all[i] for i in list_indices_validation
        ]

        sim += 1

        # split training into train and test
        #make sure validation contains each category of label
        lists_indices_per_label_validation = get_lists_indices_per_label(
            cfg.n_category, validation_label)
        label_missing_in_validation = get_label_missing(
            cfg.n_category, lists_indices_per_label_validation)

    return lists_indices_per_label_validation, validation_image, validation_label