import snappy.snap as snap
import snappy.snapext as snapext
from db.interface import *
import sys

class Graph:
	COLLAB = 'Gcollab'
	PULL = 'Gpull'
	FOLLOW = 'Gfollow'
	WATCH = 'Gwatch'
	FORK = 'Gfork'

def construct(graph_type, edgelist):
	G = snapext.EUNGraph()

	for edge in edgelist:
		u = edge[0]
		v = edge[1]
		created_at = edge[2]

		weight= edge[3] + 1

		if not G.IsNode(u):
			G.AddNode(u)

		if not G.IsNode(v):
			G.AddNode(v)

		if not G.IsEdge(u, v):
			G.AddEdge(u, v, created_at, weight)

	return G

def getKHopN(G, node, k):
	s = set()
	for i in range(2, k+1):
		v= snap.TIntV()
		snap.GetNodesAtHop(G.G, node, i, v, False)
		# add node IDs in v to s
		for item in v:
			s.add(item)
	return s

def getSubGraph(G, source, k):
	nodes= getKHopN(G, source, k)
	# to nodes add all the neighbors of source
	nodes.add(source)
	sourceNode= G.GetNI(source)
	for neighborID in sourceNode.GetOutEdges():
		nodes.add(neighborID)

	# insert nodes into TIntV
	nodesV= snap.TIntV()

	for node in nodes:
		nodesV.Add(node)
	subG= snap.GetSubGraph(G.G, nodesV)
	return subG

def get_db_graph(graph_type, uid, dbresponse):
	edgelist = []

	for row in dbresponse:
		try:
			# Follow graph has the weight after jday. Sad.
			if graph_type==Graph.FOLLOW:
				weight = row[4]
			else:
				weight = row[3]

			edgelist.append((uid[row[0]], uid[row[1]], row[2], weight))

		except KeyError:
			#sys.stderr.write("Skipping %s, %s as one of them was not in uid\n"%(row[0], row[1]))
			pass
			# Ignore if some user is not found in uid

	return construct(graph_type, edgelist)

def get_db_graphs(db, uid, min_date=None, max_date=None):
	reader = DBReader(db)

	collab = reader.collaborators(min_date, max_date)
	Gcollab = get_db_graph(Graph.COLLAB, uid, collab)

	follow = reader.followers(min_date, max_date)
	Gfollow = get_db_graph(Graph.FOLLOW, uid, follow)

	# Pull is now a part of collab graph
	# pull = reader.pull(min_date, max_date)
	# Gpull = get_db_graph(Graph.PULL, uid, pull)

	watch = reader.watch(min_date, max_date)
	Gwatch = get_db_graph(Graph.WATCH, uid, watch)

	fork = reader.fork(min_date, max_date)
	Gfork = get_db_graph(Graph.FORK, uid, fork)

	reader.close()
	return {Graph.COLLAB: Gcollab, Graph.FOLLOW:Gfollow, Graph.WATCH: Gwatch, Graph.FORK: Gfork}

def get_feat_graphs(db, uid, min_date=None, max_date=None):
	reader = DBReader(db)

	print("Getting Gfollow")
	follow = reader.followers(min_date, max_date)
	Gfollow = get_db_graph(Graph.FOLLOW, uid, follow)

	print("Getting Gwatch")
	watch = reader.watch(min_date, max_date)
	Gwatch = get_db_graph(Graph.WATCH, uid, watch)

	print("Getting Gfork")
	fork = reader.fork(min_date, max_date)
	Gfork = get_db_graph(Graph.FORK, uid, fork)

	reader.close()
	return {Graph.FOLLOW:Gfollow, Graph.WATCH: Gwatch, Graph.FORK: Gfork}

def get_base_dict(Gcollab, feature_graphs):
	base_graphs = {}

	for key,val in feature_graphs.iteritems():
		base_graphs[key] = val

	base_graphs[Graph.COLLAB] = Gcollab

	return base_graphs

def get_collab_graph(db, uid, min_date=None, max_date=None):
	reader = DBReader(db)

	print("Getting Collab graph")
	collab = reader.collaborators(min_date, max_date)
	Gcollab = get_db_graph(Graph.COLLAB, uid, collab)

	reader.close()
	return Gcollab

def get_only_collab_graph(db, uid, min_date=None, max_date=None):
	reader = DBReader(db)
	print("Getting Only Collab graph")
	collab = reader.only_collab(min_date, max_date)
	Gcollab = get_db_graph(Graph.COLLAB, uid, collab)
	reader.close()
	return Gcollab

def split_feat_graphs(base_graphs):
	feature_graphs = []

	# THIS ORDER MUST BE MAINTAINED, otherwise feature vectors will keep changing order
	# Gfollow, Gfork, Gwatch
	for graph_type in sorted(base_graphs.keys()):
		graph = base_graphs[graph_type]

		if graph_type!=Graph.COLLAB:
			feature_graphs.append(graph)

	return feature_graphs

def get_reverse_graph_map(base_graphs):
	rmap = {}

	for graph_type in sorted(base_graphs.keys()):
		graph = base_graphs[graph_type]
		rmap[graph] = graph_type

	return rmap

def print_graph_stats(name, G):
	print("Graph type: %s"%name)
	print("Graph Nodes: %d. Edges= %d"%(G.GetNodes(), G.GetEdges()))

def print_stats(base_graphs):
	for key, val in base_graphs.iteritems():
		print_graph_stats(key, val)


#! src\bin\env python3
# -*- coding: utf-8 -*-

#Importation des modules

import os, sys, pathlib, argparse, re, time, threading
from pathlib import Path
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import time
import webbrowser
from math import *

#Installation de Tkinter selon les différentes versions
try:
    try :
        import tkinter as Tk
        from tkinter import *
        from tkinter.messagebox import *
        from tkinter.filedialog import *
    except:
        import Tkinter as Tk
        from tKinter import *
        from tKinter.messagebox import *
        from tKinter.filedialog import *
except:
    raise ImportError('Tkinter non disponible')

#Définition des couleurs
orange = '#fde3d9'
orangedark = '#f47645'
blue = '#7acfdf'
bluedark = '#61a5b2'
grey = "#E3E2E1"

#Procédure de fermeture de le fenêtre princiaple
def Exit():
    if askyesno("Exit", "Do you really want to exit kofi :'( ?"):
        showwarning('Exit', "Exiting...see you soon")
        main.destroy()

    else:
        showinfo('Exit', 'Welcome back on kofi!')

#Initailisation de la fenêtre principale
main = Tk()
main.title("Kofi")
main.configure(bg = "white")
main.resizable(0,0)
#main.attributes("-fullscreen", 2)
#On récupère la largeur (ws) et la hauteur (hs) de la fenêtre
ws = main.winfo_screenwidth()
hs = main.winfo_screenheight()
#calcul la position de la fenetre adaptée à l'écra,
x = (ws/2) - (ws/2)
y = (hs/2) - (hs/2)
#applique la taille et la position
main.geometry('%dx%d+%d+%d' % (ws, hs, 0, 0))

#intercepte l'evenement quit pour informer l'utilisateur, appel la fonction Exit
main.protocol("WM_DELETE_WINDOW", Exit)

# Fonction d'affichage de l'heure

time1 = ''
time2 = ''

def hour():

    global time1
    # Heure actuelle du PC
    time2 = time.strftime('%H:%M:%S')
    # Si le temps récupéré est différent
    if time2 != time1:
        time1 = time2
        clock.config(text=time2)
    # Toute les 200ms la fonction s'appelle (récursivité) pour s'actualiser
    clock.after(200, hour)

# Fonction pour ouvrir le manuel
new = 1
url = "https://github.com/emiracherif/VCF-project/blob/master/README.md"
url2 = "https://github.com/emiracherif/VCF-project/blob/master/INSTALL.md"

def openweb():
    webbrowser.open(url,new=new)

def openweb2() :
    webbrowser.open(url2, new= new)

#######################################################
################ Ouverture du fichier! ################
#######################################################

#Récupération du chemin du fichier selectionné
def open_vcf() :
    #main.configure(bg = grey)
    filepath = askopenfilename(title="Open vcf file",filetypes=[('vcf files','.vcf')])
    # On lance la suite du programme avec le chemin récupéré
    verif_opening(filepath)

# Fonction de vérification sur le fichier sélectionnée

compteur = 0

def verif_opening(filepath) :

    ######### Initatilisation de la toplevel ##########
    loading = Toplevel(main, cursor = "watch", bg = orange)
    cpt = Toplevel(loading, bg = blue)
    loading.geometry("%dx%d%+d%+d" % (ws//7,hs//6,ws//3,hs//3))
    loading.title =("Kofi")
    loading.resizable(0,0)
    lab=Label(loading, text="Loading file...", font = ("Helvetica", 18, "bold"), bg = orange)
    lab.grid(sticky='ew')

    ########### Création d'un compteur du temps d'analyse #####
    def compteur_sec() :

        #Initialisation du compteur
        # On ne peut pas quitter le compteur manuellement
        cpt.overrideredirect(TRUE)
        cpt.title =("waiting time")
        cpt.resizable(0,0)
        cpt.geometry('%dx%d+%d+%d' % (ws//5,hs//5,ws//7,hs//6))

        compteur = 0
        text = Label(cpt, text="Waiting time :", font =("", 14), bg = blue, fg = "white")
        coff = Label(cpt, text= "Go make yourself some coffe...",font =("", 8), bg = blue, fg = "white" )
        compteur_lbl = Label(cpt, text= str(compteur), font=("", 20), bg = blue, fg = "white")
        def incremente():
        #Incrémentation
            global compteur
            compteur += 1
            compteur_lbl['text'] = str(compteur) + str("sec")
            cpt.after(1000, incremente)
        incremente()
        # appel récursif (compteur infini)

        #Placement des widgets sur la fenêtre
        text.pack(pady = 25)
        compteur_lbl.pack()
        coff.pack(side = BOTTOM)

    # On récupère le chemin du fichier dans une variable si le fichier existe

    try :
        filename = Path(filepath)
    except :
        loading.destroy()
        return showwarning("Warning", "Please select a file!")

    # Vérifie que l'utilisateur a bien choisi un fichier
    if not Path.is_file(filename):
        loading.destroy()
        return showwarning("Warning", "Please select a file!")

    if showwarning :
        loading.destroy

    prog = Label(loading, bg = orange, text = "Path OK!", font = ("Helvetica", 10))
    prog.grid(sticky='ew')

    ################################################################################
    # Normalement l'utilisteur choisi forcément un fichier de type VCF
    # des doubles vérifications sont tout de même effectuées
    ###########################################################################

    #Si le fichier séléctionné n'existe pas, ou plus
    if not filename.exists():

        #Exit de la fonction de chargement, affichage d'une erreur
        return showerror("Error", "Oops, this file dosn't exist anymore!")

    #Sinon, affichage que c'est ok!
    prog.configure(text = "Path OK! \n File OK!")

    # Appel du compteur
    compteur_sec()

    ################# Ouverture et vérification à l'intérieur du fichier ####################

    def next() :
        prog.configure(text = "")
        deeper.destroy()
        fileformat = re.findall("##fileformat=VCFv4",fd)
        chrom = re.findall("#CHROM",fd)

        # Fonction de passage à l'étape de filtration du vcf
        def ex() :
            filtre(filename)
            loading.destroy()

        # Analyse des entêtes dans le fichier
        if not chrom or not fileformat :
            #Si on ne trouve pas ces deux entêtes, on exit la fonction et on renvoit une erreur
            return showerror("Error", "Oops, mandatory header line doesn't match with vcf type")
        else :
            # Retour positif sur le fichier
            prog.configure(text = "We opened it and ... \n you're file is perfectly fine! ")
            # Affichage du bouton pour passer à l'étape suivante!
            yeah = Button(loading, text="Next step", relief = 'groove', command = ex,fg = "white", bg = orangedark, cursor='hand2')
            yeah.grid(sticky='n')

    ### Vérifie que c'est bien un vcf
    if filename.suffix==".vcf":

        #Ouverture du fichier dans la variable fd
        fd=filename.read_text()

        # Retour positif
        prog.configure(text = "Path OK! \n File OK! \n .vcf OK!")

        # Affichage du bouton pour vérifier l'intérieur du fichier
        deeper = Button(loading, relief = 'groove', text="Let's go deeper", command = next,fg = "white", cursor='hand2', bg = orangedark)
        deeper.grid(sticky = 'n') #Affiche le bouton qui lance l'analyse à l'intérieur du fichier

    else:
        # On quitte la fonction et donc la toplevel avec un message d'erreur
        return showerror("Error", "Oops, wrong extention!")

##########################################################
############# Mise en place des filtres ##################
##########################################################

m = ""
DP = ""
P = ""

def filtre(filename) :

    # Chargement de l'interface
    main.configure(bg="white")
    back.destroy()
    load.destroy()
    info = Label(main, text = "{} is loaded.".format(filename.name), font = "Helevtica 14 bold" , bg = "white" )
    info.place(relwidth = 0.9,rely = 0.15)

    # Création des cadres de séléctions
    cadre_data = LabelFrame(main, bd=1, text = "Missing data", font = "Helevtica 14", bg = orange)
    cadre_dp = LabelFrame(main, bd=1, text = "DP (Read Depth)",  bg = orange, font = "Helevtica 14")
    cadre_gen = LabelFrame(main, bd=1, text =  "Genotype percentage",  bg = orange, font = "Helevtica 14")

    #Placement des cadres
    cadre_data.place(relwidth = 0.9, relx= 0.05, rely = 0.3)
    cadre_dp.place(relwidth = 0.9,relx = 0.05, rely = 0.5)
    cadre_gen.place(relwidth = 0.9,relx = 0.05, rely= 0.7 )

    #Création des logos d'informations pour chaque valeur
    def info_data() :
        showinfo(title = 'Missing data', message = "KOFI missing data corresponds to missing genotypes specified with two dots separated by a slash './.'")

    def info_dp() :
        showinfo(title = 'DP', message = 'KOFI DP corresponds to the minimum percentage of read depth per genotype')

    def info_gen() :
        showinfo(title = 'Genotype percentage', message = 'KOFI Genotype percentage corresponds to the minimum percentage of genotype required with the predefined DP')

    #Ecritures au sein des cadres
    l1 = Label(cadre_data, text = " Please choose the percentage of tolerated missing data shall fall within 0-100 interval", bg = orange)
    b1 = Button(cadre_data,bitmap = 'question',fg="black", command = info_data).pack(side = RIGHT)
    l1.pack(side = LEFT)
    l2 = Label(cadre_dp, text = "Please choose the minimum Depth of coverage (DP) per genotype (numeric value expected)", bg = orange)
    b2 = Button(cadre_dp,bitmap = 'question', fg = "black", command = info_dp).pack(side = RIGHT)
    l2.pack(side = LEFT)
    l3 = Label(cadre_gen, text = "Please choose the minimum percentage of genotype required harboring the defined DP (numeric value expected)", bg = orange)
    b3 = Button(cadre_gen,bitmap = 'question',fg="black", command= info_gen).pack(side = RIGHT)
    l3.pack(side = LEFT)

    #Choix des données manquantes avec attribution d'une valeur par défaut

    def1 = StringVar(main)
    def1.set(5)
    choix_data = Spinbox(cadre_data, from_=0, to=50, increment=5, bg="white", textvariable = def1)
    choix_data.pack()

    #Choix de la DP
    def2 = StringVar(main)
    def2.set(10)
    choix_DP = Spinbox(cadre_dp, from_=0, to=30, increment=5, bg="white", textvariable = def2)
    choix_DP.pack()

    #Choix du % de genotype
    def3 = StringVar(main)
    def3.set(95)
    choix_gen = Spinbox(cadre_gen, from_= 0, to=100, increment = 5, bg = "white", textvariable = def3 )
    choix_gen.pack()

    #Validation des valeurs
    def val() :

        #Affectation ou non des de la valeur de la DP
        try :
            DP = int(choix_DP.get())
            #Desactivation de la séléction
            choix_DP.configure(state = 'disabled')
        except :
            return showerror(title = "Error", message = "Numeric value expected for DP")

        #Affectation ou non des de la valeur de la data
        try :
            m = int(choix_data.get())
            #Desactivation de la séléction
            choix_data.configure(state='disabled')
        except :
            return showerror(title = "Error", message = "Numeric value expected for missing data")

        #Affectation ou non des de la valeur du genotype
        try :
            P = int(choix_gen.get())
            #Desactivation de la séléction
            choix_gen.configure(state='disabled')
        except :
            return showerror(title = "Error", message = "Numeric value expected for genotpye")

        #Passage au nettoyage du fichier
        nettoyage(filename,m,P,DP)

    #Création et placement du bouton de validation
    valide = Button(main, relief = 'groove', bg = bluedark, fg = "white",font = "Helevtica 12", text = "Validate", command = val)
    valide.place(relx = 0.8, rely= 0.8 )

####################################################################
#############      ETAPE DE NETTOYAGE DU PREMIER VCF        ########
############# CREATION DE 4 NOUVEAUX FICHIERS PLUS LISIBLES ########
####################################################################

def nettoyage(filename, m, P, DP) :

    # Creation d'une toplevel de chargement
    nettoyage = Toplevel(main, cursor = "watch", bg = blue)
    nettoyage.geometry("%dx%d%+d%+d" % (ws//3,hs//3,ws//3,hs//3))
    nettoyage.title =("Kofi")
    lab=Label(nettoyage, text="\n Creating a new filtred vcf file with the choosen filtred. \n \n Conversion of the filtred vcf in tabular genotyping file.\n \n Creating a genotype distribution file.",fg = "white", font = ("Helvetica", 18, "bold"), bg = blue)
    lab.pack()

    ####### Etape 1.1 nettoyage & filtres   #######

    nb_indiv = 0

    #Creation d'un nouveau fichier vcf et on récupère son nom
    kofile = open(filename.stem+'-m'+str(m)+'-DP'+str(DP)+'-P'+str(P)+'.vcf','a')
    nameko = filename.stem+'-m'+str(m)+'-DP'+str(DP)+'-P'+str(P)+'.vcf'

    #Creation du nouveau fichier de genotypage (conversion du vcf) + son nom
    konvfile=open(filename.stem+'-m'+str(m)+'-DP'+str(DP)+'-P'+str(P)+'.geno'+'.txt','a')
    namekonv = filename.stem+'-m'+str(m)+'-DP'+str(DP)+'-P'+str(P)+'.geno'+'.txt'

    #Creation du nouveau fichier de genotypage distrib + son nom
    kogefile=open(filename.stem+'-m'+str(m)+'-DP'+str(DP)+'-P'+str(P)+'.geno'+ '.distrib'+'.txt','a')
    namekoge = filename.stem+'-m'+str(m)+'-DP'+str(DP)+'-P'+str(P)+'.geno'+ '.distrib'+'.txt'

    #Ecriture de l'entête dans le nouveau vcf kofile
    from re import finditer
    fd = open(filename,"r")

    for line in fd :

        headline = re.search("^##.+",line)

        if headline :

            kofile.write("\n"+ headline.group(0))

        chromline = re.search("#CHROM.+",line)
        if chromline :
            kofile.write("\n"+chromline.group(0))
            # Comptage du nombre d'individus dans le vcf
            liste = chromline.group(0).split("\t")
            nb_indiv = len(liste[9:])

            #Ecriture de l'entete du fichier de conversion konvfile
            konvfile.write("\n"+"SNP_ID"+"\t"+"CHROM"+"\t"+"POS"+"\t"+"REF"+"\t"+"ALT")
            liste_ind=liste[9:]
            for i in liste_ind:
                konvfile.write("\t"+i)

            #Ecriture de l'entete du fichier de geno distrib kogefile
            kogefile.write("\n"+"SNP_ID"+"\t"+"CHROM"+"\t"+"POS"+"\t"+"REF"+"\t"+"ALT"+
                           "\t"+"Ho1"+"\t"+"aHo1_Total"+"\t"+"Ho2"+"\t"+"aHo2_Total"+"\t"+"Ho3"+
                           "\t"+"aHo3_Total"+"\t"+"He1"+"\t"+"bHe1_Total"+"\t"+"He2"+"\t"+"bHe2_Total"+
                           "\t"+"He3"+"\t"+"bHe3_Total"+"\t"+"MD"+"\t"+"mD_Total"
                           )

    #Extraction des qualités et DP générales
        qual = re.search("\s(\d+\.\d+)\s",line)
        dp_g = re.search(";(DP=)(\d+);",line)

        if qual and dp_g:

            #Filtre sur la qualité et la DP générale
            if (float(qual.group(0)) > 30 and int(dp_g.group(2))>= (int(nb_indiv)*10)) :

                #Filtre sur un pourcentage max de données manquantes tolérées
                missing_iterator = finditer("\.\/\.", line)
                missing_count = 0
                for match in missing_iterator:
                    missing_count +=1

                # Condition sur le pourcentage de de données manquantes tolérées choisi par l'utilisateur ou celle par défaut
                if float((missing_count*100)/nb_indiv)<=float(m):

                    dp_geno_it=finditer(":(\d+):", line)
                    dp_geno_count=0

                    # Condition sur la DP/genotype minimale choisie par l'utilisateur ou celle par défaut
                    for mat in dp_geno_it:
                        if(int(mat.group(1))>=int(DP)):
                            dp_geno_count+=1

                    # Condition sur le pourcentage minimal d'individus choisi par l'utilisateur ayant la DP/genotype minimale  ou par défaut
                    if float((dp_geno_count*100)/nb_indiv)>=float(P):
                        # Ecriture du nouveau vcf filtré kofile
                        kofile.write("\n"+line)
    fd.close()
    kofile.close()

    ####### Etape 1.2 VCF2genofile #######

    # Ouverture du nouveau fichier vcf
    ko = open(filename.stem+'-m'+str(m)+'-DP'+str(DP)+'-P'+str(P)+'.vcf',"r").readlines()

    # Remplissage du geno file
    for line in ko :

        geno = re.search("(^[a-zA-Z]+\d+)\s+(\d+)\s+.+\s+([a-zA-Z]+)\s+([a-zA-Z]),*([a-zA-Z]*)\s",line)

        # Si la ligne n'est pas dans l'entête
        if not line.startswith('\n') and not line.startswith('#')  :
            head=line.split("\t")
            konvfile.write("\n"+head[0]+"_"+head[1]+"\t"+head[0]+"\t"+head[1]+"\t"+head[3]+"\t"+head[4])
            kogefile.write("\n"+head[0]+"_"+head[1]+"\t"+head[0]+"\t"+head[1]+"\t"+head[3]+"\t"+head[4])

    # Conversion du format des genotypes (numerique --> allélique)
        if geno:

            ref=geno.group(3)
            alt=geno.group(4)
            alts=geno.group(5)

            geno_iterator = finditer("(.)\/(.)", line)
            g1=1
            g2=1
            g3=1
            g4=1
            g5=1
            g6=1
            gg=1

            for mag in geno_iterator:

                if mag.group(0)=="0/0" or mag.group(0)=="0|0":
                    konvfile.write("\t"+ref+"/"+ref)
                    g1+=1

                elif mag.group(0)=="1/1"or mag.group(0)=="1|1":
                    g2+=1
                    konvfile.write("\t"+alt+"/"+alt)

                elif mag.group(0)=="2/2"or mag.group(0)=="2|2":
                    g3+=1
                    konvfile.write("\t"+alts+"/"+alts)

                elif mag.group(0) in ('0/1','0|1', '1/0','1|0'):
                    g4+=1
                    konvfile.write("\t"+ref+"/"+alt)

                elif mag.group(0) in ('0/2','0|2', '2/0','2|0'):
                    g5+=1
                    konvfile.write("\t"+ref+"/"+alts)

                elif mag.group(0) in ('1/2','1|2', '2/1','2|1'):
                    g6+=1
                    konvfile.write("\t"+alt+"/"+alts)

                else:
                    gg+=1
                    konvfile.write("\t"+mag.group(0))

            kogefile.write("\t"+ref+"/"+ref+"\t"+str(g1)+"\t"+alt+"/"+alt+"\t"+str(g2)+"\t"+alts+"/"+alts+"\t"+str(g3)+
                               "\t"+ref+"/"+alt+"\t"+str(g4)+"\t"+ref+"/"+alts+"\t"+str(g5)+"\t"+alt+"/"+alts+"\t"+str(g6)+
                               "\t"+mag.group(0)+"\t"+str(gg))
    konvfile.close()
    kogefile.close()

    ######################################################
    ############# STATS SUPPLEMENTAIRES ##################
    ######################################################

    ins = 0
    dell = 0
    sub = 0

    # Ouverture du fichier de génotypage
    kogefile=open(filename.stem+'-m'+str(m)+'-DP'+str(DP)+'-P'+str(P)+'.geno'+ '.distrib'+'.txt','r')

    # Ouverture d'un nouveau fichier stat
    kostatfile=open(filename.stem+'-m'+str(m)+'-DP'+str(DP)+'-P'+str(P)+'.stat'+'.txt','a')
    namekostat = str(filename.stem+'-m'+str(m)+'-DP'+str(DP)+'-P'+str(P)+'.stat'+'.txt')

    dico = {}

    for line in kogefile:

        #keys cherche le chromosome et la position, var cherche l'alt et le ref
        keys= re.search("LG(\d+)_(\d+)", line)
        var = re.search("\d+\s([a-zA-Z]+)\s(([a-zA-Z]+),?([a-zA-Z]+)?)",line)

        if keys:

            # On associe les resultats à des variables
            chrom = keys.group(1)
            pos = keys.group(2)

            ref = var.group(1)
            alt = var.group(3)
            alts = var.group(2)

            # On initialise deux compteurs : un pour le nucl ref l'autre pour l'alternatif
            count = 0
            co = 0

            # comptage des caractères du référent
            for char in ref :
                count +=1

            #Comptage des caractères de l'alt (si il y'en a plusieurs, on prend en compte uniquement le premier)
            for char in alt :
                co += 1

            #Comparaison pour trouver le nombre de substitutions, d'insertions et de délétions.
            if int(count)-int(co) == 0 :
                sub +=1
            elif int(count)-int(co) <= 0 :
                ins += 1
            else :
                dell += 1

            ######## DICTIONNAIRE ##########

            if chrom in dico.keys():
                dico[chrom][pos] = (ref, alts)

            else :
                dico[chrom] = {}
                dico[chrom][pos] = (ref, alts)

    # Taux de mutation
    tsub = ceil((sub/ (sub+ins+dell)) * 100)
    tins = ceil((ins/ (sub+ins+dell)) * 100)
    tdel = ceil((dell/ (sub+ins+dell)) * 100)

    # Ecriture dans le nouveau fichier de stat
    kostatfile.write('Types de mutation \t' + 'Substitutions' +'\t Insertions' +'\t Deletions \n')
    kostatfile.write('Nombre \t' + str(sub) +'\t' + str(ins) +'\t' + str(dell) + '\n')
    kostatfile.write('Pourcentage\t' + str(tsub) +'\t' + str(tins) +'\t' + str(tdel) + '\n\n')

    #Comptages de toutes les positions
    nb_pos = 0
    for item in dico :
        for keys in dico[chrom] :
            nb_pos += 1

    # Comptage des mutations par chromosomes
    nb_mut = 0
    liste_ch = []
    liste_sch = []

    kostatfile.write('Chromosome \t' + 'Nombre de mutations \t' + 'Pourcentage \n')

    for item in dico :
        nb_mut = 0
        for posi in dico[item] :
            nb_mut +=1

        # Arrondi supérieur
        a = ceil((nb_mut/nb_pos)*100)

        # Ecriture dans le fichier
        kostatfile.write(str(item) +'\t' + str(nb_mut) + '\t ' + str(a) + '\n')

        liste_ch.append(str(item))
        liste_ch.append(str(nb_mut))
        liste_sch.append(str(a))

    kostatfile.close()
    kogefile.close()

    showinfo(title="files created", message = "Created on your working directory : \n{}".format('-'+nameko+'\n'+'-'+namekonv +'\n'+'-'+namekoge) + '\n'+'-'+namekostat)

    if showinfo :
        # On efface les widget de la fenêtre pour passer aux stats...
        for chose in main.winfo_children():
            chose.destroy()
        # ... mais on garde le menu et l'entête du programme
        label2 = Label(main, bg=bluedark, font = "helevtica 10 italic")
        label = Label(main, text="Welcome on Kofi.dev", bg = blue, font = "helevtica 10 italic")
        label2.place(relheight = 0.05,relwidth = 1,rely = 0)
        label.place(relwidth = 1, rely = 0  )
        thend = Label(main, text = "Thanks for using KOFI!", bg = "white", fg = bluedark, font = 'helevtica 32 bold')
        thend.place(rely = 0.4, relx = 0.4)

        # Création d'un menu sur la fenêtre
        menubar = Menu(main)
        menu1 = Menu(menubar, tearoff=0)
        #menu1.add_command(label="About kofi")
        menu1.add_command(label="About kofi", command = openweb)
        menu1.add_command(label="User guide", command = openweb2)
        menubar.add_cascade(label="Help", menu=menu1)
        #Attribution du menu au main
        main.config(menu=menubar)
        return heat(m,DP,P, filename)

################################################
########### Affichage de l'heatmap #############
################################################

def heat(m, DP, P,filename) :

    #Utilisation de pandas pour réaliser un dataframe
    data = pd.read_table(filename.stem+'-m'+str(m)+'-DP'+str(DP)+'-P'+str(P)+'.geno'+ '.distrib'+'.txt', delimiter="\t")

    df=pd.DataFrame(data[['POS','CHROM','aHo1_Total','aHo2_Total','aHo3_Total','bHe1_Total','bHe2_Total','bHe3_Total','mD_Total']])

    # Default heatmap: just a visualization of this square matrix

    df = df.pivot_table(index=data[['POS']],
                        values=data[['POS','aHo1_Total','aHo2_Total','aHo3_Total','bHe1_Total','bHe2_Total','bHe3_Total','mD_Total']])

    p1 = sns.heatmap(df,vmin=0,vmax=250,cmap="RdBu_r")

    plt.show()

#############################################################################
################# INTERFACE FENETRE PRINCIPALE ##############################
#############################################################################

#Chargement interface
back = Button(main, bg = orange, bd = 0 )
#back2 = Button(main, bg = orange, bd = 0)
load = Button(main, text ='Load VCF', font = "Helevtica 16 bold", cursor="circle",command = open_vcf, bg = orangedark, fg = "white", bd = 0.2 , pady = 4, padx = 5)
label2 = Label(main, bg=bluedark, font = "helevtica 10 italic")
label = Label(main, text="Welcome on Kofi.dev", bg = blue, font = "helevtica 10 italic")

#PLacement de l'interface
back.place(rely = 0.4, relx = 0.4 , height = 300, width = 300 )
#back2.place(rely = 0.4, relx = 0.4 , height = 200, width = 200 )
load.place(rely = 0.4, relx = 0.4 , height = 150, width = 150 )
#Entête du programme
label2.place(relheight = 0.05,relwidth = 1,rely = 0)
label.place(relwidth = 1, rely = 0  )

# Création d'un menu sur la fenêtre

menubar = Menu(main)
menu1 = Menu(menubar, tearoff=0)
menu1.add_command(label="About kofi", command = openweb)
menu1.add_command(label="User guide", command = openweb2)
menubar.add_cascade(label="Help", menu=menu1)

#Attribution du menu au main
main.config(menu=menubar)

# Placement de l'horloge
clock = Label(main, font=('times', 20, 'bold'), bg='white', fg= bluedark)
clock.place(rely = 0.85, relx = 0.85)
hour()

main.mainloop()


from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import time
import math

def calc(x):
	return math.log(abs(12*math.sin(x)))

try:
	browser = webdriver.Chrome()
	browser.get('http://suninjuly.github.io/explicit_wait2.html')
	price = WebDriverWait(browser, 15).until(
		EC.text_to_be_present_in_element((By.ID, "price"), '$100') #Задаем явное ожидание 15 секунд или пока в елементе текст не будет равен $100
	)
	browser.find_element(By.ID, 'book').click()

	x = browser.find_element_by_id('input_value').text
	result = calc(int(x))

	field = browser.find_element_by_id("answer")
	browser.execute_script('return arguments[0].scrollIntoView(true);',field) #Выполнение кастомного js скрипта который проскролит страницу до поля ввода
	field.send_keys(str(result))

	browser.find_element(By.ID, 'solve').click()

finally:
	time.sleep(3)
	browser.quit()


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import RidgeCV
from sklearn.metrics import mean_absolute_error, mean_squared_error
from sklearn.linear_model import Ridge
from sklearn.model_selection import train_test_split, KFold
from sklearn.preprocessing import PolynomialFeatures
from sklearn.preprocessing import PolynomialFeatures

df = pd.read_csv ("final.csv")
print (df.head( ))
Goals_pG=df.iloc[:,0]
Assists_pG=df.iloc[:,1]
Yel_pG=df.iloc[:,2]
Red_pG=df.iloc[:,3]
Shot_pG=df.iloc[:,4]
Pass_success_rate=df.iloc[:,5]
AerialsWon_pG=df.iloc[:,6]
X=np.column_stack((Goals_pG,Assists_pG,Yel_pG,Red_pG,Shot_pG,Pass_success_rate,AerialsWon_pG))
Y=df.iloc [:,7]

X_train, X_test, y_train, y_test = train_test_split(X, Y, random_state=50, test_size=0.3)

mean_error=[]; std_error=[]
polynomial=[1,2]
for i in polynomial:
    po=PolynomialFeatures(degree=i)
    features_power=po.fit_transform(X)
    Xtrain=features_power
    Ytrain=Y
    model=Ridge(alpha=2,max_iter=10000)
    model.fit(Xtrain,Ytrain)
    temp=[]
    from sklearn.model_selection import KFold
    kf = KFold(n_splits=5)
    for train, test in kf.split(X):
        model.fit(Xtrain[train], Ytrain[train])
        ypred = model.predict(Xtrain[test])
        from sklearn.metrics import mean_squared_error
        temp.append(mean_squared_error(Ytrain[test],ypred))
    mean_error.append(np.array(temp).mean())
    std_error.append(np.array(temp).std())
import matplotlib.pyplot as plt
plt.rc('font', size=18)
plt.rcParams['figure.constrained_layout.use'] = True
plt.errorbar(polynomial,mean_error,yerr=std_error,linewidth=3)
plt.xlabel('pow'); plt.ylabel('Mean square error')
plt.xlim((0,3))
plt.show()

## Ridge

mean_error=[];std_error=[]
Ci_range=[0.01,0.1,0.5,1,2,3,4,5]
for Ci in Ci_range:
    from sklearn.linear_model import Ridge
    model = Ridge(alpha=1/(2*Ci))
    temp=[]
    kf=KFold(n_splits=5)
    for train,test in kf.split(X):
        model.fit(X[train],Y[train])
        ypred=model.predict(X[test])
        temp.append(mean_squared_error(Y[test],ypred))
    mean_error.append(np.array(temp).mean())
    std_error.append(np.array(temp).std())

print("temp is",temp)
print("std_error is ",std_error)

plt.rc('font', size=18)
plt.rcParams['figure.constrained_layout.use'] = True
plt.errorbar(Ci_range,mean_error,yerr=std_error,linewidth=3)
plt.title('Ridge', fontsize=25)
plt.xlabel('C')
plt.ylabel('Mean square error')
plt.show()
# alpha_range=[50,5,1,0.5,0.25,0.167,0.125,0.1]
# ridgeCV= RidgeCV(alphas=alpha_range,cv=5).fit(X,Y)
# best_alpha = ridgeCV.alpha_
# print("For Ridge, best_alpha is",best_alpha)
# print("best C is",1/(2*best_alpha))

ridge=Ridge(alpha=0.25)
po=PolynomialFeatures(degree=2)
X_train_2=po.fit_transform(X_train)
X_test_2=po.fit_transform(X_test)
X_2=po.fit_transform(X)
ridge.fit(X_train_2,y_train)
y_tr=ridge.predict(X_train_2)
y_pred=ridge.predict(X_test_2)
y=ridge.predict(X_2)
print('Ridge coefficients: ',ridge.coef_)
print('train Mean Absolute Error:', mean_absolute_error(y_train, y_tr))
print('test Mean Absolute Error:', mean_absolute_error(y_test, y_pred))
print('train Mean Squared Error:', mean_squared_error(y_train, y_tr))
print('test Mean Squared Error:', mean_squared_error(y_test, y_pred))
# print('Root Mean Squared Error:', np.sqrt(mean_squared_error(y_test, y_pred)))
print('Ridge train score: ',ridge.score(X_train_2,y_train))
print('Ridge test score: ',ridge.score(X_test_2,y_test))
plt.ylabel('Rating',fontsize=15)
plt.xlabel('Goals_pG',fontsize=15)
plt.title('Ridge', fontsize=25)
plt.scatter(X_train[:,0],y_train,color='red')
plt.scatter(X_train[:,0],y_tr,color='blue')
plt.legend(["training data","prediction"])
plt.show()
plt.ylabel('Rating',fontsize=15)
plt.xlabel('Goals_pG',fontsize=15)
plt.title('Ridge', fontsize=25)
plt.scatter(X_test[:,0],y_test,color='red')
plt.scatter(X_test[:,0],y_pred,color='blue')
plt.legend(["training data","prediction"])
plt.show()


import platform
from setuptools import setup

version = '1.4.3.dev0'

if platform.system() in ('Linux','Unix'):
    man1pages = [('share/man/man1', [
           'csv2ods/csv2ods.1',
           'mailodf/mailodf.1',
           'odf2xhtml/odf2xhtml.1',
           'odf2mht/odf2mht.1',
           'odf2xml/odf2xml.1',
           'odfimgimport/odfimgimport.1',
           'odflint/odflint.1',
           'odfmeta/odfmeta.1',
           'odfoutline/odfoutline.1',
           'odfuserfield/odfuserfield.1',
           'xml2odf/xml2odf.1'])]
else:
    man1pages = []
# Currently no other data files to add
datafiles = [] + man1pages

setup(name='odfpy',
      version=version,
      classifiers = [
        'Development Status :: 5 - Production/Stable',
        'Environment :: Console',
        'Intended Audience :: End Users/Desktop',
        'Intended Audience :: Developers',
        'Intended Audience :: System Administrators',
        'License :: OSI Approved :: Apache Software License',
        'License :: OSI Approved :: GNU General Public License (GPL)',
        'License :: OSI Approved :: GNU Library or Lesser General Public License (LGPL)',
        'Operating System :: MacOS :: MacOS X',
        'Operating System :: Microsoft :: Windows',
        'Operating System :: POSIX',
        'Programming Language :: Python :: 2',
        'Programming Language :: Python :: 2.7',
        'Programming Language :: Python :: 3',
        'Programming Language :: Python :: 3.4',
        'Programming Language :: Python :: 3.5',
        'Programming Language :: Python :: 3.6',
        'Programming Language :: Python :: 3.7',
        'Topic :: Office/Business',
        'Topic :: Software Development :: Libraries :: Python Modules',
      ],
      description='Python API and tools to manipulate OpenDocument files',
      long_description = (
"""
Odfpy is a library to read and write OpenDocument v. 1.2 files.
The main focus has been to prevent the programmer from creating invalid
documents. It has checks that raise an exception if the programmer adds
an invalid element, adds an attribute unknown to the grammar, forgets to
add a required attribute or adds text to an element that doesn't allow it.

These checks and the API itself were generated from the RelaxNG
schema, and then hand-edited. Therefore the API is complete and can
handle all ODF constructions.

In addition to the API, there are a few scripts:

- csv2odf - Create OpenDocument spreadsheet from comma separated values
- mailodf - Email ODF file as HTML archive
- odf2xhtml - Convert ODF to (X)HTML
- odf2mht - Convert ODF to HTML archive
- odf2xml - Create OpenDocument XML file from OD? package
- odfimgimport - Import external images
- odflint - Check ODF file for problems
- odfmeta - List or change the metadata of an ODF file
- odfoutline - Show outline of OpenDocument
- odfuserfield - List or change the user-field declarations in an ODF file
- xml2odf - Create OD? package from OpenDocument in XML form

The source code is at https://github.com/eea/odfpy

Visit https://github.com/eea/odfpy/wiki for documentation and examples.

The code at https://joinup.ec.europa.eu/software/odfpy/home is obsolete."""
),
      author='Soren Roug',
      author_email='soren.roug@eea.europa.eu',
      url='https://github.com/eea/odfpy',
      packages=['odf'],
      scripts=[
          'csv2ods/csv2ods',
          'mailodf/mailodf',
          'odf2xhtml/odf2xhtml',
          'odf2mht/odf2mht',
          'odf2xml/odf2xml',
          'odfimgimport/odfimgimport',
          'odflint/odflint',
          'odfmeta/odfmeta',
          'odfoutline/odfoutline',
          'odfuserfield/odfuserfield',
          'xml2odf/xml2odf'],
      data_files=datafiles,
      install_requires=['defusedxml', ]
      )


import colorama
from colorama import Fore, Style
import pyfiglet
import time

# Initialize colorama
colorama.init()
# Display the message when opening the command prompt
print(Fore.YELLOW + "Thanks for downloading our Converter!")
print("This tool is very easy to use, just follow readme.txt for all the steps!")
print("If you have any questions, feel free to message me on Telegram @blugoshop\n" + Style.RESET_ALL)

# Generate the ASCII art
ascii_art = pyfiglet.figlet_format("BluGo Token Converter")
print(Fore.BLUE + ascii_art + Style.RESET_ALL)

# Open the tokens file
with open("tokens.txt", "r") as file:
    # Prompt the user to press button 1 to start
    input(Fore.CYAN + "Press button 1 to start the token cutting process." + Style.RESET_ALL)

    # Start the timer
    start_time = time.time()

    try:
        tokens = file.readlines()

        for token in tokens:
            # Remove spaces from the token
            cut_token = token.split(":")[2].strip()
            print(token)

            # Simulate cutting a token by waiting for 1 second
            time.sleep(0)

            # Display success message in green along with the token
            print(Fore.GREEN + "Success!" + Style.RESET_ALL + " Token: " + cut_token)

            # Write the cut token to the file
            open("cut-tokens.txt", "a").write(f"{cut_token}\n")

        # Calculate the elapsed time
        elapsed_time = time.time() - start_time
        print(Fore.YELLOW + f"Finished in {elapsed_time:.2f} seconds." + Style.RESET_ALL)

    except Exception as e:
        print(Fore.RED + f"Error occurred: {e}" + Style.RESET_ALL)


from abc import ABCMeta, abstractmethod
import chess
import math
import time
import numpy as np
from collections import namedtuple

from search_helpers import quickselect, material_balance
import guerilla.data_handler as dh

# Note:
#   Moves are stored as UCI strings
#   Leaf FEN is stripped
Transposition = namedtuple("Transposition", "best_move value leaf_fen node_type")

PV_NODE = 0
CUT_NODE = 1
ALL_NODE = 2
LEAF_NODE = 3  # Additional type, we use it to demark nodes which were leafs to reduce the # of evaluation

class TranspositionTable:
    def __init__(self, exact_depth=True):
        self.table = {}  # Main transposition table {FEN: Transposition Entry}
        self.exact_depth = exact_depth  # If we require an exact depth match

        self.cache_hits = {} # Cache hits by depth
        self.cache_miss = 0
        self.num_transpositions = 0

    def __str__(self):
        return "[TT] {} entries | {} transpositions".format(len(self.table), self.num_transpositions)

    def fetch(self, fen, requested_depth):
        """
        Fetches the transposition for the input FEN
        :param fen: [String] Input FEN for which depth is queried.
        :param requested_depth: [Int] Requested depth. Effect depends on self.exact_depth:
            (True): Returns a transposi tion for which the input FEN was searched to EXACTLY requested_depth.
            (False): Return a transposition for which the input FEN was search to AT LEAST requested_depth.
        :return:
        """
        # Returns {Best Move, Value, Type, Depth}

        # Check for entry in table
        white_fen = dh.flip_to_white(fen)
        # print "Fetching ORIGINAL {} WHITE {}".format(fen, white_fen)
        entry = self._get_entry(white_fen)
        if not entry or entry.deepest < requested_depth:
            self.cache_miss += 1
            return

        # If not using exact depth, then get deepest depth
        if not self.exact_depth:
            requested_depth = entry.deepest

        # Fetch depth result from table
        if requested_depth in entry.value_dict:
            if requested_depth not in self.cache_hits:
                self.cache_hits[requested_depth]= 0
            self.cache_hits[requested_depth] += 1

            transpo = entry.value_dict[requested_depth]

            # If black is next then flip move AND leaf fen
            if dh.black_is_next(fen):
                transpo = self._flip_transposition(transpo)

            return transpo
        else:
            self.cache_miss += 1

    def update(self, fen, search_depth, best_move, score, leaf_fen, node_type):

        # Flip to white
        transpo = Transposition(best_move, score, leaf_fen, node_type)

        # Flip transposition if necessary
        if dh.black_is_next(fen):
            transpo = self._flip_transposition(transpo)

        entry = self._get_entry(dh.flip_to_white(fen), create=True)
        if search_depth not in entry.value_dict:
            self.num_transpositions += 1
        entry.add_depth(search_depth, transpo)

    def exists(self, fen):
        assert (dh.white_is_next(fen))
        return dh.strip_fen(fen) in self.table

    def _get_entry(self, fen, create=False):
        assert (dh.white_is_next(fen))
        key = dh.strip_fen(fen)
        if key not in self.table and create:
            self.table[key] = TranpositionEntry()
        return self.table.get(key)

    def _flip_transposition(self, entry):
        best_move = dh.flip_move(entry.best_move) if entry.best_move is not None else entry.best_move
        leaf_fen = dh.flip_board(entry.leaf_fen)
        return Transposition(best_move, entry.value, leaf_fen, entry.node_type)

    def clear(self):
        # Clears the transposition table
        self.table.clear()
        self.cache_hits = {}
        self.cache_miss = 0

    def get_value(self, fen):
        """
        Returns the value for the input FEN, and the depth where that value comes from.
        :param fen: [String] FEN.
        :return:
            [(Transposition, depth)] If FEN exists in cache else [(None, None)]
        """

        result = (None, None)

        # Check for entry entry
        white_fen = dh.flip_to_white(fen)
        entry = self._get_entry(white_fen)

        # Get deepest transposition with an exact value
        if entry is not None and entry.deepest_value is not None:
            transpo = entry.value_dict[entry.deepest_value]

            # Flip if necessary
            if dh.black_is_next(fen):
                transpo = self._flip_transposition(transpo)

            result = (transpo, entry.deepest_value)

        return result

class TranpositionEntry:
    def __init__(self):
        self.value_dict = {}  # {Depth: Transposition}
        self.deepest = None # Deepest transposition
        self.deepest_value = None # Deepest transposition for which we have an exact value (Leaf or PV Node)

    def add_depth(self, depth, transposition):
        # Update value dict and deepest
        self.value_dict[depth] = transposition
        self.deepest = max(depth, self.deepest)

        # Update deepest value
        if transposition.node_type == PV_NODE or transposition.node_type == LEAF_NODE:
            self.deepest_value = max(depth, self.deepest_value)

class Search:
    __metaclass__ = ABCMeta

    def __init__(self, evaluation_function):
        """
        Initializes the Player abstract class.
        Input:
            evaluation_function [function]:
                function to evaluate leaf nodes (usually the neural net)
        """
        self.evaluation_function = evaluation_function
        self.win_value = dh.WIN_VALUE
        self.lose_value = dh.LOSE_VALUE
        self.draw_value = dh.DRAW_VALUE
        self.num_evals = 0
        self.num_visits = []  # Index: depth, Value: Number of vists at that depth
        self.depth_reached = 0
        # Transposition table
        self.tt = TranspositionTable(exact_depth=True)  # TODO: We can change this

    def evaluate(self, fen):
        # Returns the value of the input FEN from the perspective of the next player to play
        # Depth is used to update information

        self.num_evals += 1

        # Evaluate
        board = chess.Board(fen)
        if board.is_checkmate():
            return self.lose_value
        elif board.can_claim_draw() or board.is_stalemate():
            return self.draw_value
        else:
            return self.evaluation_function(dh.flip_to_white(fen))

    @abstractmethod
    def __str__(self):
        raise NotImplementedError("You should never see this")

    @abstractmethod
    def run(self, board, time_limit=None, clear_cache=False):
        """
        Runs search based on parameter.
        Inputs:
                board [chess.Board]:
                    current state of board
            Outputs:
                best_move [chess.Move]:
                    Best move to play
                best_value [float]:
                    value achieved by best move
                best_leaf [String]
                    FEN of the board of the leaf node which yielded the highest value.
        """
        raise NotImplementedError("You should never see this")

    def reset(self):
        # Clears the transposition table
        self.tt.clear()

        # Reset some logging variables
        self.num_evals = 0
        self.num_visits = []
        self.depth_reached = 0

class IterativeDeepening(Search):
    """
    Searches game tree in an Iterative Deepening Depth search.
    At each depth optionally prune from remaining possibilities
    Implements alpha_beta pruning by default
    """

    def __init__(self, evaluation_function, time_limit=10,
                 max_depth=None, h_prune=False, prune_perc=0.0,
                 ab_prune=True, verbose=True, use_partial_search=False):

        """
            Constructor
            Inputs:
                evaluation_function[function]:
                    function used to evaluate leaf nodes
                time_limit[float]:
                    time limit per move
                max_depth[int]:
                    If not None, limit depth to max_depth
                h_prune[bool]:
                    Heuristic_prune. If true, prune between between depth-limited-searches
                prune_perc[float range([0,1])]:
                    Percent of nodes to prune for heuristic prune
                ab_prune[bool]:
                    Alpha-beta pruning. Same results on or off, only set to off for td training
                use_partial_search [Bool]:
                    Whether or not to use partial search results (i.e. when a timeout occurs during DFS).
        """
        super(IterativeDeepening, self).__init__(evaluation_function)

        self.time_limit = time_limit
        self.buff_time = time_limit * 0.02
        self.depth_limit = 1  # Depth limit for DFS
        self.max_depth = max_depth
        self.order_moves = True  # Whether moves should be ordered
        self.h_prune = h_prune
        self.prune_perc = prune_perc
        self.ab_prune = ab_prune
        self.root = None  # holds root node
        self.is_partial_search = False  # Marks if the last DFS call was a partial search
        self.use_partial_search = use_partial_search

        # Holds the Killer Moves by depth. Each Entry is (set of moves, sorted array of (value, moves)).
        self.killer_table = None
        self.num_killer = 2  # Number of killer moves store for each depth
        self._reset_killer()

        # Move value for ordering when board not found in transposition table
        self.order_fn_fast = material_balance

    def __str__(self):
        return "IterativeDeepening"

    def _uci_to_move(self, uci):
        return chess.Move.from_uci(uci) if uci is not None else uci

    def DLS(self, node, alpha, beta):
        """
        Recursive depth limited negamax search with alpha_beta pruning.
        Source: https://en.wikipedia.org/wiki/Negamax#Negamax_with_alpha_beta_pruning_and_transposition_tables
        :param node: [SearchNode] Roote node for search.
        :param alpha: [Float] Lower bound.
        :param beta: [Float] Upper bound.
        :return:
            best_score [float]:
                Score achieved by best move
            best_move [chess.Move]:
                Best move to play
            best_leaf [String]
                FEN of the board of the leaf node which yielded the highest value.
        """

        # Track some info
        if not node.visited:
            # We've never seen this node before -> Track some info
            try:
                self.num_visits[node.depth] += 1
            except IndexError:
                self.num_visits.append(1)

            self.depth_reached = max(self.depth_reached, node.depth)

            # Check if a new killer table entry should be created
            if node.depth >= len(self.killer_table):
                self.killer_table.append({'moves': set(), 'values': list()})

            node.visited = True

        alpha_original = alpha

        # Check transposition table
        result = self.tt.fetch(node.fen, requested_depth=self.depth_limit - node.depth)
        if result:
            if result.node_type == PV_NODE:
                return result.value, self._uci_to_move(result.best_move), result.leaf_fen
            elif result.node_type == CUT_NODE:
                # lower bound
                alpha = max(alpha, result.value)
            elif result.node_type == ALL_NODE:
                # upper bound
                beta = min(beta, result.value)

            if alpha >= beta:
                return result.value, self._uci_to_move(result.best_move), result.leaf_fen

        # Check children
        if not node.children:
            node.gen_children()

        # Check if limit reached
        self.time_left = (time.time() - self.start_time) <= self.time_limit - self.buff_time
        if node.depth == self.depth_limit or not node.expand or not self.time_left or not node.children:
            # Evaluate IF: depth limit reached, pruned, no time left OR no children
            if not self.time_left:
                self.is_partial_search = True

            # Check if we have previously evaluated this node as a leaf node
            if result and result.node_type == LEAF_NODE:
                return result.value, None, node.fen

            # Otherwise evaluate node
            node.value = self.evaluate(node.fen)

            # Update transposition table
            self.tt.update(node.fen, 0, None, node.value, node.fen, LEAF_NODE)

            return node.value, None, node.fen

        # Get Ordered children
        moves = self.get_ordered_moves(node) if self.order_moves else node.child_moves

        # Find best move (recursive)
        best_value = float("-inf")
        best_move = None
        leaf_fen = None

        for move in moves:
            # Get best score for opposing player and flip it to your perspective
            value, _, lf = self.DLS(node.children[move], alpha=-beta, beta=-alpha)
            value = -value

            if value > best_value:
                best_move = move
                best_value = value
                leaf_fen = lf

            # Check for pruning
            alpha = max(alpha, value)
            if alpha >= beta:
                self.update_killer(move, value, node.depth)
                break

        # Update transposition table
        if best_value <= alpha_original:
            # ALL NODES searched, no good moves found -> value is an upper bound
            node_type = ALL_NODE
        elif best_value >= beta:
            # CUT NODE, pruning occurred -> value is a lower bound
            node_type = CUT_NODE
        else:
            # PV NODE Otherwise its potentially part of the the principal variation
            node_type = PV_NODE

        # Update transposition table
        self.tt.update(node.fen, self.depth_limit - node.depth, best_move, best_value, leaf_fen, node_type)

        # Return result of search
        return best_value, self._uci_to_move(best_move), leaf_fen

    def get_ordered_moves(self, node):
        """
        Orders the child moves of the node.
        Ordering is based on:
            (1) Killer moves
            (2) Moves for which we have a value, ordering by (-depth, value) in increasing order
            (3) Other moves
        :param node: [SearchNode] Node who's child moves we need to order.
        :return: [List of Strings] Ordered moves
        """
        killer_moves = []  # Children that are "killer" moves
        value_moves = [] # Moves with values
        other_moves = []

        for move in node.child_moves:
            child_fen = node.children[move].fen

            # Favor killer moves
            if self.is_killer(move, node.depth):
                killer_moves.append((move, self.order_fn_fast(child_fen)))
                continue

            # Check if we have an estimate for the move value
            #   Assign it to a group accordingly
            transpo, depth = self.tt.get_value(child_fen)
            if transpo:
                # Note: take negative of depth since want to look at moves scored deeper first
                value_moves.append((move, (-depth, transpo.value)))
            else:
                other_moves.append((move, self.order_fn_fast(child_fen)))

        # Order in ascending order (want to check boards which are BAD for opponent first)
        killer_moves.sort(key=lambda x: x[1])
        value_moves.sort(key=lambda x: x[1])
        other_moves.sort(key=lambda x: x[1])

        moves = killer_moves + value_moves + other_moves

        assert(len(moves) == len(node.child_moves))

        move_order = [x[0] for x in moves]

        return move_order

    def update_killer(self, killer_move, value, depth):
        """
        Updates the killer move table.
        Input:
            killer_move [Chess.Move]
                The move which caused the A-B pruning to trigger.
            value [Float]
                The value yielded by the killer move.
            depth [Int]
                The depth FROM which the move was played.
        """

        k_tab = self.killer_table[depth]

        # Skip if already in killer table
        if killer_move in k_tab['moves']:
            return

        # Check if table is full
        if len(k_tab['moves']) < self.num_killer:
            # Not Full
            self._add_killer_move(depth, value, killer_move)
        else:
            # Full
            # Check if move is better than worst current move
            if value > k_tab['values'][0]:
                # Remove Worst
                _, worst_move = k_tab['values'].pop(0)
                k_tab['moves'].remove(worst_move)

                # Add Item
                self._add_killer_move(depth, value, killer_move)

    def _add_killer_move(self, depth, value, killer_move):
        """
        Adds killer move to the table.
        """
        # Update moves
        self.killer_table[depth]['moves'].add(killer_move)

        # Update values
        self.killer_table[depth]['values'].append((value, killer_move))
        self.killer_table[depth]['values'].sort(key=lambda x: x[0])  # Sorting each time is OK since length is small.

    def is_killer(self, move, depth):
        """
        Checks if the current move is a killer move.
        Input:
            move [chess.Move]
                Move to check.
        Output:
            output [Boolean]
                True if it is a kill move, False if not.
        """
        return move in self.killer_table[depth]['moves']

    def _reset_killer(self):
        """
        Resets the killer moves table.
        :return:
        """
        self.killer_table = [{'moves': set(), 'values': list()}]

    def prune(self, node):
        """
            Recursive pruning of nodes
        """
        if not node.expand or not node.children:
            return

        children = list(node.get_child_nodes())
        # k = number of nodes that I keep
        k = max(min(len(children), 2),
                int(math.ceil(len(children) * (1 - self.prune_perc))))
        quickselect(children, 0, len(children) - 1, k - 1, key=lambda x: x.value)

        for child in children[:k]:
            self.prune(child)
            child.expand = True
        for child in children[k:]:
            child.expand = False

    def run(self, board, time_limit=None, reset=False):
        """
            For the duration of the time limit and depth limit:
                1. Depth Limited Search
                2. If enabled: Prune nodes
                3. Increase max depth
            Inputs:
                board[chess.Board]:
                    Chess board to search the best move for
                time_limit[float]:
                    time limit for search. If None, defaults to time_limit set in init
                reset [bool]:
                    Resets the search instance. Used for training.
            Outputs:
                best_value [float]:
                    value achieved by best move
                best_move [chess.Move]:
                    Best move to play
                best_leaf [String]
                    FEN of the board of the leaf node which yielded the highest value.
        """
        if time_limit is None:
            time_limit = self.time_limit

        if reset:
            self.reset()
        self.num_evals = 0
        self.eval_time = 0
        self.num_visits = []
        self._reset_killer()

        # Start timing
        if time_limit is not None:
            self.time_limit = time_limit
        self.start_time = time.time()
        self.time_left = True

        self.root = SearchNode(board.fen(), 0)

        self.depth_limit = 1
        value = best_move = leaf_board = None
        while self.time_left and (self.max_depth is None or self.depth_limit <= self.max_depth):

            # Run search
            new_results = self.DLS(self.root, alpha=float("-inf"), beta=float("inf"))
            if not self.is_partial_search or (self.is_partial_search and self.use_partial_search):
                value, best_move, leaf_board = new_results
            self.is_partial_search = False

            # Prune if necessary
            if self.h_prune:
                self.prune(self.root)

            # Increase depth
            self.depth_limit += 1

        return value, best_move, leaf_board

class SearchNode:
    def __init__(self, fen, depth, value=None):
        """
        Generic node used for searching in 'rank_prune'.
        Input:
            fen [String]
                FEN of the board.
            depth [Int]
                Depth at which the node occurs.
            value [Float] (Optional)
                Value of the node.
        Non-Input Class Attributes
            children [Dict of SearchNode's]
                Children of the current SearchNode, key is Move which yields that SearchNode.
        """
        self.fen = fen
        self.depth = depth
        self.children = {}
        self.child_moves = []
        self.expand = True
        self.visited = False

    def _add_child(self, move, child):
        assert (isinstance(child, SearchNode))

        # Convert move to uci
        uci_move = move.uci()

        self.children[uci_move] = child
        self.child_moves.append(uci_move)

    def gen_children(self):
        # Create children
        board = chess.Board(self.fen)
        for move in board.legal_moves:
            board.push(move)
            self._add_child(move, SearchNode(board.fen(), self.depth + 1))
            board.pop()

    def get_child_nodes(self):
        """
        Returns a list of child nodes.
        """
        return self.children.values()

    def __str__(self):
        return "Node[{}, {}, {} children]".format(self.fen, self.depth, len(self.children))

class Minimax(IterativeDeepening):
    """
        Uses a recursive function to perform a simple minimax with
        alpha-beta pruning.
    """

    def __init__(self, leaf_eval, max_depth=2, ab_prune=True):
        """
            Constructor
            Inputs:
                leaf_eval[function]:
                    function used to evaluate leaf nodes
                max_depth[int]:
                    depth to go to
                ab_prune[bool]:
                    Alpha-beta pruning. Same results on or off, only set to off for td training
        """
        super(Minimax, self).__init__(leaf_eval, time_limit=np.inf, max_depth=max_depth,
                                      prune_perc=0.0, ab_prune=ab_prune, use_partial_search=False)

        # Set depth limit to max depth
        self.depth_limit = max_depth

        # Turn off propagation
        self.use_prop = False

        # No buffer time needed
        self.buff_time = 0

    def __str__(self):
        return "Minimax"

    def run(self, board, time_limit=None, reset=False):
        """ Reset some variables, call recursive function """
        if reset:
            self.reset()
        self.num_evals = 0
        self.eval_time = 0
        self.num_visits = []
        self.time_left = True

        # Run to depth limit
        self.start_time = 0
        self.root = SearchNode(board.fen(), 0)  # Note: store unflipped fen
        return self.DLS(self.root, alpha=float("-inf"), beta=float("inf"))


# ====================================================================================================
#  Brief:  Fitting code and likelihood definitions, specific to the "one bkg, one sig" case studied here
#  Author: Stephen Menary (stmenary@cern.ch)
# ====================================================================================================

import math

import numpy as np
from iminuit import Minuit
from scipy.stats import poisson, chi2
from scipy.special import gamma
import matplotlib.pyplot as plt

import Confidence.messaging as msg
import Confidence.stats as stats
from .Models.Models import Model
import Confidence.Plotting.plotting as plotting

#  Global: minimisation model
model_to_fit = None
data_to_fit = None
cov_to_fit = None
inv_cov_to_fit = None
l_norm = None

'''
#          This is the likelihood following a chi2 distribution
#  Likelihood
#
def likelihood (params) :
	if len(params) is not 2 :
		msg.fatal("Confidence.fitting.likelihood","Likelihood can only be evaluated for parameter lists of length 2, whereas {0} was provided".format(len(params)))
	global model_to_fit, data_to_fit, inv_cov_to_fit, l_norm
	model_to_fit.coefficients = params
	pred = model_to_fit.generate_prediction()
	res = data_to_fit - pred
	v = np.matmul(res, np.matmul(inv_cov_to_fit, res))
	return l_norm * np.exp(-0.5*v)
	return chi2.pdf(v, len(res))   #   doesn't work -- why?
'''

#          This is the likelihood calculated as the product of Poissons (unbiased, but slow, often results in overflow errors for large values, and only works for integer k unless we approximate k! as Gamma(k))
#  Likelihood
#
def likelihood (params) :
	if len(params) is not 2 :
		msg.fatal("Confidence.fitting.likelihood","Likelihood can only be evaluated for parameter lists of length 2, whereas {0} was provided".format(len(params)))
	global model_to_fit, data_to_fit, inv_cov_to_fit, l_norm
	model_to_fit.coefficients = params
	pred = model_to_fit.generate_prediction()
	L = 1.
	for i in range(len(pred)) :
		mu = pred[i]
		k = data_to_fit[i]
		if k % 1 == 0 :
			L = L * poisson.pmf(k, mu)
		else :
			P = np.exp(-1.*mu) * (mu ** k) / gamma(k)
			L = L * P
	return L

'''
                 #   L = L * poisson.pmf(data_to_fit[i], pred[i])
'''

#  -2logL
#
def TNLL (params, **kwargs) :
	l = likelihood(params)
	return -2. * np.log(l)

#  TNLL --> likelihood
#
def TNLL_to_likelihood (tnll, **kwargs) :
	return np.exp(-0.5*tnll)

#  Do fit and return the test statistic q ( q = L(s+b) / L(b) )
#
def do_fit (values, errors, fix, **kwargs) :
	m = Minuit.from_array_func(
		TNLL,
		values,
		error = errors,
		errordef = 1.,
		fix = fix,
		limit = tuple( (None,None) for i in range(len(values)) ),
		name = tuple( ["mu_bkg" , "mu_sig"] ),
		pedantic=True)
	m.migrad()
	if kwargs.get("minos",False) is True :
		m.minos()
	return m

#  Do fit and return the minimised TNLL
#
def get_fitted_TNLL (values, errors, fix) :
	m = do_fit(values, errors, fix)
	return m.get_fmin().fval

#  Do fit and return the minimised TNLL as well as mu_sig and it's error
#
def get_fitted_TNLL_and_signal_strength (values, errors, fix, **kwargs) :
	m = do_fit(values, errors, fix)
	if "profile" in kwargs :
		m.draw_profile(kwargs["profile"], subtract_min=True)
	return m.get_fmin().fval, m.values["mu_sig"], m.errors["mu_sig"]

#  Mkae copy of model and set it as the fit model
#
def set_fit_model (model) :
	global model_to_fit
	model_to_fit = Model(model.templates)

#  Set global data and covariance information
#
def set_data_and_covariance (values, errs, correlation) :
	global data_to_fit, cov_to_fit, inv_cov_to_fit, l_norm
	data_to_fit = values
	cov_to_fit = np.zeros(shape=correlation.shape,dtype=np.double)
	for i in range(len(errs)) :
		for j in range(len(errs)) :
			cov_to_fit[i][j] = errs[i] * errs[j] * correlation[i][j]
	inv_cov_to_fit = np.linalg.inv(cov_to_fit)
	l_norm = np.fabs(np.linalg.det(cov_to_fit)) * ( (2*math.pi) ** len(values) )
	l_norm = 1. / np.sqrt(l_norm)

#  Expected uncertainty on sig
#
def get_expected_mu_sig_error (model, correlation) :
	set_fit_model(model)
	y, ey = model.generate_asimov()
	set_data_and_covariance(y, ey, correlation)
	m = do_fit((1.,1.), (0.5,0.5), (False,False))
	return m.errors["mu_sig"]

#  Throw toys around a model, fit them and return the distribution of q = L(s+b) / L(b). Plot the pulls on mu_sig if desired.
#
def throw_and_fit_toys (model, correlation, n_toys, **kwargs) :
	set_fit_model(model)
	mu_true = model.coefficients[1]
	toys_q = []
	pulls_mu_sig = []
	for i in range(n_toys) :
		if i>0 and (10*i) % n_toys == 0 :
			msg.info("Confidence.DM_fitting.throw_and_fit_toys", "Processing toy {0} out of {1}".format(i,n_toys), verbose_level=1)
		n_tries = 0
		while n_tries > -1 :
			if n_tries == 1000 : msg.fatal("Confidence.DM_fitting.throw_and_fit_toys","Covariance matrix is still singular after 1000 attempts")
			try :
				y, ey = model.throw_toy()
				set_data_and_covariance(y, ey, correlation)
				n_tries = -1
			except Exception as e :
				msg.warning("Confidence.DM_fitting.throw_and_fit_toys","Error {0} occurred when throwing toy. Retrying.".format(e),verbose_level=1)
				n_tries = n_tries + 1
		tnll_s_b, mu_sig, mu_sig_err = get_fitted_TNLL_and_signal_strength( model.coefficients, (0.5,0.5), (False,False) )
		if kwargs.get("plot_first_profile",False) is True :
			x, y = [], []
			for i in range(21) :
				this_x = mu_sig + (i*0.2 - 2.0)*mu_sig_err
				x.append(this_x)
				y.append( get_fitted_TNLL( (1.0,this_x), (0.5,0.5), (False,True) ) )
			fig  = plt.figure(figsize=(7,7))
			ax = fig.add_subplot(111)
			ax.plot(x, y, "-", c="r")
			plt.show()
			if type(plotting.document) is not None :
				fig.savefig(plotting.document, format='pdf')
			plt.close(fig)
		tnll_b = get_fitted_TNLL( (1.,0.), (0.5,0.5), (False,True)  )
		q = TNLL_to_likelihood(tnll_s_b) / TNLL_to_likelihood(tnll_b)
		toys_q.append( float(q) )
		pulls_mu_sig.append( (mu_sig-mu_true) / mu_sig_err )
	toys_q.sort()
	mean, sem, sd, se = stats.get_mean_sem_sd_se(pulls_mu_sig)
	msg.info("Confidence.DM_fitting.throw_and_fit_toys", "    Pulls mean  =  {0:.4f}  +/-  {1:.4f}".format(mean, sem))
	msg.info("Confidence.DM_fitting.throw_and_fit_toys", "    Pulls sd    =  {0:.4f}  +/-  {1:.4f}".format(sd, se))
	if kwargs.get("plot_pulls",False) is False :
		return toys_q
	msg.info("Confidence.DM_fitting.throw_and_fit_toys","Plotting pulls for mu_true = {0:.5f}".format(mu_true))
	plotting.plot_histo_from_data(400, [pulls_mu_sig], xmin=-5, xmax=5, ymin=0., xlabel=kwargs.get("pulls_xlabel","Pull"), ylabel="Num. toys", labels=[kwargs.get("pulls_label","")], overlay_gauss=True)
	return toys_q


from settings import *
import pygame as pg
import pygame.freetype as ft
from modules.player import Player
from modules.floor import Floor
import sys

class App:
    def __init__(self):
        pg.init()
        self.screen = pg.display.set_mode(WIN_SIZE)
        self.clock = pg.time.Clock()
        self.font = ft.SysFont('Verdana', FONT_SIZE)
        pg.display.set_caption('Forbidden Treasures PyGame')
        self.load_modules()
        self.start_time = 0
        self.game_over = False
        self.menu = True
        self.score = 0
    def increase_score(self, amount):
        self.score += amount
    def draw_score(self):
        score_str = f'Score: {self.score}'
        score_surface, score_rect = self.font.render(score_str, fgcolor='white', bgcolor='black')
        score_pos = (10, 10)
        self.screen.blit(score_surface, score_pos)

    def check_events(self):
        for event in pg.event.get():
            if event.type == pg.QUIT or (event.type == pg.KEYDOWN and event.key == pg.K_ESCAPE):
                pg.quit()
                sys.exit()
            elif event.type == pg.KEYDOWN:
                if self.menu:
                    self.menu = False
                    self.start_time = pg.time.get_ticks()
                else:
                    self.player.check_events(event)

    def load_modules(self):
        self.player = Player()
        self.floor = Floor(19, 13, 1)

    def update(self):
        if not self.menu:
            self.player.update(self.floor.getFloors())

    def draw_fps(self):
        fps = f'{self.clock.get_fps():.0f} FPS'
        fps_surface, fps_rect = self.font.render(fps, fgcolor='white', bgcolor='black')
        fps_pos = (10, 10)
        self.screen.blit(fps_surface, fps_pos)

    def draw_timer(self):
        elapsed_time = (pg.time.get_ticks() - self.start_time) // 1000
        remaining_time = max(0, 120 - elapsed_time)
        time_str = f'Oxigeno: {remaining_time}s'
        time_surface, time_rect = self.font.render(time_str, fgcolor='white', bgcolor='black')
        time_pos = (self.screen.get_width() - time_rect.width - 10, 10)
        self.screen.blit(time_surface, time_pos)

    def draw_menu(self):
        self.screen.fill('black')
        title_str = 'Forbidden Treasures PyGame'
        title_surface, title_rect = self.font.render(title_str, fgcolor='white', bgcolor='black')
        title_pos = ((self.screen.get_width() - title_rect.width) // 2, 100)
        self.screen.blit(title_surface, title_pos)
        start_str = 'Press any key to start'
        start_surface, start_rect = self.font.render(start_str, fgcolor='white', bgcolor='black')
        start_pos = ((self.screen.get_width() - start_rect.width) // 2, 300)
        self.screen.blit(start_surface, start_pos)
        pg.display.flip()

    def draw(self):
        if self.menu:
            self.draw_menu()
        else:
           self.screen.fill('black')
           self.player.draw(self.screen)
           self.floor.draw(self.screen)
           self.draw_timer()
           fps_surface, fps_rect = self.font.render(f'{self.clock.get_fps():.0f} FPS', fgcolor='white', bgcolor='black')
           score_surface, score_rect = self.font.render(f'Score: {self.score}', fgcolor='white', bgcolor='black')

           fps_pos = (10, self.screen.get_height() - fps_rect.height - 10)
           score_pos = (fps_pos[0] + fps_rect.width + 20, self.screen.get_height() - score_rect.height - 10)

           self.screen.blit(fps_surface, fps_pos)
           self.screen.blit(score_surface, score_pos)

           pg.display.flip()

    def run(self):
        while True:
            self.check_events()
            self.update()
            self.draw()
            self.clock.tick(60)

            elapsed_time = (pg.time.get_ticks() - self.start_time) // 1000
            if elapsed_time >= 120:
                self.game_over = True

            if self.game_over:
                self.game_over_screen()
                break

    def game_over_screen(self):
        game_over_image = pg.image.load('assets/sprites/playerMuerte.png').convert_alpha()
        game_over_rect = game_over_image.get_rect(center=self.screen.get_rect().center)
        self.screen.blit(game_over_image, game_over_rect)
        pg.display.flip()
        pg.time.wait(3000)  # Esperar 3 segundos antes de salir
        pg.quit()
        sys.exit()

if __name__ == '__main__':
    app = App()
    app.run()


from starlette.testclient import TestClient
from pyruby_backend import main

def test_app():
    client = TestClient(main.app)
    response = client.post("/graphql/", json={'query': '{ hello }'})
    assert response.status_code == 200
    assert "hello" in response.json()["data"].keys()

def test_user_registration_journey(admin_user, base_user):
    client = TestClient(main.app)
    create_user_gql = """
    mutation createUser($email: String!, $name: String!) {
        createUser(input: {primaryEmail: $email, fullName: $name}) {
            id
            primaryEmail
            fullName
        }
    }
    """
    user = _run_gql(client, create_user_gql, {"email": base_user['email'], "name": "Bobby McBob"}, "createUser", admin_user)
    assert user["fullName"] == "Bobby McBob"

    invite_gql = """
    mutation sendInvitation($userId: ID!, $email: String!) {
        sendInvitation(input: {userId: $userId, email: $email}) {
            id
            user {
                fullName
            }
            email
        }
    }
    """

    invitation = _run_gql(client, invite_gql, {"userId": user['id'], "email": user['primaryEmail']}, 'sendInvitation', admin_user)
    assert invitation['user']['fullName'] == 'Bobby McBob'

    accept_invitation_gql = """
        mutation acceptInvitation($invitationId: ID!) {
            acceptInvitation(input: { invitationId: $invitationId}) {
                uid
                email
                user {
                    id
                    fullName
                }
            }
        }
    """

    new_login = _run_gql(client, accept_invitation_gql, {"invitationId": invitation['id']}, 'acceptInvitation', base_user)
    assert new_login['email'] == base_user['email']

    get_me_gql = """
    query me {
        me {
            id
            fullName
            primaryEmail
        }
    }
    """
    me = _run_gql(client, get_me_gql, {"cacheBuster": base_user['localId']}, 'me', base_user)
    assert me['primaryEmail'] == base_user['email']
    assert me['fullName'] == 'Bobby McBob'

def test_only_admin_users_can_create_new_users(base_user):
    client = TestClient(main.app)
    create_user_gql = """
        mutation createUser($email: String!, $name: String!) {
            createUser(input: {primaryEmail: $email, fullName: $name}) {
                id
                primaryEmail
                fullName
            }
        }
        """
    response = client.post("/graphql/", headers={"authorization": f"Bearer {base_user['idToken']}"}, json={
        'query': create_user_gql,
        'operationName': "createUser",
        'variables': {"email": base_user['email'], "name": "Bobby McBob"}
    })
    assert response.json().get("errors") is not None
    assert response.json().get("errors")[0]["path"][0] == "createUser"

def _run_gql(client, query, args, operation, user):
    response = client.post("/graphql/", headers={"authorization": f"Bearer {user['idToken']}"}, json={
        'query': query,
        'operationName': operation,
        'variables': args
    })
    assert response.json().get("errors", None) is None
    assert response.status_code == 200
    assert response.json()["data"][operation] is not None
    return response.json()["data"][operation]


import json
import os
from distutils.util import strtobool
from pathlib import Path
from typing import List, Optional, Union

import click
from anymotion_sdk import Client, ClientValueError
from anymotion_sdk.session import HttpSession

from . import __version__
from .exceptions import ClickException, SettingsValueError
from .output import echo_request, echo_response, echo_warning
from .settings import Settings
from .state import State

def get_client(state: State) -> Client:
    """Get client from state."""
    settings = get_settings(state.profile)
    if not settings.is_ok:
        command = click.style(f"{state.cli_name} configure", fg="cyan")
        message = (
            "The credentials is invalid or not set. "
            f"Run {command} to set credentials."
        )
        raise ClickException(message)

    session = HttpSession()
    if hasattr(session, "user_agent"):
        session.user_agent = f"{state.cli_name}/{__version__}"

    try:
        client = Client(
            client_id=str(settings.client_id),
            client_secret=str(settings.client_secret),
            api_url=settings.api_url,
            interval=settings.interval,
            timeout=settings.timeout,
            session=session,
        )
    except ClientValueError as e:
        raise ClickException(str(e))

    if state.verbose:
        client.session.add_request_callback(echo_request)
        client.session.add_response_callback(echo_response)

    # TODO: refactor
    state.is_download = settings.is_download
    state.is_open = settings.is_open

    return client

# TODO: remove?
def get_settings(profile: str, use_env: bool = True) -> Settings:
    """Get settings from profile."""
    try:
        settings = Settings(profile, use_env=use_env)
    except SettingsValueError as e:
        raise ClickException(str(e))
    return settings

def parse_rule(rule: Optional[str]) -> Optional[Union[list, dict]]:
    """Convert string type analysis and drawing rules to list or dict type."""
    if rule is None:
        return rule

    try:
        rule = json.loads(rule)
    except json.decoder.JSONDecodeError:
        raise ClickException("Rule format is invalid. Must be in JSON format.")

    if not isinstance(rule, (list, dict)):
        message = "Rule format is invalid. Must be in list or object format."
        raise ClickException(message)

    return rule

def echo_invalid_option_warning(condition: str, target_options: List[str]) -> None:
    """Output warning message."""
    if len(target_options) == 0:
        return
    elif len(target_options) == 1:
        be, pronoun, s = "is", "This", ""
    else:
        be, pronoun, s = "are", "These", "s"
    opt_str = ", ".join([f"'{option}'" for option in target_options])
    message = (
        f"{opt_str} {be} only available when {condition}. "
        f"{pronoun} option{s} {be} ignored.\n"
    )
    echo_warning(message)

def color_id(number: int) -> str:
    """Set color to ID."""
    return click.style(str(number), fg="cyan")

def color_path(path: Union[str, Path]) -> str:
    """Set color to path."""
    return click.style(str(path), fg="blue")

def get_bool_env(key: str, default: bool) -> bool:
    """Get boolean environment variable."""
    str_env = os.getenv(key)
    if str_env:
        try:
            bool_env = bool(strtobool(str_env))
        except ValueError:
            bool_env = default
    else:
        bool_env = default
    return bool_env


from tables.developer import Developer
from django.db import models

class APIManagement(models.Model):
    developer = models.ForeignKey(
        Developer,
        related_name='api_management',
        on_delete=models.CASCADE
        )
    key = models.CharField(max_length=64)
    quota = models.BigIntegerField()

    def __str__(self):
        return "{} - {} - {}".format(
            self.pk,
            self.developer.user.username,
            self.quota
            )


from pyserini.search import LuceneSearcher
from typing import Optional, List
import logging
import json
from abc import ABCMeta, abstractmethod

logger = logging.getLogger(__name__)

class PyseriniRetriever:
    def __init__(self, index_location: str, use_bm25: bool = True, k1: float = float(0.9), b: float = float(0.4)):
        """
        Initialize Pyserini retriever

        Args:
            index_path (str): Path to a Pyserini index
            use_bm25 (bool, optional): set BM25 as the scoring function. Defaults to True.
            k1 (float, optional): bm25 parameter to tune impact of term frequency Defaults to float(0.9).
            b (float, optional): bm25 constant to fine tune the effect of document length   Defaults to float(0.4).
        """
        self.index_location = index_location
        self.searcher = LuceneSearcher(index_location)
        self.searcher.set_bm25()
        if use_bm25:
            self.searcher.set_bm25(k1=k1,b=b)
        self.topK = 10
        logger.info(f'Initialized LuceneSearcher index_dir: {self.searcher.index_dir}  num_docs: {self.searcher.num_docs} use_bm25: {use_bm25} k1: {k1} b: {b}')

    def retrieve(self, query: str, topK: Optional[int] = 10):
        """

        Run queries against the index to retrieve ranked list of documents
        Return documents that are most relevant to the query.

        Args:
             query: search
             top_k: number of hits to return, defaults to 10

        Returns:
             List of hits, each hit is a dict containing :
             {
                "rank": i,
                "score": hit.score,
                "doc_id": docid,
                "title": title,
                "text": text
            }

        """

        hits = self.searcher.search(query, topK)
        search_results = self._collect_hits(hits)
        return search_results

    def batch_retrieve(self,  queries: List[str], qids: List[str], topK: int = 10, threads: int = 1):

        """
           Run a batch of queries

           Args:
                queries:  list of query strings
                qids:     list of qid strings corresponding to queries
                top_k:    number of hits to return, defaults to 10
                threads:  maximum number of threads to use

            Returns:
                Dict of qid to hits

        """

        hits = self.searcher.batch_search(queries, qids, k=topK, threads=threads)
        query_to_hits = {}
        for q, hits in hits.items():
            query_to_hits[q] = self._collect_hits(hits)
        return query_to_hits

    def _collect_hits(self, hits: List):
        search_results = []
        for i, hit in enumerate(hits):
            title, text = json.loads(hit.raw)['contents'].split("\t")
            title = title.replace('\n',' ')
            text = text.replace('\n',' ')
            docid = hit.docid
            search_result = {
                "rank": i,
                "score": hit.score,
                "doc_id": docid,
                "title": title,
                "text": text
            }
            search_results.append(search_result)
        return search_results