from time import sleep
import threading
import tkinter as tk
from tkinter import messagebox
from customtkinter import *
import cv2
import imutils
from PIL import Image
from PIL import ImageTk
import time                             #Importamos el paquete time
import math
import datetime
from tktooltip import ToolTip
import sqlite3
import matplotlib.pyplot as plt
import os.path

from w1thermsensor import W1ThermSensor #Importamos el paquete W1ThermSensor
import bme280
import smbus2

import board
import busio
import adafruit_ads1x15.ads1115 as ADS
from adafruit_ads1x15.analog_in import AnalogIn

import RPi.GPIO as GPIO
import numpy as np

set_appearance_mode("Dark")
set_default_color_theme("dark-blue")

window = CTk()
window.title('Hidroponia')
width = window.winfo_screenwidth()
height = window.winfo_screenheight()
window.geometry("%dx%d" %(width,height))

def cerrar_interfaz():
    continuous_threading.shutdown(0)
    leer_datos()
    GPIO.cleanup()
    window.destroy()

frame1 = CTkFrame(master=window, width=1800, height= 200)
frame1.pack(fill=BOTH, side=LEFT, expand=True)
frame2 = CTkFrame(master=window, width=100)
frame2.pack(fill=BOTH, side=LEFT)
frame2.rowconfigure(0,weight=4)
frame2.rowconfigure(1,weight=1)#label ph
frame2.rowconfigure(2,weight=1)#cuadro ph
frame2.rowconfigure(3,weight=1)
frame2.rowconfigure(4,weight=1)#temp
frame2.rowconfigure(5,weight=1)
frame2.rowconfigure(6,weight=1)
frame2.rowconfigure(7,weight=1)#hum
frame2.rowconfigure(8,weight=1)
frame2.rowconfigure(9,weight=1)
frame2.rowconfigure(10,weight=1)#temp
frame2.rowconfigure(11,weight=1)
frame2.rowconfigure(12,weight=1)
frame2.rowconfigure(13,weight=1)#cond
frame2.rowconfigure(14,weight=1)
frame2.rowconfigure(15,weight=2000)
frame2.columnconfigure(0,weight=1)
frame2.columnconfigure(1,weight=1)
frame3 = CTkFrame(master=window, width=100)
frame3.pack(fill=BOTH, side=LEFT)
framebotones = CTkFrame(frame1)
framebotones.pack(side=TOP)
frameSelPlanta = CTkFrame(frame1)
frameSelPlanta.pack(side=TOP)

espacio1=CTkLabel(frame2,text="")
espacio1.grid(row=3,column=0)
espacio2=CTkLabel(frame2,text="")
espacio2.grid(row=6,column=0)
espacio3=CTkLabel(frame2,text="")
espacio3.grid(row=9,column=0)
espacio4=CTkLabel(frame2,text="")
espacio4.grid(row=12,column=0)

i2c = busio.I2C(board.SCL, board.SDA)
ads = ADS.ADS1115(i2c)

titulo2 = CTkLabel(frame2, text="Sensores",font=("Georgia Pro",18))
titulo2.grid(column=0,row=0,padx=10, pady=(20,25))
SaveDataLabel = CTkLabel(frame2, text="¿Guardar datos?")
#SaveDataLabel.grid(column=0,row=1,padx=5)
switch_var = StringVar(value = "off")
#def switch_event():
SaveDataSwitch = CTkSwitch(frame2, text ="¿Guardar datos?", variable=switch_var, onvalue="on", offvalue="off",font=("Georgia Pro",12))
SaveDataSwitch.grid(column=0,row=1,sticky='e',padx=(0,20))
label1 = CTkLabel(frame2, text="pH",font=("Georgia Pro",12))
label1.grid(column=0,row=4,padx=5,pady=5)
#(column=1,row=14,sticky='e',padx=(0,20))

def SpH():

    phsensor = AnalogIn(ads, ADS.P3).value
    phsensor = round((phsensor+11149.625)/3680.709,2)
    pHSV.set(round(phsensor,1))
    window.after(1000,SpH)

pHSV=StringVar()

Entry1 = CTkEntry(frame2,textvariable=pHSV, justify=CENTER)
Entry1.grid(column=0,row=5,padx=5, pady=1)
ToolTip(Entry1,msg="Rango ideal de 5.5 a 6.5",delay=0,padx=10,pady=10,font=("Georgia Pro",12))
SpH()
Entry1.configure(state= DISABLED)

label2 = CTkLabel(frame2,text="Temperatura del agua (°C)",font=("Georgia Pro",12))
label2.grid(column=0,row=7,padx=5, pady=5)

def Stemp():
    temperatura = sensor.get_temperature()
    TempSV.set(round(temperatura,0))
    window.after(1000,Stemp)

TempSV=StringVar()

sensor = W1ThermSensor()

Entry2 = CTkEntry(frame2,textvariable=TempSV, justify=CENTER,font=("Georgia Pro",12))
Entry2.grid(column=0,row=8,padx=5, pady=1)
ToolTip(Entry2,msg="Rango ideal de 18 a 24°C",delay=0,padx=10,pady=10,font=("Georgia Pro",12))
Stemp()
Entry2.configure(state= DISABLED)

port = 1
#address = 0x77 # Adafruit BME280 address. Other BME280s may be different
#bus = smbus2.SMBus(port)
#bme280.load_calibration_params(bus,address)

labeltempamb=CTkLabel(frame2,text="Temperatura ambiental (°C)",font=("Georgia Pro",12))
labeltempamb.grid(column=0,row=10,padx=5, pady=1)
'''
def Stempamb():
    tempambsensor = bme280.sample(bus,address).temperature
    tempambsensor = round(tempambsensor,1)
    tempambSV.set(tempambsensor)
    window.after(1000,Stempamb)

tempambSV=StringVar()

EntryTempamb = CTkEntry(frame2,textvariable=tempambSV, justify=CENTER,font=("Georgia Pro",12))
EntryTempamb.grid(column=0,row=11,padx=5, pady=1)
EntryTempamb.configure(state= DISABLED)
Stempamb()
'''

label8 = CTkLabel(frame2,text="Conductividad (mS/cm)",font=("Georgia Pro",12))
label8.grid(column=0,row=13,padx=5, pady=1)

def Sconductividad():
    conducsensor = AnalogIn(ads, ADS.P2).value
    #conducsensor = round((conducsensor-127.658)/1222.464 + 0.07,2)
    conducsensor = round((conducsensor*0.000764306)-0.0265299,2)# sensor dfrobot
    #conducsensor = round((conducsensor*0.000166867)-0.115138,2)# sensor tds
    conducSV.set(conducsensor)
    window.after(1000,Sconductividad)

conducSV=StringVar()

Entry8 = CTkEntry(frame2,textvariable=conducSV, justify=CENTER,font=("Georgia Pro",12))
Entry8.grid(column=0,row=14, pady=1)

ToolTip(Entry8,msg="Rango ideal de 1.5 a 2.5 mS/cm",delay=0,padx=10,pady=10,font=("Georgia Pro",12))
Sconductividad()
Entry8.configure(state= DISABLED)

Datosgraf = {
    'ph' : [0, [0,14], 'pH' , 'Horas'],
    'temp' : [1, [15,25], 'Temperatura agua (°C)', 'Horas'],
    'tempa': [2, [15, 25], 'Temperatura ambiente (°C)', 'Horas'],
    'cond' : [3, [0,3], 'Conductividad (mS/cm)', 'Horas'],
    }

def graficar(sensor):
    conn = sqlite3.connect('historial_de_sensores4.db')
    c = conn.cursor()
    c.execute("SELECT *, oid FROM addresses")
    valores=[]
    horas=[]
    for row in c.fetchall():
        valores.append(row[Datosgraf[sensor][0]])
        horas.append(row[4])
    plt.plot(horas,valores)
    plt.xlabel(Datosgraf[sensor][3])
    plt.ylabel(Datosgraf[sensor][2])
    plt.grid(True)
    plt.ylim(Datosgraf[sensor][1])
    plt.xticks(rotation=45)
    plt.show()
    conn.commit()
    conn.close()

icono_graph = CTkImage(Image.open(r"/home/damian/tt/icono3.png").resize((15,15),Image.ANTIALIAS))
btnGraphPH=CTkButton(frame2,text="",image=icono_graph,width=30,height=30,fg_color="gray",hover_color="white",command=lambda: graficar('ph'))
btnGraphPH.grid(column=1,row=5,sticky='e',padx=(0,20))
btnGraphTemp=CTkButton(frame2,text="",image=icono_graph,width=30,height=30,fg_color="gray",hover_color="white",command=lambda: graficar('temp'))
btnGraphTemp.grid(column=1,row=8,sticky='e',padx=(0,20))
btnGraphTempA=CTkButton(frame2,text="",image=icono_graph,width=30,height=30,fg_color="gray",hover_color="white",command=lambda: graficar('tempa'))
btnGraphTempA.grid(column=1,row=11,sticky='e',padx=(0,20))
btnGraphCond=CTkButton(frame2,text="",image=icono_graph,width=30,height=30,fg_color="gray",hover_color="white",command=lambda: graficar('cond'))
btnGraphCond.grid(column=1,row=14,sticky='e',padx=(0,20))

lblcolumna3 = CTkLabel(frame3,text="Control de iluminación",font=("Georgia Pro",18)).pack(padx=10,pady=20)

intensidad_var = StringVar()
intensidad_var.set('Intensidad 100%')
labelIntensidad = CTkLabel(frame3, textvariable=intensidad_var,font=("Georgia Pro",12))
labelIntensidad.pack(padx=10)

GPIO.setmode(GPIO.BCM)
GPIO.setup(21,GPIO.OUT)
p=GPIO.PWM(21,100)
p.start(0)
def evento_slider(value):
    intensidad_var.set('Intensidad '+str(int(value))+'%')

sliderIntensidad = CTkSlider(frame3,from_=0,to=100, width=150,command=evento_slider)
sliderIntensidad.pack(padx=10,pady=(0,16))
sliderIntensidad.set(100)

label4 = CTkLabel(frame3,text="Fotoperiodo (h)",font=("Georgia Pro",12)).pack(padx=10,pady=1)
fotoperiodo=StringVar()
fotoperiodo.set(0)
Entry4 = CTkEntry(frame3,textvariable =fotoperiodo,justify=CENTER)
Entry4.pack(padx=10,pady=(5,35))
ToolTip(Entry4,msg="Para albahaca verde se \n recomiendan 16 horas \n       de iluminación".center(25," "),delay=0,padx=10,pady=10)
valorfoto = Entry4.get()

label5 = CTkLabel(frame3,text="Hora de inicio",font=("Georgia Pro",12)).pack(padx=10,pady=1)
horainicio=StringVar()
horainicio.set('00:00')
Entry5 = CTkEntry(frame3,textvariable =horainicio,justify=CENTER)
Entry5.pack(padx=10,pady=5)
Entry5.configure(state= DISABLED)

fecha_inicial=datetime.datetime.now()
hora_inicial=datetime.time(hour=fecha_inicial.hour,minute=fecha_inicial.minute)
lista = [hora_inicial] #.time()
for i in range(1,24):
    if i + fecha_inicial.hour > 23:
        lista.append(datetime.time(hour=i+fecha_inicial.hour-24, minute=fecha_inicial.minute))
    else:
        lista.append(datetime.time(hour=i+fecha_inicial.hour, minute=fecha_inicial.minute))
print(lista)
def suma():
    seleccion = CTkToplevel()
    btnSelHoras.configure(state=DISABLED)
    seleccion.title('Seleccionar hora de inicio')
    seleccion.geometry("350x185")

    if not valorfoto.isdigit():
        destruir_ventana()

    frame_tiempo=CTkFrame(seleccion)
    frame_tiempo.grid(row=0,column=0,padx=10,pady=20)
    hora= CTkLabel(frame_tiempo, text="Hora")
    hora.grid(row=0,column=0,padx=30,pady=20)
    valores_horas=["00"]

    for i in range(1,24):
        if len(str(i)) == 1:
            valores_horas.append(f'0{str(i)}')
        else:
            valores_horas.append(str(i))
        i+=1
    valores_horas=tuple(valores_horas)

    spin_horas=tk.Spinbox(master=frame_tiempo,values=valores_horas,justify=CENTER, width=10, state="readonly",wrap=True)
    spin_horas.grid(row=0,column=1,padx=10,pady=20)
    valores_minutos=["00"]

    for i in range(1,60):
        if len(str(i)) == 1:
            valores_minutos.append(f'0{str(i)}')
        else:
            valores_minutos.append(str(i))
        i+=1
    valores_minutos=tuple(valores_minutos)

    spin_minutos=tk.Spinbox(master=frame_tiempo,values=valores_minutos,justify=CENTER, width=10, state="readonly",wrap=True)
    spin_minutos.grid(row=0,column=2,padx=10,pady=20)

#spin.bind("<MouseWheel>",scroll copiar scroll
    def destruir_ventana():
        seleccion.destroy()
        btnSelHoras.configure(state=ACTIVE)

    def aceptar():
        valorfoto = Entry4.get()
        valorinicio=spin_horas.get()
        horainicio.set(str(valorinicio)+':'+str(spin_minutos.get()))
        try:
            suma=int(valorfoto)+int(valorinicio)
        except ValueError:
            tk.messagebox.showerror("Advertencia","Advertencia, el fotoperiodo no es un número entero")

        if int(valorfoto)>24:
            tk.messagebox.showerror("Advertencia","Favor de ingresar un valor de fotoperiodo no mayor a 24 h.")
            suma=0
        if suma>= 24:
            suma-=24
        if len(str(suma)) == 1:
            tiempoFin.set('0'+str(suma)+':'+str(spin_minutos.get()))
        else:
            tiempoFin.set(str(suma)+':'+str(spin_minutos.get()))
        #crear lista de horas en la cual comparar para tomar mediciones
        destruir_ventana()

    seleccion.protocol("WM_DELETE_WINDOW",destruir_ventana)
    btnOK=CTkButton(seleccion, text= "Aceptar",command=aceptar)
    btnOK.grid(row=1,column=0,padx=(20,10),pady=15,sticky='sw')
    btnCancel=CTkButton(seleccion, text= "Cancelar",command=destruir_ventana)
    btnCancel.grid(row=1,column=0,padx=(10,20),pady=15,sticky='se')

btnSelHoras = CTkButton(frame3, text="Seleccionar hora", width=25, command=suma,font=("Georgia Pro",12))

btnSelHoras.pack(pady=(5,35),padx=10)

label6 = CTkLabel(frame3,text="Hora de termino",font=("Georgia Pro",12)).pack(pady=1)
tiempoFin=StringVar()
tiempoFin.set('00:00')
Entry6 = CTkEntry(frame3,textvariable=tiempoFin, justify=CENTER,font=("Georgia Pro",12))
Entry6.pack(padx=10,pady=(5,35))
Entry6.configure(state= DISABLED)

EstadoLED=0
def comparacion_encendido():
    global now
    now = datetime.datetime.now().time()
    horafinal = datetime.datetime.strptime(tiempoFin.get(),"%H:%M").time()
    hora1 = datetime.datetime.strptime(horainicio.get(),"%H:%M").time()
        #if (now < horafinal and now > hora1) or (horafinal < hora1 and now < datetime:
    if (now < horafinal and now > hora1) or (horafinal < hora1 and ((now < datetime.time(23,59,0) and now > hora1) or (now > datetime.time(0,0,0) and now <horafinal))):
            #sliderIntensidad.configure(state= ACTIVE)
        p.ChangeDutyCycle(int(sliderIntensidad.get()))
        EstadoLED=1
    else:
        p.ChangeDutyCycle(0)
        EstadoLED=0
    window.after(100,comparacion_encendido)
    #comparar la hora con las horas en las que se van a tomar las muestras
def comparacion_escritura():
    global lista
    global hora_actual

    hora_actual=datetime.time(hour=now.hour,minute=now.minute)
    if switch_var.get() == "on" and hora_actual in lista:
        escribir_datos()
    window.after(60000,comparacion_escritura)

comparacion_encendido()
comparacion_escritura()
def iniciar():
    global cap
    cap = cv2.VideoCapture(-1)
    btnIniciar.configure(state=DISABLED)#evitamos que el usuario seleccione botones que no debería para evitar errores
    btnVisualizar.configure(state=DISABLED)
    visualizar()

def visualizar():
    global cap
    if cap is not None:
        ret, frame = cap.read()
        if ret == True:
            frame = cv2.flip(frame,0)
            frame = imutils.resize(frame, width=640)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            im = Image.fromarray(frame)
            img = ImageTk.PhotoImage(image=im)

            #img = CTkImage(dark_image=frame,size=(200,200))
            lblVideo.configure(image=img)
            lblVideo.image = img
            lblVideo.after(10, visualizar)
        else:
            lblVideo.image = ""
            cap.release()

def finalizar():
    global cap
    cap.release()#Si hay problemas con esto ponerlo dentro de un try
    btnIniciar.configure(state=ACTIVE)
    btnCapturar.configure(state=ACTIVE)
    btnVisualizar.configure(state=ACTIVE)

def tomarfoto():
    global cap
    ret, frame = cap.read()
    frame = cv2.flip(frame,0)
    i=1
    while (os.path.isfile(segmented_button_var.get()+"-"+str(i)+".jpg")):#checar si ya existe la imagen
        i+=1
    else:
        cv2.imwrite(segmented_button_var.get()+"-"+str(i)+".jpg",frame) #si no existe la imagen, escribir la imagen sin sobreescribir

def visualizarfoto():
    btnIniciar.configure(state=DISABLED)
    btnCapturar.configure(state=DISABLED)
    #btnFinalizar.configure(state=DISABLED)
    i=1
    lista_de_fotos=[]
    while (os.path.isfile(segmented_button_var.get()+"-"+str(i)+".jpg")):#checar si ya existe la imagen
        lista_de_fotos.append(ImageTk.PhotoImage(Image.open(segmented_button_var.get()+"-"+str(i)+".jpg")))
        i+=1
    #Creamos ventana donde vamos a ver las fotos de la planta seleccionada
    ventana_fotos = CTkToplevel()
    ventana_fotos.title(segmented_button_var.get())
    global lblFotos
    lblFotos = tk.Label(ventana_fotos,image=lista_de_fotos[0])
    lblFotos.grid(row=0,column=0,columnspan=3)
    def sig(image_number):
        global lblFotos
        global btnSig
        global btnAnt
        lblFotos.grid_forget()
        lblFotos = tk.Label(ventana_fotos,image=lista_de_fotos[image_number-1])
        btnSig = CTkButton(ventana_fotos, text=">>", command=lambda: sig(image_number+1),font=("Georgia Pro",12))
        btnAnt = CTkButton(ventana_fotos, text="<<", command=lambda: ant(image_number-1),font=("Georgia Pro",12))

        if image_number == len(lista_de_fotos):
            btnSig=CTkButton(ventana_fotos, text=">>",state=DISABLED,font=("Georgia Pro",12))
        lblFotos.grid(row=0,column=0,columnspan=3)
        btnAnt.grid(row=1,column=0)
        btnSig.grid(row=1,column=2)
    def ant(image_number):
        global lblFotos
        global btnSig
        global btnAnt
        lblFotos.grid_forget()
        lblFotos = tk.Label(ventana_fotos,image=lista_de_fotos[image_number-1])
        btnSig = CTkButton(ventana_fotos, text=">>", command=lambda: sig(image_number+1),font=("Georgia Pro",12))
        btnAnt = CTkButton(ventana_fotos, text="<<", command=lambda: ant(image_number-1),font=("Georgia Pro",12))

        if image_number == 1:
            btnAnt=CTkButton(ventana_fotos, text="<<",state=DISABLED,font=("Georgia Pro",12))

        lblFotos.grid(row=0,column=0,columnspan=3)
        btnAnt.grid(row=1,column=0)
        btnSig.grid(row=1,column=2)

    btnSig = CTkButton(ventana_fotos, text=">>", command=lambda: sig(2),font=("Georgia Pro",12))
    btnAnt = CTkButton(ventana_fotos, text="<<", command=ant,state=DISABLED,font=("Georgia Pro",12))
    btnAnt.grid(row=1,column=0)
    btnSig.grid(row=1,column=2)

def graficar_crecimiento():
    btnIniciar.configure(state=DISABLED)
    btnCapturar.configure(state=DISABLED)
cap = None

btnIniciar = CTkButton(framebotones, text="Iniciar", width=45, command=iniciar,font=("Georgia Pro",12))
btnIniciar.pack(side=LEFT,padx=20)
btnCapturar = CTkButton(framebotones, text="Tomar Foto", width=45, command= tomarfoto,font=("Georgia Pro",12))
btnCapturar.pack(side=LEFT,pady=10,padx=20)
btnVisualizar = CTkButton(framebotones, text="Visualizar", width=45, command= visualizarfoto,font=("Georgia Pro",12))
btnVisualizar.pack(side=LEFT,pady=10,padx=20)
btnVisualizar = CTkButton(framebotones, text="Graficar", width=45, command= graficar_crecimiento,font=("Georgia Pro",12))
btnVisualizar.pack(side=LEFT,pady=10,padx=20)
btnFinalizar = CTkButton(framebotones, text="Finalizar", width=45, command=finalizar,font=("Georgia Pro",12))
btnFinalizar.pack(side=LEFT,padx=20,pady=10)

selPlantaLabel = CTkLabel(frameSelPlanta, text="Selección de planta",font=("Georgia Pro",12)).pack()
segmented_button_var = StringVar(value="A1")
btnsSelPlanta =CTkSegmentedButton(frameSelPlanta,values=["A1","A2","B1","B2"], variable=segmented_button_var,font=("Georgia Pro",12))
btnsSelPlanta.pack()
lblVideo = tk.Label(frame1,text="") #Se usa label de la lib original de tkinter ya que la de CTk ya no es compatible en la version 5
lblVideo.pack()
lblVideo.configure(bg="#212121")#se buscó un color que fuera similar al del fondo para tapar el color blanco que es el default
# Base de datos

# Creación de base de datos o conexión a una base de datos
conn = sqlite3.connect('historial_de_sensores4.db')

# Creación de cursor
c = conn.cursor()

# Creación de tabla
'''
c.execute("""CREATE TABLE addresses(
        pH real,
        temp real,
        temp_ambiental real,
        conductividad real,
        hora text,
        luz text
        )""")
'''

def escribir_datos():
    conn = sqlite3.connect('historial_de_sensores4.db')
    c = conn.cursor()
    c.execute("INSERT INTO addresses VALUES(:ph,:temp, :temp_ambiental, :conductividad, :hora, :luz)",
              {
                  'ph': Entry1.get(),
                  'temp': Entry2.get(),
                  'temp_ambiental': EntryTempamb.get(),
                  'conductividad': Entry8.get(),
                  'hora': str(hora_actual),
                  'luz': EstadoLED
              }
              )
    conn.commit()
    conn.close()

def leer_datos():
    conn = sqlite3.connect('historial_de_sensores4.db')
    c = conn.cursor()
    c.execute("SELECT *, oid FROM addresses")
    records = c.fetchall()
    print(records)
    conn.commit()
    conn.close()

conn.commit()
conn.close()

window.protocol("WM_DELETE_WINDOW",cerrar_interfaz)

window.mainloop()


import sqlite3
conn=sqlite3.connect('newdb')
c=conn.cursor()
#c.execute("CREATE TABLE if not exists record(name text,age int)")

for i in range(1):
    a=input('enter_name:')
    b=input('enter_age:')
c.execute("insert into record values(?,?)",(a,b))
conn.commit()
conn.close()


# Standard library imports
import csv
import json
from pathlib import Path
from typing import Dict, List, Tuple
# Third party imports
import pandas as pd
import pymongo
from pymongo import MongoClient
from sklearn.model_selection import train_test_split, KFold
from sklearn.tree import DecisionTreeClassifier
# Local imports
from utils import diff_cols, grid_search, prepare_dataset, pymongo_loads

# File definitions
csv_src = Path("/import/SpeedDatingData.csv")

# Preprocessing inputs
query_cols = {
    "iid": 1,
    "pid": 1,
    "age": 1,
    "age_o": 1,
    "dec": 1,
    "dec_o": 1,
    "attr": 1,
    "attr_o": 1,
    "sinc": 1,
    "sinc_o": 1,
    "intel": 1,
    "intel_o": 1,
    "fun": 1,
    "fun_o": 1,
    "amb": 1,
    "amb_o": 1,
    "shar": 1,
    "shar_o": 1,
    "match": 1,
}
diff_cols_ = ["age", "dec", "attr", "sinc", "intel", "fun", "amb", "shar"]
y_cols = ["match"]

# model inputs
test_size = 0.2
random_state = 777
max_depths = [1, 2, 3, 4, 5, 6, 7, 8]

if __name__ == "__main__":
    print(80 * "~")
    print(f"Load data from {csv_src}")
    collection = pymongo_loads(csv_src, "speeddating", "events")

    print(80 * "~")
    print(f"Prepare data for classifier")
    df = prepare_dataset(collection, query_cols, diff_cols_)

    print(80 * "~")
    print("Split into test and train sets")
    X_tr, X_te, y_tr, y_te = train_test_split(df.drop(axis=1, labels=y_cols), df[y_cols[0]], random_state=random_state, test_size=test_size)

    print(80 * "~")
    print("Conduct cross validation")
    validation_scores, validation_summary = grid_search(X_tr, y_tr, max_depths)
    print("Cross validation results")
    print(validation_summary)

    print(80 * "~")
    print("Select best max depth")
    best_max_depth = validation_summary.index.values[validation_summary["acc_val"].argmax()]
    print(f"Best max depth={best_max_depth}")

    print(80 * "~")
    print("Train model on training data using best max depth")
    best_model = DecisionTreeClassifier(max_depth=best_max_depth)
    best_model.fit(X_tr, y_tr)

    print(80 * "~")
    print("Score model on test data")
    acc_te = round(best_model.score(X_te, y_te), 2)
    print(f"Test accuracy={acc_te}")

    print(80 * "~")
    feature_importances = pd.Series(best_model.feature_importances_, index=list(X_tr)).sort_values(ascending=False)
    print(f"Feature importances:\n{feature_importances}")


import sys
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary
from Param import *
from Param_GMAN import *

class conv2d_(nn.Module):
    def __init__(self, input_dims, output_dims, kernel_size, stride=(1, 1),
                 padding='SAME', use_bias=True, activation=F.relu,
                 bn_decay=None):
        super(conv2d_, self).__init__()
        self.activation = activation
        if padding == 'SAME':
            self.padding_size = math.ceil(kernel_size)
        else:
            self.padding_size = [0, 0]
        self.conv = nn.Conv2d(input_dims, output_dims, kernel_size, stride=stride,
                              padding=0, bias=use_bias)
        self.batch_norm = nn.BatchNorm2d(output_dims, momentum=bn_decay)
        torch.nn.init.xavier_uniform_(self.conv.weight)

        if use_bias:
            torch.nn.init.zeros_(self.conv.bias)

    def forward(self, x):
        x = x.permute(0, 3, 2, 1)
        x = F.pad(x, ([self.padding_size[1], self.padding_size[1], self.padding_size[0], self.padding_size[0]]))
        x = self.conv(x)
        x = self.batch_norm(x)
        if self.activation is not None:
            x = F.relu_(x)
        return x.permute(0, 3, 2, 1)

class FC(nn.Module):
    def __init__(self, input_dims, units, activations, bn_decay, use_bias=True):
        super(FC, self).__init__()
        if isinstance(units, int):
            units = [units]
            input_dims = [input_dims]
            activations = [activations]
        elif isinstance(units, tuple):
            units = list(units)
            input_dims = list(input_dims)
            activations = list(activations)
        assert type(units) == list
        self.convs = nn.ModuleList([conv2d_(
            input_dims=input_dim, output_dims=num_unit, kernel_size=[1, 1], stride=[1, 1],
            padding='VALID', use_bias=use_bias, activation=activation,
            bn_decay=bn_decay) for input_dim, num_unit, activation in
            zip(input_dims, units, activations)])

    def forward(self, x):
        for conv in self.convs:
            x = conv(x)
        return x

class STEmbedding(nn.Module):
    """
    spatio-temporal embedding
    SE:     [num_vertex, D]
    TE:     [batch_size, num_hist + num_pred, 2] (dayofweek, timeofday)
    T:      num of time steps in one day
    D:      output dims
    return: [batch_size, num_his + num_pred, num_vertex, D]
    """

    def __init__(self, D, bn_decay, device):
        super(STEmbedding, self).__init__()
        self.FC_se = FC(
            input_dims=[D, D], units=[D, D], activations=[F.relu, None],
            bn_decay=bn_decay
        )

        self.FC_te = FC(
            input_dims=[295, D], units=[D, D], activations=[F.relu, None],
            bn_decay=bn_decay
        )
        self.device = device

    def forward(self, SE, TE, T=288):
        # spatial embedding
        SE = SE.unsqueeze(0).unsqueeze(0)
        SE = self.FC_se(SE)
        # temporal embedding
        dayofweek = torch.empty(TE.shape[0], TE.shape[1], 7)
        timeofday = torch.empty(TE.shape[0], TE.shape[1], T)
        for i in range(TE.shape[0]):
            dayofweek[i] = F.one_hot(TE[..., 0][i].to(torch.int64) % 7, 7)
        for j in range(TE.shape[0]):
            timeofday[j] = F.one_hot(TE[..., 1][j].to(torch.int64) % 288, T)
        TE = torch.cat((dayofweek, timeofday), dim=-1)
        TE = TE.unsqueeze(dim=2).to(device=self.device)
        TE = self.FC_te(TE)
        del dayofweek, timeofday
        return SE + TE

class spatialAttention(nn.Module):
    """
    spatial attention mechanism
    X:      [batch_size, num_step, num_vertex, D]
    STE:    [batch_size, num_step, num_vertex, D]
    K:      number of attention heads
    d:      dimension of each attention outputs
    return: [batch_size, num_step, num_vertex, D]
    """

    def __init__(self, K, d, bn_decay):
        super(spatialAttention, self).__init__()
        D = K * d
        self.d = d
        self.K = K
        self.FC_q = FC(input_dims=2 * D, units=D, activations=F.relu,
                       bn_decay=bn_decay)
        self.FC_k = FC(input_dims=2 * D, units=D, activations=F.relu,
                       bn_decay=bn_decay)
        self.FC_v = FC(input_dims=2 * D, units=D, activations=F.relu,
                       bn_decay=bn_decay)
        self.FC = FC(input_dims=D, units=D, activations=F.relu,
                     bn_decay=bn_decay)

    def forward(self, X, STE):
        batch_size = X.shape[0]
        X = torch.cat((X, STE), dim=-1)
        # [batch_size, num_step, num_vertex, K * d]
        query = self.FC_q(X)
        key = self.FC_k(X)
        value = self.FC_v(X)
        # [K * batch_size, num_step, num_vertex, d]
        query = torch.cat(torch.split(query, self.d, dim=-1), dim=0)
        key = torch.cat(torch.split(key, self.d, dim=-1), dim=0)
        value = torch.cat(torch.split(value, self.d, dim=-1), dim=0)
        # [K * batch_size, num_step, num_vertex, num_vertex]
        attention = torch.matmul(query, key.transpose(2, 3))
        attention /= (self.d ** 0.5)
        attention = F.softmax(attention, dim=-1)
        # [batch_size, num_step, num_vertex, D]
        X = torch.matmul(attention, value)
        X = torch.cat(torch.split(X, batch_size, dim=0), dim=-1)
        X = self.FC(X)
        del query, key, value, attention
        return X

class temporalAttention(nn.Module):
    """
    temporal attention mechanism
    X:      [batch_size, num_step, num_vertex, D]
    STE:    [batch_size, num_step, num_vertex, D]
    K:      number of attention heads
    d:      dimension of each attention outputs
    return: [batch_size, num_step, num_vertex, D]
    """
    def __init__(self, K, d, bn_decay, mask=True):
        super(temporalAttention, self).__init__()
        D = K * d
        self.d = d
        self.K = K
        self.mask = mask
        self.FC_q = FC(input_dims=2 * D, units=D, activations=F.relu,
                       bn_decay=bn_decay)
        self.FC_k = FC(input_dims=2 * D, units=D, activations=F.relu,
                       bn_decay=bn_decay)
        self.FC_v = FC(input_dims=2 * D, units=D, activations=F.relu,
                       bn_decay=bn_decay)
        self.FC = FC(input_dims=D, units=D, activations=F.relu,
                       bn_decay=bn_decay)

    def forward(self, X, STE):
        batch_size_ = X.shape[0]
        X = torch.cat((X, STE), dim=-1)
        # [batch_size, num_step, num_vertex, K * d]
        query = self.FC_q(X)
        key = self.FC_k(X)
        value = self.FC_v(X)
        # [K * batch_size, num_step, num_vertex, d]
        query = torch.cat(torch.split(query, self.d, dim=-1), dim=0)
        key = torch.cat(torch.split(key, self.d, dim=-1), dim=0)
        value = torch.cat(torch.split(value, self.d, dim=-1), dim=0)
        # query: [K * batch_size, num_vertex, num_step, d]
        # key:   [K * batch_size, num_vertex, d, num_step]
        # value: [K * batch_size, num_vertex, num_step, d]
        query = query.permute(0, 2, 1, 3)
        key = key.permute(0, 2, 3, 1)
        value = value.permute(0, 2, 1, 3)
        # [K * batch_size, num_vertex, num_step, num_step]
        attention = torch.matmul(query, key)
        attention /= (self.d ** 0.5)
        # mask attention score
        if self.mask:
            batch_size = X.shape[0]
            num_step = X.shape[1]
            num_vertex = X.shape[2]
            mask = torch.ones(num_step, num_step)
            mask = torch.tril(mask)
            mask = torch.unsqueeze(torch.unsqueeze(mask, dim=0), dim=0)
            mask = mask.repeat(self.K * batch_size, num_vertex, 1, 1)
            mask = mask.to(torch.bool)
            attention = torch.where(mask, attention, -1 ** 15 + 1)
        # softmax
        attention = F.softmax(attention, dim=-1)
        # [batch_size, num_step, num_vertex, D]
        X = torch.matmul(attention, value)
        X = X.permute(0, 2, 1, 3)
        X = torch.cat(torch.split(X, batch_size_, dim=0), dim=-1)
        X = self.FC(X)
        del query, key, value, attention
        return X

class gatedFusion(nn.Module):
    """
    gated fusion
    HS:     [batch_size, num_step, num_vertex, D]
    HT:     [batch_size, num_step, num_vertex, D]
    D:      output dims
    return: [batch_size, num_step, num_vertex, D]
    """

    def __init__(self, D, bn_decay):
        super(gatedFusion, self).__init__()
        self.FC_xs = FC(input_dims=D, units=D, activations=None,
                        bn_decay=bn_decay, use_bias=False)
        self.FC_xt = FC(input_dims=D, units=D, activations=None,
                        bn_decay=bn_decay, use_bias=True)
        self.FC_h = FC(input_dims=[D, D], units=[D, D], activations=[F.relu, None],
                        bn_decay=bn_decay)

    def forward(self, HS, HT):
        XS = self.FC_xs(HS)
        XT = self.FC_xt(HT)
        z = torch.sigmoid(torch.add(XS, XT))
        H = torch.add(torch.mul(z, HS), torch.mul(1 - z, HT))
        H = self.FC_h(H)
        del XS, XT, z
        return H

class STAttBlock(nn.Module):
    def __init__(self, K, d, bn_decay, mask=False):
        super(STAttBlock, self).__init__()
        self.spatialAttention = spatialAttention(K, d, bn_decay)
        self.temporalAttention = temporalAttention(K, d, bn_decay, mask=mask)
        self.gatedFusion = gatedFusion(K * d, bn_decay)

    def forward(self, X, STE):
        HS = self.spatialAttention(X, STE)
        HT = self.temporalAttention(X, STE)
        H = self.gatedFusion(HS, HT)
        del HS, HT
        return torch.add(X, H)

class transformAttention(nn.Module):
    """
    transform attention mechanism
    X:          [batch_size, num_his, num_vertex, D]
    STE_his:    [batch_size, num_his, num_vertex, D]
    STE_pred:   [batch_size, num_pred, num_vertex, D]
    K:          number of attention heads
    d:          dimension of each attention outputs
    return:     [batch_size, num_pred, num_vertex, D]
    """

    def __init__(self, K, d, bn_decay):
        super(transformAttention, self).__init__()
        D = K * d
        self.K = K
        self.d = d
        self.FC_q = FC(input_dims=D, units=D, activations=F.relu,
                       bn_decay=bn_decay)
        self.FC_k = FC(input_dims=D, units=D, activations=F.relu,
                       bn_decay=bn_decay)
        self.FC_v = FC(input_dims=D, units=D, activations=F.relu,
                       bn_decay=bn_decay)
        self.FC = FC(input_dims=D, units=D, activations=F.relu,
                     bn_decay=bn_decay)

    def forward(self, X, STE_his, STE_pred):
        batch_size = X.shape[0]
        # [batch_size, num_step, num_vertex, K * d]
        query = self.FC_q(STE_pred)
        key = self.FC_k(STE_his)
        value = self.FC_v(X)
        # [K * batch_size, num_step, num_vertex, d]
        query = torch.cat(torch.split(query, self.d, dim=-1), dim=0)
        key = torch.cat(torch.split(key, self.d, dim=-1), dim=0)
        value = torch.cat(torch.split(value, self.d, dim=-1), dim=0)
        # query: [K * batch_size, num_vertex, num_pred, d]
        # key:   [K * batch_size, num_vertex, d, num_his]
        # value: [K * batch_size, num_vertex, num_his, d]
        query = query.permute(0, 2, 1, 3)
        key = key.permute(0, 2, 3, 1)
        value = value.permute(0, 2, 1, 3)
        # [K * batch_size, num_vertex, num_pred, num_his]
        attention = torch.matmul(query, key)
        attention /= (self.d ** 0.5)
        attention = F.softmax(attention, dim=-1)
        # [batch_size, num_pred, num_vertex, D]
        X = torch.matmul(attention, value)
        X = X.permute(0, 2, 1, 3)
        X = torch.cat(torch.split(X, batch_size, dim=0), dim=-1)
        X = self.FC(X)
        del query, key, value, attention
        return X

class GMAN(nn.Module):
    """
    GMAN
        X:          [batch_size, num_his, num_vertex]
        TE:         [batch_size, num_his + num_pred, 2] (time-of-day, day-of-week)
        SE:         [num_vertex, K * d]
        num_his:    number of history steps
        num_pred:   number of prediction steps
        T:          one day is divided into T steps
        L:          number of STAtt blocks in the encoder/decoder
        K:          number of attention heads
        d:          dimension of each attention head outputs
        return:     [batch_size, num_pred, num_vertex]
    """

    def __init__(self, SE, timestep_in, device, statt_layers=1, att_heads=8, att_dims=8, bn_decay=0.1):
        super(GMAN, self).__init__()
        L = statt_layers
        K = att_heads
        d = att_dims
        D = K * d
        self.num_his = timestep_in
        self.SE = SE
        self.STEmbedding = STEmbedding(D, bn_decay, device)
        self.STAttBlock_1 = nn.ModuleList([STAttBlock(K, d, bn_decay) for _ in range(L)])
        self.STAttBlock_2 = nn.ModuleList([STAttBlock(K, d, bn_decay) for _ in range(L)])
        self.transformAttention = transformAttention(K, d, bn_decay)
        self.FC_1 = FC(input_dims=[1, D], units=[D, D], activations=[F.relu, None],
                       bn_decay=bn_decay)
        self.FC_2 = FC(input_dims=[D, D], units=[D, 1], activations=[F.relu, None],
                       bn_decay=bn_decay)

    def forward(self, X, TE):

        # input
        X = torch.unsqueeze(X, -1)
        X = self.FC_1(X)
        # STE
        STE = self.STEmbedding(self.SE, TE)
        STE_his = STE[:, :self.num_his]
        STE_pred = STE[:, self.num_his:]
        # encoder
        for net in self.STAttBlock_1:
            X = net(X, STE_his)
        # transAtt
        X = self.transformAttention(X, STE_his, STE_pred)
        # decoder
        for net in self.STAttBlock_2:
            X = net(X, STE_pred)
        # output
        X = self.FC_2(X)
        del STE, STE_his, STE_pred
        return torch.squeeze(X, 3)

def main():
    GPU = sys.argv[-1] if len(sys.argv) == 2 else '3'
    device = torch.device("cuda:{}".format(GPU)) if torch.cuda.is_available() else torch.device("cpu")
    SE = torch.zeros((N_NODE, SE_DIM), dtype=torch.float32).to(device=device)
    model = GMAN(SE, TIMESTEP_IN, device).to(device)
    summary(model, [(TIMESTEP_IN, N_NODE),(TIMESTEP_IN+TIMESTEP_OUT, TE_DIM)], device=device)

if __name__ == '__main__':
    main()


import random
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Dataset

class Sharding(Dataset):
    """A sharding of a dataset"""

    def __init__(self, data, index):
        self.data = data
        self.index = index

    def __len__(self):
        return len(self.index)

    def __getitem__(self, index):
        data_idx = self.index[index]
        return self.data[data_idx]

class DatasetCreator(object):
    """Dataset creator loads and divides dataset"""

    def __init__(self, num_clients, dataset_name, partition_method, seed=5):
        random.seed(seed)  # fix seed for reproduction
        self.n_clients = num_clients
        self.dt_name = dataset_name
        self.pt_method = partition_method
        self.shardings = []
        self.training_data, self.test_data = self.__get_dataset()
        self.__divide_dataset()

    def get_loader(self, client_id=0, batch_sz=32, is_test=False, pin_memory=True, time_out=0, num_workers=0):
        """Get DataLoader object."""
        if not is_test:
            drop_last = shuffle = True
            sharding = Sharding(self.training_data, self.shardings[client_id])
            loader = DataLoader(sharding, batch_sz, shuffle=shuffle, pin_memory=pin_memory, timeout=time_out,
                                num_workers=num_workers, drop_last=drop_last)
        else:
            index = [_ for _ in range(len(self.test_data))]
            non_sharding = Sharding(self.test_data, index)
            loader = DataLoader(non_sharding, batch_sz, shuffle=False, pin_memory=pin_memory, timeout=time_out,
                                num_workers=num_workers, drop_last=False)
        return loader

    def get_training_data_len(self, client_id=0):
        return len(self.shardings[client_id])

    def get_test_data_len(self):
        return len(self.test_data)

    def __divide_dataset(self):
        if self.pt_method == 'uniform':
            data_len = len(self.training_data)
            train_indexes = list(range(data_len))

            random.shuffle(train_indexes)

            part_len = int(1. / self.n_clients * len(train_indexes))

            for _ in range(self.n_clients):
                self.shardings.append(train_indexes[0:part_len])
                train_indexes = train_indexes[part_len:]

    def __get_dataset(self):
        training_dataset = test_dataset = None
        training_transform, test_transform = self.__get_transform()

        if self.dt_name == 'CIFAR10':
            training_dataset = datasets.CIFAR10(
                root='data/CIFAR10',
                train=True,
                download=True,
                transform=training_transform
            )
            test_dataset = datasets.CIFAR10(
                root='data/CIFAR10',
                train=False,
                download=True,
                transform=test_transform
            )

        return training_dataset, test_dataset

    def __get_transform(self):
        if self.dt_name == 'CIFAR10':
            transform_train = transforms.Compose([
                transforms.RandomCrop(32, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
            ])

            transform_test = transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
            ])

        return transform_train, transform_test


import matplotlib.pyplot as plt
import numpy as np

dataArray = np.genfromtxt('csv_file.csv', delimiter=',', names=True)

plt.figure()
for col_name in dataArray.dtype.names:
    if(col_name == dataArray.dtype.names[1]):
        continue
    plt.plot(dataArray[col_name], label=col_name)
plt.legend()
plt.show()

ECE 484 MPs/Final Project/Dec 6/gen_graph.py


import jwt
import json
import requests
from time import time
from datetime import timedelta
from hashlib import sha256
from api import cfg
from flask import Flask, request, redirect, session, url_for
from app import app

app.secret_key = cfg['secret_key']

clannel_id = cfg['line_login']['clannel_id']
clannel_secret = cfg['line_login']['clannel_secret']
login_url = 'http://%s/login' % (cfg['web_url'])

line_login_url = 'https://access.line.me/oauth2/v2.1/authorize?' + '&'.join([
    'redirect_uri=%s' % login_url,
    'client_id=%s' % clannel_id,
    'response_type=code',
    'scope=openid+profile',
    'state=%s' % cfg['line_login']['state'],
    'nonce=%s' % cfg['line_login']['nonce'],
    'bot_prompt=normal',
])

@app.route('/login')
def login():
    code = request.args.get('code')
    state = request.args.get('state')
    error = request.args.get('error', None)

    if error is not None:
        return '登入失敗=%s' % error

    token_time = time()
    headers = {'Content-Type': 'application/x-www-form-urlencoded'}
    token = requests.post('https://api.line.me/oauth2/v2.1/token', headers=headers, data={
        'redirect_uri':login_url,
        'client_id':clannel_id,
        'client_secret':clannel_secret,
        'grant_type':'authorization_code',
        'code':code,
    })

    token = json.loads(token.text)
    if 'error' in token:
        return '登入失敗=%s' % token.get('error_description')

    token_id = jwt.decode(token['id_token'], clannel_secret, audience=clannel_id, issuer='https://access.line.me', algorithms=['HS256'])
    if token_id['iss'] != 'https://access.line.me': return '驗證失敗1'
    if token_id['aud'] != clannel_id: return '驗證失敗2'
    if token_id['nonce'] != cfg['line_login']['nonce']: return '驗證失敗3'
    if token_id['exp'] < token_time: return '驗證失敗4'

    token = sha256((cfg['line_login']['state'] + token_id['sub']).encode()).hexdigest()
    session['token'] = token

    session['uid'] = token_id['sub']
    #session['user_name'] = token_id['name']
    #session['user_picture'] = token_id['picture']

    #session['group_id'] = None if state == str(None) else state

    return redirect('/') #redirect('/g/%s' % session['group_id']) if 'group_id' in session else redirect('/')

'''
@app.before_request
def make_session_permanent(): #session 有效時間
    session.permanent = True
    app.permanent_session_lifetime = timedelta(days=7)

'''


import cv2
import matplotlib.pyplot as plt

img = cv2.imread("images/child.jpg")

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) #Convert BGR to RGB
ax = plt.subplot(221)
plt.imshow(img)

img2 = cv2.flip(img, 2)
ax = plt.subplot(222)
plt.imshow(img2)

img3 = cv2.flip(img, 0)
ax = plt.subplot(223)
plt.imshow(img3)

img4 = cv2.flip(img, -1)
ax = plt.subplot(224)
plt.imshow(img4)

plt.show()


#!/usr/bin/env python
import numpy as np
import random
import itertools as it
import os
import copy

import gym
from gym import spaces

class ConnectNEnv(gym.Env):
    """
    Methods this class has to contain:

    def reset() --> state
    def step(action) --> state, reward, done, info
    def render() --> # Render board

    """
    metadata = {'render.modes': ['human']}

    VISUALISATION = {
        1.: '\x1b[0;31;40m' + 'X' + '\x1b[0m',  # Color red
        -1.: '\x1b[0;32;40m' + 'O' + '\x1b[0m',  # Color green
        0.: ' '
    }

    def __init__(self, game_mode=['learner', 'random'], agent=None):
        # Game mode must be lenght of tow
        assert len(game_mode) == 2
        # Only specific game modes are allowed
        for gm in game_mode:
            assert gm in ['learner', 'random', 'human', 'agent']

        self.game_mode = game_mode

        self.__version__ = "0.1.0"

        self.n = 3  # Required disks to win
        self.m = 4  # Size of field

        self.curr_step = -1

        self.agent = agent

        self.name = f'Connect{self.n}'

        self.action_space = spaces.Discrete(self.m)
        self.observation_space = spaces.Box(low=-1, high=1, dtype=np.int, shape=(self.m, self.m))
        #self.observation_space = spaces.Discrete(self.m*self.m)

        self.curr_episode = -1
        self.action_episode_memory = []

        # Manipulate rewards
        self.reward_win = 1
        self.reward_lost = -10
        self.reward_set_on_full_column = -0.5
        self.reward_draw = 0

        self.create_winning_templates()
        self.reset()

    def _step_learner(self, action):
        # Check for valid turn. If not, add a small negative reward
        if action not in self.possible_turns:
            return self.board, self.reward_set_on_full_column, True, {}

        # Place disk
        self._place_disk(action)

        # Check for winning
        if self.check_for_winning():
            return self.board, self.reward_win, True, {}

        # Board is full --> no winner
        elif self.disks_set >= self.m * self.m:
            return self.board, self.reward_draw, True, {}

        # Game goes on, toggle the disk
        self.toggle_disk()

        # Default return
        return self.board, 0, False, {}

    def _step_agent(self):
        # Evaluate action from agent
        self.agent.training = False
        action = self.agent.forward(self.board)
        self.agent.training = True

        # If turn is invalid, choose random turn
        if action not in self.possible_turns:
            action = self._get_random_action()

        # Place disk and evaluate
        self._evaluate_opponent_turn(action)

    def _step_random(self):
        # Choose random (possible) action
        action = self._get_random_action()

        # Place disk and evaluate
        self._evaluate_opponent_turn(action)

    def _step_human(self, action):
        # TODO: Write human interface here
        return

    def _evaluate_opponent_turn(self, action):
        # Place disk
        self._place_disk(action)

        # Check for winning
        if self.check_for_winning():
            return self.board, self.reward_lost, True, {}

        # Board is full --> no winner
        elif self.disks_set >= self.m * self.m:
            return self.board, self.reward_draw, True, {}

        # Game goes on, toggle the disk
        self.toggle_disk()

        # Default return
        return self.board, 0, False, {}

    def _place_disk(self, action):
        # Place disk
        disks_in_column = np.sum(np.abs(self.board[:, action]))
        self.board[int(self.m - disks_in_column - 1), action] = \
            self.current_disk

        # Adjust turn options, of column is full
        if disks_in_column + 1 == self.m:
            self.possible_turns.remove(action)

        # Count up
        self.disks_set += 1

    def step(self, action):
        result = self._step_learner(action)
        if result is not None:
            return result

        if 'random' in game_mode:
            result = self._step_random()
            if result is not None:
                return result

    def single_step(self, action):
        # Check for valid turn
        if action not in self.possible_turns:
            return self.board, -0.5, True, {}

        # Turn is valid
        self.disks_set += 1

        # Place disk
        disks_in_column = np.sum(np.abs(self.board[:, action]))
        self.board[int(self.m - disks_in_column - 1), action] = \
            self.current_disk

        # Adjust turn options, of column is full
        if disks_in_column + 1 == self.m:
            self.possible_turns.remove(action)

        # Check if learner won the game
        if self.check_for_winning():
            return self.board, 1, True, {}

        # Board is full --> no winner
        elif self.disks_set >= self.m * self.m:
            return self.board, 0, True, {}

        self.toggle_disk()

    def reset(self):
        """
        Reset the state of the environment and returns an initial observation.
        Returns
        -------
        observation (object): the initial observation of the space.
        """
        # Update statistics
        self.curr_episode += 1
        self.action_episode_memory.append([])

        self.possible_turns = set(range(self.m))
        self.board = self.create_empty_board()

        self.disks = it.cycle([1, -1])
        self.toggle_disk()

        self.disks_set = 0
        self.game_over = False

        self.winner = None

        self.reward = 0
        self.terminal = False

        return self.board

    def seed(self, seed):
        random.seed(seed)
        np.random.seed

    def observe(self):
        return self.board, self.reward, self.terminal

    def execute_action(self, action):
        self.update(action)

    def toggle_disk(self):
        self.current_disk = next(self.disks)

    def create_empty_board(self):
        return np.zeros((self.m, self.m))

    def create_winning_templates(self):
        # Horizontal and vertical
        winning_templates = []
        for row in range(self.m):
            for col in range(self.m):

                # Horizontal lines
                if col + self.n <= self.m:
                    board_hor = self.create_empty_board()
                    board_hor[row, col:int(col+self.n)] = 1
                    winning_templates.append(board_hor)

                # Vertical lines
                if row + self.n <= self.m:
                    board_ver = self.create_empty_board()
                    board_ver[row:int(row+self.n), col] = 1
                    winning_templates.append(board_ver)

                # Diagonal lines (up and down)
                if (col + self.n <= self.m) & (row + self.n <= self.m):
                    board_diag_down = self.create_empty_board()
                    board_diag_down[row:, col:][np.diag_indices(self.n)] = 1
                    board_diag_up = np.fliplr(board_diag_down)
                    winning_templates.append(board_diag_down)
                    winning_templates.append(board_diag_up)

        self.winning_templates = np.array(winning_templates)

    def check_for_winning(self):
        if np.max(
            np.abs(
                np.sum(
                    np.sum(self.board * self.winning_templates, axis=1),
                    axis=1
                )
            )
        ) == self.n:
            return True
        return False

    def set_disk(self, column):
        # Check for valid turn
        if column not in self.possible_turns:
            print(
                f'Not possible, player {self.current_disk}. Please choose one of the following '
                f'columns: {self.possible_turns}. Your choice: {column}'
            )
            return False

        # Place disk
        disks_in_column = np.sum(np.abs(self.board[:, column]))
        self.board[int(self.m - disks_in_column - 1), column] = \
            self.current_disk

        # Adjust turn options, of column is full
        if disks_in_column + 1 == self.m:
            self.possible_turns.remove(column)

        return True

    def display_board(self):
        for row in self.board:
            text = '|' + ''.join([self.VISUALISATION[val] for val in row]) + '|'
            print(text)
        # Make it more human readable
        print(f' {"".join([str(i) for i in range(self.m)])} ')

    def _get_random_action(self):
        self.set_disk(np.random.choice(list(self.possible_turns)))

    def render(self, mode='human', close=False):
        self.display_board()
        return

    def stack_on_top(self):
        # Try to find a good turn
        for row in range(self.m - 1, -1, -1):
            row_ = self.board[row]
            for i, col in enumerate(row_):
                if col != 0 and i in self.possible_turns:
                    self.set_disk(int(i))
                    return

        # Otherwise, apply random turn
        self._get_random_action()


import numpy as np
from torch.utils.data import DataLoader, Dataset, random_split, ConcatDataset
import torch
from typing import Tuple, List

class DIADataSet(Dataset):
    def __init__(
        self,
        libMsMz: np.ndarray,
        libMs: np.ndarray,
        libSpectrumIfShare: np.ndarray,
        libTarget: np.ndarray,
    ) -> None:
        r"""
        Input:
        * libMsMz 肽段匹配到的 s 个最相关的混合图谱
        * libMs 肽段理论二级图谱
        * libSpectrumIfShare 肽段特异峰向量
        * libTarget 标签
        """
        super(DIADataSet, self).__init__()
        self.libSpectrumIfShare = libSpectrumIfShare
        self.libMsMz = libMsMz
        self.libMs = libMs
        self.libTarget = libTarget
        self.len = len(self.libTarget)

    def __len__(self) -> int:
        return self.len

    def __getitem__(
        self, index: int
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        r""" """
        return (
            self.libMsMz[index],
            self.libMs[index],
            self.libSpectrumIfShare[index],
            self.libTarget[index],
        )

def SplitTrainAndValDataSet(
    TrainDataSet: DIADataSet,
    TrainDataDecoySet: DIADataSet,
    lengthValDataSet: int,
    lengthValDataDecoySet: int,
) -> Tuple[DIADataSet, DIADataSet, DIADataSet]:
    r"""
    Input:
    * TrainDataSet: 正样本训练数据
    * TrainDataDecoySet 负样本训练数据
    * lengthValDataSet 正样本验证集大小
    * lengthValDataDecoySet 负样本验证集大小
    1. 根据输入的训练数据, 从正样本和负样本中各抽取数条数据作为验证集，每个 epoch 结束前进行 FDR 控制
       选取大于分数阈值最多的模型
    2. 将剩余训练数据中的正样本和负样本进行合并并且进行随机化选取

    Output:
        返回训练集, 验证集, 诱饵验证集
    """
    trainDataSet, valDataSet = random_split(
        TrainDataSet,
        [len(TrainDataSet) - lengthValDataSet, lengthValDataSet],
        generator=torch.Generator().manual_seed(0),
    )
    trainDataDecoySet, valDataDecoySet = random_split(
        TrainDataDecoySet,
        [len(TrainDataDecoySet) - lengthValDataDecoySet, lengthValDataDecoySet],
        generator=torch.Generator().manual_seed(0),
    )
    trainDataSet = ConcatDataset([trainDataSet, trainDataDecoySet])
    return trainDataSet, valDataSet, valDataDecoySet

def GetAllDataLoader(
    batch_size: int,
    TrainDataSet: DIADataSet,
    TrainDataSetDecoy: DIADataSet,
    testDataSet: DIADataSet,
    testDataSetDecoy: DIADataSet,
    lengthValDataSet: int,
    lengthValDataSetDecoy: int,
) -> Tuple[DataLoader, DataLoader, DataLoader, DataLoader, DataLoader]:
    r"""
    把 GetTrainDataLoader 和 GetTestDataLoader 两个函数操作合并在一起
    """
    trainDataSet, valDataSet, valDataDecoySet = SplitTrainAndValDataSet(
        TrainDataSet, TrainDataSetDecoy, lengthValDataSet, lengthValDataSetDecoy
    )
    trainData, valData, valDataDecoy = GetTrainDataLoader(
        batch_size, trainDataSet, valDataSet, valDataDecoySet
    )
    testData, testDataDecoy = GetTestDataLoader(
        batch_size, testDataSet, testDataSetDecoy
    )
    return trainData, valData, valDataDecoy, testData, testDataDecoy

def GetTrainDataLoader(
    batch_size: int,
    trainDataSet: DIADataSet,
    valDataSet: List[DIADataSet],
    valDataDecoySet: List[DIADataSet],
) -> Tuple[DataLoader, List[DataLoader], List[DataLoader]]:
    r"""
    Input:
    * batch_size: batch 大小
    * trainDataSet: 训练集
    * valDataSet: 验证集
    * valDataDecoySet: 诱饵验证集

    OutPut: return
    * 训练集
    * 正样本验证集
    * 诱饵验证集的 dataloader
    """
    trainData = DataLoader(
        dataset=trainDataSet, shuffle=True, batch_size=batch_size, num_workers=8
    )
    valData = [
        DataLoader(dataset=dataset, shuffle=False,
                   batch_size=batch_size, num_workers=8)
        for dataset in valDataSet
    ]
    valDataDecoy = [
        DataLoader(dataset=dataset, shuffle=False,
                   batch_size=batch_size, num_workers=8)
        for dataset in valDataDecoySet
    ]
    return trainData, valData, valDataDecoy

def GetTestDataLoader(
    batch_size: int, testDataSet: DIADataSet, testDataDecoySet: DIADataSet
) -> Tuple[DataLoader, DataLoader]:
    r"""
    Input:
    * batch_size: batch 大小
    * testDataSet: 测试集
    * testDataDecoySet: 诱饵测试集

    Output: return
    * 测试集
    * 诱饵测试集的 dataloader
    """
    testData = DataLoader(
        dataset=testDataSet, shuffle=False, batch_size=batch_size, num_workers=8
    )
    testDataDecoy = DataLoader(
        dataset=testDataDecoySet, shuffle=False, batch_size=batch_size, num_workers=8
    )
    return testData, testDataDecoy

def LibSpectrumIfSharePreTreatMent(LibSpectrumIfShare: np.ndarray) -> np.ndarray:
    r"""
    对特异峰向量进行预处理操作
    [0, 0, 0, 0, 1, 1] -> [[0], [0], [0], [0], [1], [1]]
    """
    x = LibSpectrumIfShare
    y = []
    for item in x:
        y.append(item)
    testLibSpectrumIfShare = np.array(y)
    testLibSpectrumIfShare = np.reshape(
        testLibSpectrumIfShare, (np.shape(testLibSpectrumIfShare)[0], 6, 1)
    )
    return testLibSpectrumIfShare

def LibMsMzPreTreatMent(LibMsMz: np.ndarray) -> np.ndarray:
    r"""
    1. 匹配到的谱图进行预处理操作
    2. 归一化等或不进行操作
    """
    Ms = []
    for data in LibMsMz:
        Ms.append(data)
    test = np.array(Ms)
    msIntensity = test[:, :, :, [1]]
    msIntensity = msIntensity
    msMz = test[:, :, :, [0]]
    Ms = np.concatenate((msMz, msIntensity), axis=3)
    LibMsMz = Ms
    return LibMsMz

def Pretreatment(
    data: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    r"""
    对图谱和特异峰向量进行预处理操作
    """
    LibMsMz, LibMz, LibTarget, LibSpectrumIfShare, fileWinLib = (
        data[0],
        data[1],
        data[2],
        data[3],
        data[4],
    )
    LibSpectrumIfShare = LibSpectrumIfSharePreTreatMent(LibSpectrumIfShare)
    LibMsMz = LibMsMzPreTreatMent(LibMsMz)
    return LibMsMz, LibMz, LibTarget, LibSpectrumIfShare, fileWinLib

def GetTrainDataSet(data: np.ndarray) -> DIADataSet:
    r"""
    Input:
    * data 输入的初始训练数据

    事先对肽段匹配上的图谱、肽段参考图谱、肽段特异峰向量进行预处理操作, 详情见 Pretreatment 函数

    Output: return
    * 处理之后的数据集
    """
    LibMsMz, LibMz, LibTarget, LibSpectrumIfShare, _ = Pretreatment(data)
    return DIADataSet(LibMsMz, LibMz, LibSpectrumIfShare, LibTarget)

def GetTestDataSetAndFilewindow(
    data: np.ndarray,
) -> Tuple[DIADataSet, Tuple[str, str, Tuple[str, int]]]:
    r"""
    Input:
    * data 输入的初始测试数据

    事先对肽段匹配上的图谱、肽段参考图谱、肽段特异峰向量进行预处理操作, 详情见 Pretreatment 函数

    Output: return
    * 处理之后的数据集
    * (文件名、窗口信息、肽段名)
    """
    LibMsMz, LibMz, LibTarget, LibSpectrumIfShare, FileWinLib = Pretreatment(
        data)
    return DIADataSet(LibMsMz, LibMz, LibSpectrumIfShare, LibTarget), FileWinLib

def GetTestData(
    batch_size: int, testPath: str, testDecoyPath: str
) -> Tuple[DataLoader, DataLoader, Tuple[str], Tuple[str]]:
    r"""
    Input:
    * batch_size batch 大小
    * testPath 测试数据路径
    * testDecoyPath 测试诱饵数据路径

    Output: return
    * 测试数据 DataLoader
    * 测试诱饵数据 DataLoader
    * 测试数据 (文件, 窗口, (肽段名, 带电荷数))
    * 测试诱饵数据 (文件, 窗口, (肽段名, 带电荷数))
    """
    TestData = np.load(testPath, allow_pickle=True)
    TestDataDecoy = np.load(testDecoyPath, allow_pickle=True)

    testDataSet, fileWinLib = GetTestDataSetAndFilewindow(TestData)
    testDataDecoySet, fileWinLibDecoy = GetTestDataSetAndFilewindow(
        TestDataDecoy)

    testData, testDataDecoy = GetTestDataLoader(
        batch_size, testDataSet, testDataDecoySet
    )

    return testData, testDataDecoy, fileWinLib, fileWinLibDecoy

def GetLFQAndTrypsin_HFXData(
    batch_size: int,
    lfqPath: str,
    lfqDecoyPath: str,
    Trypsin_HFXpath: str,
    Trypsin_HFXdecoyPath: str,
) -> Tuple[DataLoader, List[DataLoader], List[DataLoader]]:
    r"""
    Input:
    * batch_size: batch 大小
    * lfqPath: lfq 正样本训练数据路径  (可能为 lfq32 和 lfq64)
    * lfqDecoyPath: lfq 负样本训练数据路径 (可能为 lfq32 和 lfq64)
    * Trypsin_HFXPath: 血浆正样本训练数据路径
    * Trypsin_HFXDecoyPath: 血浆负样本训练数据路径

    获取 lfq 和 trupsin_HFX 血浆的训练数据

    Output: return
    * 训练集
    * 正样本验证集
    * 诱饵验证集的 dataloader
    """
    # lfq 和血浆数据
    lfqTrainData = np.load(lfqPath, allow_pickle=True)
    lfqDecoyData = np.load(lfqDecoyPath, allow_pickle=True)
    Trypsin_HFXTrainData = np.load(Trypsin_HFXpath, allow_pickle=True)
    Trypsin_HFXDecoyData = np.load(Trypsin_HFXdecoyPath, allow_pickle=True)
    # 获取数据的 DataSet
    lfqTrainDataSet = GetTrainDataSet(lfqTrainData)
    lfqTrainDataDecoySet = GetTrainDataSet(lfqDecoyData)
    Trypsin_HFXTrainDataSet = GetTrainDataSet(Trypsin_HFXTrainData)
    Trypsin_HFXTrainDataDecoySet = GetTrainDataSet(Trypsin_HFXDecoyData)

    """
    由于血浆和 lfq 样本数量不一致, 有训练集和验证集拆分可以采用两种不同的方法
    一、
        实现将两者进行合并处理, 再从正样本和负样本中各抽取 2w 条数据作为验证集和诱饵验证集

    二、
    1. 血浆从正样本和负样本中各抽取 5k 条作为验证集和诱饵验证集, lfq 从正样本和负样本中抽取 1w 条作为验证集和诱饵验证集
    2. 把两者训练集、验证集、诱饵验证集进行合并, 训练集随机化
    """
    # TrainDataSet = ConcatDataset([lfqTrainDataSet, Trypsin_HFXTrainDataSet])
    # TrainDataDecoySet = ConcatDataset(
    #     [lfqTrainDataDecoySet, Trypsin_HFXTrainDataDecoySet]
    # )

    # trainDataSet, valDataSet, valDataDecoySet = SplitTrainAndValDataSet(
    #     TrainDataSet, TrainDataDecoySet, 20000, 20000
    # )

    lfqtrainDataSet, lfqvalDataSet, lfqvalDataDecoySet = SplitTrainAndValDataSet(
        lfqTrainDataSet, lfqTrainDataDecoySet, 10000, 10000
    )

    (
        Trypsin_HFXtrainDataSet,
        Trypsin_HFXvalDataSet,
        Trypsin_HFXvalDataDecoySet,
    ) = SplitTrainAndValDataSet(
        Trypsin_HFXTrainDataSet, Trypsin_HFXTrainDataDecoySet, 1000, 1000
    )

    trainDataSet = ConcatDataset([lfqtrainDataSet, Trypsin_HFXtrainDataSet])
    # ConcatDataset([lfqvalDataSet, Trypsin_HFXvalDataSet]),
    # ConcatDataset([lfqvalDataDecoySet, Trypsin_HFXvalDataDecoySet]),

    valDataSet, valDataDecoySet = [lfqvalDataSet, Trypsin_HFXvalDataSet], [
        lfqvalDataDecoySet,
        Trypsin_HFXvalDataDecoySet,
    ]

    trainData, valData, valDataDecoy = GetTrainDataLoader(
        batch_size, trainDataSet, valDataSet, valDataDecoySet
    )
    return trainData, valData, valDataDecoy

def GetAllTrainData(batch_size: int,
                    lfq32Path: str,
                    lfq32DecoyPath: str,
                    lfq64Path: str,
                    lfq64DecoyPath: str,
                    tims20211002Path: str,
                    tims20211002DecoyPath: str,
                    Trypsin_HFXpath: str,
                    Trypsin_HFXdecoyPath: str) -> Tuple[DataLoader, List[DataLoader], List[DataLoader]]:
    r"""
    Input:
    * batch_size: batch 大小
    * lfq32Path: lfq32 正样本训练数据路径
    * lfq32DecoyPath: lfq32 负样本训练数据路径
    * lfq64Path: lfq64 正样本训练数据路径
    * lfq64DecoyPath: lfq64 负样本训练数据路径
    * tims20211002Path: tims 血浆正样本训练数据路径
    * tims20211002Path: tims 血浆负样本训练数据路径
    * Trypsin_HFXPath: Trypsin_HFX 血浆正样本训练数据路径
    * Trypsin_HFXDecoyPath: Trypsin_HFX 血浆负样本训练数据路径

    获取 lfq 和两类血浆的训练数据

    Output: return
    * 训练集
    * 正样本验证集
    * 诱饵验证集的 dataloader
    """
    # lfq 和血浆数据
    lfq32TrainData = np.load(lfq32Path, allow_pickle=True)
    lfq32DecoyData = np.load(lfq32DecoyPath, allow_pickle=True)
    lfq64TrainData = np.load(lfq64Path, allow_pickle=True)
    lfq64DecoyData = np.load(lfq64DecoyPath, allow_pickle=True)
    tims20211002TrainData = np.load(tims20211002Path, allow_pickle=True)
    tims20211002DecoyData = np.load(tims20211002DecoyPath, allow_pickle=True)
    Trypsin_HFXTrainData = np.load(Trypsin_HFXpath, allow_pickle=True)
    Trypsin_HFXDecoyData = np.load(Trypsin_HFXdecoyPath, allow_pickle=True)
    # 获取数据的 DataSet
    lfq32TrainDataSet = GetTrainDataSet(lfq32TrainData)
    lfq32TrainDataDecoySet = GetTrainDataSet(lfq32DecoyData)
    lfq64TrainDataSet = GetTrainDataSet(lfq64TrainData)
    lfq64TrainDataDecoySet = GetTrainDataSet(lfq64DecoyData)
    tims20211002TrainDataSet = GetTrainDataSet(tims20211002TrainData)
    tims20211002TrainDataDecoySet = GetTrainDataSet(tims20211002DecoyData)
    Trypsin_HFXTrainDataSet = GetTrainDataSet(Trypsin_HFXTrainData)
    Trypsin_HFXTrainDataDecoySet = GetTrainDataSet(Trypsin_HFXDecoyData)

    lfq32trainDataSet, lfq32valDataSet, lfq32valDataDecoySet = SplitTrainAndValDataSet(
        lfq32TrainDataSet, lfq32TrainDataDecoySet, 10000, 10000
    )

    lfq64trainDataSet, lfq64valDataSet, lfq64valDataDecoySet = SplitTrainAndValDataSet(
        lfq64TrainDataSet, lfq64TrainDataDecoySet, 10000, 10000
    )

    tims20211002trainDataSet, tims20211002valDataSet, tims20211002valDataDecoySet = SplitTrainAndValDataSet(
        tims20211002TrainDataSet, lfq64TrainDataDecoySet, 1000, 1000
    )

    (
        Trypsin_HFXtrainDataSet,
        Trypsin_HFXvalDataSet,
        Trypsin_HFXvalDataDecoySet,
    ) = SplitTrainAndValDataSet(
        Trypsin_HFXTrainDataSet, Trypsin_HFXTrainDataDecoySet, 1000, 1000
    )

    trainDataSet = ConcatDataset(
        [lfq32trainDataSet, lfq64trainDataSet, tims20211002trainDataSet, Trypsin_HFXtrainDataSet])

    valDataSet, valDataDecoySet = [lfq32valDataSet, lfq64valDataDecoySet, tims20211002valDataSet, Trypsin_HFXvalDataSet], [
        lfq32valDataDecoySet,
        lfq64valDataDecoySet,
        tims20211002valDataSet,
        Trypsin_HFXvalDataDecoySet,
    ]

    trainData, valData, valDataDecoy = GetTrainDataLoader(
        batch_size, trainDataSet, valDataSet, valDataDecoySet
    )
    return trainData, valData, valDataDecoy

def GetTrainDataPath() -> Tuple[str, str, str, str, str, str, str, str]:
    r"""
    获取各训练数据路径, 包含以下各类数据
    * LFQ32fix
    * LFQ64var
    * Trypsin_HFX 血浆
    * tims 血浆
    """
    root = "/data/xp/SpectronautLabelTrainAndTest/"

    lfq32TrainPath = (
        root
        + "LFQ32/trainData/HYE124_TTOF6600_32fix_train_win_data_minmz5_identify.npy"
    )

    lfq32TrainDecoyPath = (
        root
        + "LFQ32/trainData/HYE124_TTOF6600_32fix_train_win_data_minmz5_identify_decoy.npy"
    )

    lfq64TrainPath = (
        root
        + "LFQ64/trainData/HYE124_TTOF6600_64var_train_win_data_minmz5_identify.npy"
    )

    lfq64TrainDecoyPath = (
        root
        + "LFQ64/trainData/HYE124_TTOF6600_64var_train_win_data_minmz5_identify_decoy.npy"
    )

    Trypsin_HFXtrainpath = (
        root
        + "20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX/trainData/20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX_train_win_data_minmz5_identify.npy"
    )

    Trypsin_HFXtrainDecoyPath = (
        root
        + "20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX/trainData/20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX_train_win_data_minmz5_identify_decoy.npy"
    )

    tims20211002TrainPath = (
        root
        + "tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc/trainData/tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc_train_win_data_minmz5_identify.npy"
    )

    tims20211002TrainDecoyPath = (
        root
        + "tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc/trainData/tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc_train_win_data_minmz5_identify_decoy.npy"
    )

    return (
        lfq32TrainPath,
        lfq32TrainDecoyPath,
        lfq64TrainPath,
        lfq64TrainDecoyPath,
        Trypsin_HFXtrainpath,
        Trypsin_HFXtrainDecoyPath,
        tims20211002TrainPath,
        tims20211002TrainDecoyPath,
    )

def GetTestDataPath() -> (
    Tuple[str, str, str, str, str, str, str, str, str, str, str, str]
):
    r"""
    获取各训练数据路径, 包含以下各类数据
    * LFQ32fix
    * LFQ64var
    * Trypsin_HFX 血浆
    * tims 血浆
    * Hela 数据
    * 无色谱 lfq
    """
    root = "/data/xp/SpectronautLabelTrainAndTest/"

    lfq32TestPath = (
        root + "LFQ32/testData/HYE124_TTOF6600_32fix_test_win_data_minmz5_identify.npy"
    )

    lfq32TestDecoyPath = (
        root
        + "LFQ32/testData/HYE124_TTOF6600_32fix_test_win_data_minmz5_identify_decoy.npy"
    )

    lfq64TestPath = (
        root + "LFQ64/testData/HYE124_TTOF6600_64var_test_win_data_minmz5_identify.npy"
    )

    lfq64TestDecoyPath = (
        root
        + "LFQ64/testData/HYE124_TTOF6600_64var_test_win_data_minmz5_identify_decoy.npy"
    )

    Trypsin_HFXtestpath = (
        root
        + "20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX/testData/20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX_test_win_data_minmz5_identify.npy"
    )

    Trypsin_HFXtestDecoyPath = (
        root
        + "20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX/testData/20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX_test_win_data_minmz5_identify_decoy.npy"
    )

    tims20211002TestPath = (
        root
        + "tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc/testData/tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc_test_win_data_minmz5_identify.npy"
    )

    tims20211002TestDecoyPath = (
        root
        + "tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc/testData/tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc_test_win_data_minmz5_identify_decoy.npy"
    )

    HelaTestPath = root + "Hela/testData/Hela_test_win_data_minmz5_identify.npy"

    HelaTestDecoyPath = (
        root + "Hela/testData/Hela_test_win_data_minmz5_identify_decoy.npy"
    )

    NoRTtestPath = root + "NoRTLfq/testData/NoRT_LFQ_test_win_data_minmz5_identify.npy"

    NoRTtestDecoyPath = (
        root + "NoRTLfq/testData/NoRT_LFQ_test_win_data_minmz5_identify_decoy.npy"
    )

    return (
        lfq32TestPath,
        lfq32TestDecoyPath,
        lfq64TestPath,
        lfq64TestDecoyPath,
        Trypsin_HFXtestpath,
        Trypsin_HFXtestDecoyPath,
        tims20211002TestPath,
        tims20211002TestDecoyPath,
        HelaTestPath,
        HelaTestDecoyPath,
        NoRTtestPath,
        NoRTtestDecoyPath,
    )

def GetSpectronautResultPath() -> Tuple[str, str, str, str]:
    r"""
    获取 Spectronaut 结果的路径
    """
    root = "/data/xp/SpectronautLabelTrainAndTest/"
    lfq32Path = root + "LFQ32/LFQ32_Spectronaut_HYE124_TTOF6600_32fix_identify_all_pep.npy"
    lfq64Path = root + "LFQ64/LFQ64_Spectronaut_HYE124_TTOF6600_64var_quantify_all_pep.npy"
    tims20211002Path = root + \
        "tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc/plasma_1_2_3_4_6_7_8_9_Spectronaut_identify_all_pep.npy"
    Trypsin_HFXPath = root + \
        "20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX/20230111_plasma_24QC_Spectronaut_identify_all_pep.npy"
    return lfq32Path, lfq64Path, tims20211002Path, Trypsin_HFXPath


import heapq
from pathlib import Path
from UnionFind import UnionFind

def cluster(adj_list, n, k):
    """Creates k clusters of the n points listed in the provided adjacency list."""

    # Union-find data structure to represent clusters
    parents = [i for i in range(1, n+1)]
    clusters = UnionFind(parents)

    # Heap for quick access of minimum separation
    priority_tuples = [priority_tuple(edge) for edge in adj_list]
    heapq.heapify(priority_tuples)

    min_separation, endpoints = heapq.heappop(priority_tuples)
    while n > k:
        if clusters.find(endpoints[0]) != clusters.find(endpoints[1]):
            clusters.union(endpoints[0], endpoints[1])

            while clusters.find(endpoints[0]) == clusters.find(endpoints[1]):
                min_separation, endpoints = heapq.heappop(priority_tuples)

            n -= 1

    return min_separation

def priority_tuple(edge):
    return (edge[2], [edge[0], edge[1]])

if __name__ == '__main__':
    path = Path(__file__ + '../..').resolve()
    adj_list = []
    with open(Path(path, 'clustering1.txt')) as file:
        for line in file:
            adj_list.append([int(x) for x in line.split()])

    n = adj_list.pop(0)[0]
    min_separation = cluster(adj_list, n, 4)

    print(min_separation)