#!/usr/bin/env python3.8

## LIBRARIES
import argparse
from Bio import SeqIO
import sys
import time
import datetime
import glob
import os
from PDOR_lib import core_snps_2_fasta
from PDOR_lib import core_snps_list_path

## FUNCTIONS
class MyParser(argparse.ArgumentParser):
    def error(self, message):
        sys.stderr.write('error: %s\n' % message)
        self.print_help()
        sys.exit(2)

def parse_args():

    parser=MyParser(usage='%(prog)s -q QUERY_FOLDER -sd SD_SKETCH -ref REFERENCE_GENOME -snp_thr SNP_THRESHOLD [options]',
        prog='P-DOR.py',
        argument_default=argparse.SUPPRESS,
        epilog="According to the legend, P-dor is the Son of K-mer, but it also likes SNPs")
    requiredNamed = parser.add_argument_group('Input data (all required, except when in TEST mode)')
    requiredNamed.add_argument('-q', help='query folder containing genomes in .fna format',metavar="<dirname>",dest="query_folder",type=str,required=False,default=None)
    requiredNamed.add_argument("-sd", help="Source Dataset (SD) sketch file",metavar="<dirname>",dest="db_sketch", required=False,default=None)
    requiredNamed.add_argument("-ref", help="reference genome", dest="ref", metavar="<filename>",required=False,default=None)
    requiredNamed.add_argument("-snp_thr", dest="snp_threshold", help="Threshold number of SNPs to define an epidemic cluster",required=False,metavar="<int>",default=None)

    optional = parser.add_argument_group('Additional arguments')
    optional.add_argument("-amrf", help="if selected, P-DOR uses amrfinder-plus to search for antimicrobial resistance and virulence genes in the entire Analysis Dataset [Default: %(default)s]",action="store_true",required=False)
    optional.add_argument("-meta", help="metadata file; see example file for formatting [Default: %(default)s]",default=None,required=False)
    optional.add_argument("-sd_folder", help="folder containing the genomes from which the Source Dataset (SD) sketch was created [Default: P-DOR downloads the background genomes from BV-BRC]",required=False,default="",dest="bkg_folder",nargs="?")
    optional.add_argument("-borders", type=int, help="length of the regions at the contig extremities from which SNPs are not called [Default: %(default)s]",metavar="<int>",default=20)
    optional.add_argument("-min_contig_length", type=int, help="minimum contig length to be retained for SNPs analysis [Default: %(default)s]",metavar="<int>",default=500)
    optional.add_argument("-snp_spacing", type=int, help="Number of bases surrounding the mutated position to call a SNP [Default: %(default)s]",metavar="<int>",default=10)
    optional.add_argument('-species', help="full species name of the genomes in the query folder [Default: %(default)s]\nNeeds to be written within single quotes, e.g.: -species 'Klebsiella pneumoniae'.\nThis option is used for a quicker and more precise gene search with AMRFinderPlus", default="")
    optional.add_argument('-n', type=int,help='Maximum closest genomes from Source Dataset (SD) [Default: %(default)s]',metavar="<int>",dest="near",default=20)
    optional.add_argument('-t', type=int,help='number of threads [Default: %(default)s]',metavar="<int>",dest="threads",default=2)
    parser.add_argument('-TEST', action= 'store_true',help='TEST MODE. Warning: this option overrides all other arguments', required=False)
    parser.add_argument('-v', '--version', action='version', version='P-DOR v1.1')

    parser.set_defaults(amrf=False)
    parser.set_defaults(TEST=False)
    args = parser.parse_args()
    if len(sys.argv)==1:
        parser.print_help()
        sys.exit()
    elif (not args.TEST and (args.query_folder is None or args.db_sketch is None or args.ref is None or args.snp_threshold is None)):
        parser.error("\n\nERROR!! If not in TEST mode, the following arguments are required: -q, -sd, -ref, -snp_thr\n\n")
    return args

def logfile():
	with open("PDOR.log","w") as p:
		for arg, value in sorted(vars(args).items()):
			p.write("%s\t%s\n" %(arg,value))

def parse_species(sp):
    orglist=["Acinetobacter_baumannii", "Burkholderia_cepacia", "Burkholderia_pseudomallei", \
    "Campylobacter", "Clostridioides_difficile", "Enterococcus_faecalis", "Enterococcus_faecium", \
    "Escherichia", "Klebsiella_oxytoca", "Klebsiella_pneumoniae", "Neisseria_gonorrhoeae", \
    "Neisseria_meningitidis", "Pseudomonas_aeruginosa", "Salmonella", "Staphylococcus_aureus", \
    "Staphylococcus_pseudintermedius", "Streptococcus_agalactiae", "Streptococcus_pneumoniae", \
    "Streptococcus_pyogenes", "Vibrio_cholerae"]
    organism=""
    if sp.strip()!="":
        test="_".join(sp.strip().split())
        if test in orglist:
            organism=test
        else:
            test=sp.strip().split()[0]
            if test in orglist:
                organism=test
    return organism

def check_res_vir(threads,AD_folder,i_ref):
	print ("Checking for resistance and virulence genes...")
	path_res=path_dir+"/"+Results_folder_name+"/"+AD_folder
	os.chdir(path_res)
	os.system("cp ../%s ." %(i_ref))
	os.system("amrfinder --update")
	organism=parse_species(args.species)
	if organism=="":
		os.system("ls *fna | parallel -j %i 'amrfinder -n {} --plus --threads 1 -o {}.txt --name {}'" %(threads))
	else:
		os.system("ls *fna | parallel -j %i 'amrfinder -n {} --plus --threads 1 -o {}.txt -O %s --name {}'" %(threads,organism))
	os.system("cat *fna.txt >summary_resistance_virulence.txt")
	os.system("rm *fna.txt")
	os.system("rm %s/%s" %(path_res,i_ref.split("/")[-1]))

	os.system("mv summary_resistance_virulence.txt %s/%s" %(path_dir,Results_folder_name))
	os.chdir("..")

def Mummer_snp_call(threads,i_ref,Align_folder):
	print("Aligning genomes with Mummer4...\n")
	os.chdir(path_dir+"/"+Results_folder_name+"/"+Align_folder)
	os.system('ls *fna | parallel -j %i "nucmer --maxgap=500 -p {} %s {}; delta-filter -1 {}.delta > {}_filtered.delta; show-snps -H -C -I -r -T {}_filtered.delta | cut -f1,2,3 >{}.snp"' %(threads,i_ref))
	os.system('ls *_filtered.delta | parallel -j %i "dnadiff -d {} -p {}_info >/dev/null 2>&1"' %(threads))
	reportFiles=glob.glob("*.report")
	rFdict={}
	for rF in reportFiles:
		rFcontent=open(rF,"r").read().split("\n")
		rFdict[rF.strip().replace(".fna_filtered.delta_info.report","")]=float([t.strip().split()[-1].split("(")[1].split("%")[0] for t in rFcontent if "AlignedBase" in t][0])
	with open("Coverage_report.txt","w") as rFout:
		sleep=False
		for t in rFdict.keys():
			if rFdict[t]<70:
				sleep=True
				print("\n\nWARNING! Genome %s aligns poorly to the REFERENCE GENOME (%f %% horizontal coverage). Please, check!\n\n" %(t,rFdict[t]))
			rFout.write("%s\t%f\n" %(t,rFdict[t]))
		if sleep==True:
			time.sleep(10)
	#sys.exit()
	os.chdir(path_dir+"/"+Results_folder_name)
	os.system("ls %s/*.snp >SNP_genomes.list" %(Align_folder))
	core_snps_list_path("SNP_genomes.list", snp_spacing, "SNP_positions")
	core_snps_2_fasta("SNP_genomes.list", "SNP_positions.core.list",i_ref,"SNP_alignment")
	if os.stat("SNP_alignment.core.fasta").st_size > 10:
		print ("Core-SNPs aligmment detected!\n")
	else:
                sys.exit("\nERROR: Core-SNPs alignment not present, check your input files!\n")

### MAIN
print('\n\n      (/T\)  _____\n      (-,-) ()____)\n      \(o)/   -|3       BEHOLD! This is P-DOR!\n     /="""===//||\n,OOO//|   |    ""\nO:O:O LLLLL\n\OOO/ || ||\n     C_) (_D\n\n')
args = parse_args()
time_now=datetime.datetime.now()
datestamp=time_now.strftime("%Y-%m-%d_%H-%M-%S")
start_time = time.time()
prog_dir = os.path.abspath( os.path.dirname( sys.argv[0]))
Results_folder_name="Results_%s" %(datestamp)

##CHECK IF TEST MODE
if args.TEST==True:
    Results_folder_name="Test_%s" %(datestamp)
    args.amrf=False
    args.bkg_folder='%s/data/test_DB/' %(prog_dir)
    args.borders=20
    args.db_sketch='%s/data/sketches.msh' %(prog_dir)
    args.meta='%s/data/sample_metadata_table.txt' %(prog_dir)
    args.min_contig_length=500
    args.near=5
    args.query_folder='%s/data/test_query/' %(prog_dir)
    args.ref='%s/data/NJST258.fna' %(prog_dir)
    args.snp_spacing=10
    args.snp_threshold='20'
    args.species='Klebsiella pneumoniae'
    args.threads=2
    print("\n\nTesting AMRFinderPlus function on a single genome")
    os.system("amrfinder --update")
    organism=parse_species(args.species)
    os.system("amrfinder -n %s --plus --threads %i -o test.amrftest -O %s --name REF" %(args.ref,args.threads,organism))
    os.system("head -n2 test.amrftest")
    os.system("rm test.amrftest")

os.mkdir(Results_folder_name)
logfile()
path_dir = os.path.abspath( os.path.dirname( Results_folder_name))

os.system("mv PDOR.log %s" %(Results_folder_name))

query_folder=args.query_folder
db_sketch=os.path.abspath(args.db_sketch)
ref=args.ref
threads=args.threads
nearest=args.near
borders=args.borders
contig_length=args.min_contig_length
snp_spacing=args.snp_spacing
bkg_folder=args.bkg_folder
if args.meta is not None:
    ABSmeta=os.path.abspath(args.meta)

ref=os.path.abspath(ref)
ABS_query_folder=os.path.abspath(query_folder)

exts=['*.fasta', '*.fna', '*.fa']

files = [f for ext in exts for f in glob.glob(os.path.join(ABS_query_folder, ext))]

genomes_query = [i.strip().split("/")[-1] for i in files]
#genomes_query=[i.replace(".fasta",".fna").replace(".fa",".fna") for i in genomes_query]

print ("\nChecking input formats...\n")
### CHECK N 1
# REF has 1 contig
num_cont_c1=int(os.popen('grep -c ">" %s' %(ref)).read().strip())
if num_cont_c1>1:
	sys.exit('\nERROR: the reference file needs to contain only one sequence!')
### CHECK N 2
# REF is in fasta format
elif num_cont_c1==0:
	sys.exit('\nERROR: the reference file needs to be a FASTA file!')
### CHECK N 3
if db_sketch.split(".")[-1]!="msh":
	sys.exit("\n\nERROR: please check you sketch file!!!")
### CHECK N 4
for gen in genomes_query:
	if gen.endswith(".fasta") or gen.endswith(".fna") or gen.endswith(".fa"):
		n_genomes=len(genomes_query)
		print ("Detected %i genomes in the query folder\n" %n_genomes)
		break
	else:
		sys.exit("\n\nERROR: your query genome folder does not contain fasta files!!!")

### format REF
alt_ref=SeqIO.read(ref,"fasta")
alt_ref.id="REF_%s" %(str(alt_ref.id).split(".")[0])
SeqIO.write(alt_ref,"%s/%s.fna" %(Results_folder_name,alt_ref.id),"fasta")

for gen_names in genomes_query:
	ext=gen_names.strip().split(".")[-1]
	fname=".".join(gen_names.strip().split(".")[:-1])
	if (ext == "fa" or ext == "fasta") :
		os.rename ("%s/%s" %(ABS_query_folder,gen_names), "%s/%s.fna" %(ABS_query_folder,fname))

NEAREST=[]
for i in genomes_query:
	cmd="mash dist %s/%s %s -p %i 2>/dev/null | sort -gr -k5 | head -n %i | cut -f2"  %(ABS_query_folder,i,db_sketch,threads,nearest)
	NEAREST=NEAREST+(os.popen(cmd).read().strip().strip('\n').split('\n'))

NEAREST=list(set(NEAREST))
NEAREST=[i.split("/")[-1] for i in NEAREST]
print ("Sketches completed!\n")

os.chdir(Results_folder_name)
with open("query_genome_list.txt","w") as q:
	for i in genomes_query:
		q.write("%s\n" %(i))
with open("backgroud_genome_list.txt","w") as q:
	for i in NEAREST:
		q.write("%s\n" %(i))
os.mkdir("Background")

if args.bkg_folder:
	print ("Retrieving nearest genomes from the %s folder\n" %args.bkg_folder)
	os.chdir(path_dir)
	background_folder=os.path.abspath(bkg_folder)
	os.chdir(Results_folder_name)
	for N in NEAREST:
		cmd="cp %s/%s Background/" %(background_folder,N)
		os.system(cmd)
else:
	os.chdir("Background")
	print ("Retrieving nearest genomes from the BV-BRC Database...\n")
	for N in NEAREST:
		N=N.replace(".fna","").replace(".fasta","").replace(".fa","")
		cmd = "wget -qN ftp://ftp.bvbrc.org/genomes/%s/%s.fna" %(N,N)
		os.system(cmd)
	os.chdir(path_dir+"/"+Results_folder_name)

os.mkdir("Align")
os.mkdir("Analysis_Dataset")

for i in glob.glob("Background/*"):
	name=i.strip().split("/")[-1]
	oF=open("Align/BD_%s" %(name),"w")
	origsize=0
	trimmedsize=0
	for x in SeqIO.parse(i,"fasta"):
		origsize+=len(str(x.seq))
		if len(str(x.seq))>=contig_length+2*borders:
			x.seq=x.seq[borders:-borders]
			trimmedsize+=len(str(x.seq))
			SeqIO.write(x,oF,"fasta")
	oF.close()
	if trimmedsize<float(origsize)*0.90:
		print("\nDatabase genome %s was too short after contig polishing and was excluded from the Analysis Dataset (AD)" %(name))
		os.system("rm Align/BD_%s" %(name))
	else:
		os.system("cp %s Analysis_Dataset/BD_%s" %(i,name))

for i in glob.glob("%s/*" %(ABS_query_folder)):
	name=i.strip().split("/")[-1]
	oF=open("Align/%s" %(name),"w")
	origsize=0
	trimmedsize=0
	for x in SeqIO.parse(i,"fasta"):
		origsize+=len(str(x.seq))
		if len(str(x.seq))>=contig_length+2*borders:
			x.seq=x.seq[borders:-borders]
			trimmedsize+=len(str(x.seq))
			SeqIO.write(x,oF,"fasta")
	oF.close()
	if trimmedsize<float(origsize)*0.90:
		sys.exit("\n\nERROR: your query genome %s was too short after contig polishing, please check your genome quality! EXECUTION ABORTED" %(name))
	else:
		os.system("cp %s Analysis_Dataset/%s" %(i,name))

aln_folder_size=os.popen("ls Align/ | wc -l").read().strip()
os.system("rm -rf Background")

if args.amrf==True:
    check_res_vir(threads,"Analysis_Dataset","%s.fna" %(alt_ref.id))

alt_ref_abs=os.path.abspath("%s.fna" %(alt_ref.id))
Mummer_snp_call(threads,alt_ref_abs,"Align")
os.mkdir("SNPs")
os.system("mv Align/*.snp SNPs/")
os.system("mv Align/Coverage_report.txt .")
snp_num=os.popen("ls SNPs/*.snp | wc -l").read().strip()
if int(snp_num)!=int(aln_folder_size):
    sys.exit("\n\nERROR! Something went wrong! Not all .snp files were produced")
os.system("rm -rf Align")

print("Performing phylogenetic reconstruction...\n")
cmd="iqtree -s SNP_alignment.core.fasta -pre SNP_alignment -m MFP+GTR+ASC -bb 1000 -nt %i" %(threads)
os.system(cmd)

os.system("$CONDA_PREFIX/bin/Rscript %s/Snpbreaker.R SNP_alignment.core.fasta %s" %(prog_dir,args.snp_threshold))

if args.amrf==True:
    emptyAMRF=os.popen("cat summary_resistance_virulence.txt | sort | uniq | wc -l").read().strip()
    if int(emptyAMRF)<=1:
        os.system("rm summary_resistance_virulence.txt")

os.system("$CONDA_PREFIX/bin/Rscript %s/annotated_tree.R SNP_alignment.treefile" %(prog_dir))

if args.meta is not None:
	os.system("$CONDA_PREFIX/bin/Rscript %s/contact_network.R %s" %(prog_dir,ABSmeta))

###ORGANIZE OUTPUT
os.mkdir("Other_output_files")
os.system("mv SNPs Other_output_files/SNPs")
os.system("mv Analysis_Dataset Other_output_files/Analysis_Dataset/")
os.system("mv SNP_alignment* Other_output_files/")

os.system("mv backgroud_genome_list.txt Other_output_files/")
os.system("mv clusters_manual_threshold_%s_snps.csv Other_output_files/" %(args.snp_threshold))
os.system("mv query_genome_list.txt Other_output_files/")
os.system("mv snp_distance_matrix.tsv Other_output_files/")
os.system("mv SNP_genomes.list Other_output_files/")
os.system("mv SNP_positions.core.list Other_output_files/")
if args.amrf==True:
    os.system("mv summary_resistance_virulence.txt Other_output_files/")

os.chdir(path_dir)
os.system("rm %s" %(alt_ref_abs))

time=float(time.time() - start_time)
hours=int(time/3600)
minutes=((time/3600)-hours)*60
print ("\n\n\n\n")
print ("####################\n")
print ("Analysis completed in %i hours and %f minutes" %(hours,minutes))
print ("Results are stored in folder %s/%s" %(path_dir,Results_folder_name))
print ("####################\n")


#!/usr/bin/env python3
# -*-coding:utf8-*-
# !/bin/bash

from enum import Enum

class BoundaryType(Enum):
    # Inlet
    VELOCITY_INLET	    = 'velocityInlet'
    FLOW_RATE_INLET	    = 'flowRateInlet'
    PRESSURE_INLET	    = 'pressureInlet'
    ABL_INLET	        = 'ablInlet'
    OPEN_CHANNEL_INLET  = 'openChannelInlet'
    FREE_STREAM	        = 'freeStream'
    FAR_FIELD_RIEMANN	= 'farFieldRiemann'
    SUBSONIC_INFLOW	    = 'subsonicInflow'
    SUPERSONIC_INFLOW	= 'supersonicInflow'

    # Outlet
    PRESSURE_OUTLET	    = 'pressureOutlet'
    OPEN_CHANNEL_OUTLET = 'openChannelOutlet'
    OUTFLOW	            = 'outflow'
    SUBSONIC_OUTFLOW	= 'subsonicOutflow'
    SUPERSONIC_OUTFLOW	= 'supersonicOutflow'

    # Wall
    WALL	            = 'wall'
    THERMO_COUPLED_WALL	= 'thermoCoupledWall'
    POROUS_JUMP	        = 'porousJump'
    FAN	                = 'fan'

    # Internal
    SYMMETRY	        = 'symmetry'
    INTERFACE	        = 'interface'
    EMPTY	            = 'empty'
    CYCLIC	            = 'cyclic'
    WEDGE	            = 'wedge'

# https://www.openfoam.com/documentation/user-guide/4-mesh-generation-and-conversion/4.2-boundaries
class BasicPatchType(Enum):
    BASIC_TYPE_PATCH            = 'patch'
    BASIC_TYPE_SYMMETRY_PLANE   = 'symmetryPlane'
    BASIC_TYPE_EMPTY            = 'empty'
    BASIC_TYPE_WEDGE            = 'wedge'
    BASIC_TYPE_CYCLIC           = 'cyclic'
    BASIC_TYPE_PROCESSOR        = 'processor'


from datetime import datetime
from django.db import models
from django.contrib.auth.models import User

class ResetData(models.Model):
    date_time = models.DateTimeField()
    end_date_time = models.DateTimeField(default=datetime.utcfromtimestamp(2147483640))
    max_clips = models.IntegerField(default=2)
    ranks = models.IntegerField(default=5)
    user_created_clip = models.BooleanField(default=True)
    clip_newer_than = models.DateTimeField(default=datetime.utcfromtimestamp(1))

class AllowedChannel(models.Model):
    broadcaster_id = models.CharField(max_length=100)
    reset_data = models.ForeignKey(ResetData, on_delete=models.CASCADE)

    class Meta:
        unique_together = (
            "broadcaster_id",
            "reset_data",
        )

class Clip(models.Model):
    id = models.CharField(primary_key=True, max_length=100, unique=True)
    url = models.URLField(null=True)
    embed_url = models.URLField(null=True)
    broadcaster_id = models.CharField(null=True, max_length=100)
    broadcaster_name = models.CharField(null=True, max_length=100)
    creator_id = models.CharField(null=True, max_length=100)
    creator_name = models.CharField(null=True, max_length=100)
    video_id = models.CharField(null=True, max_length=100)
    game_id = models.CharField(null=True, max_length=100)
    title = models.CharField(null=True, max_length=200)
    view_count = models.IntegerField(null=True)
    created_at = models.DateTimeField(null=True)
    thumbnail_url = models.URLField(null=True)
    duration = models.FloatField(null=True)
    vod_offset = models.IntegerField(null=True)
    date_added = models.DateTimeField(null=True)
    account = models.ForeignKey(User, on_delete=models.CASCADE)
    rank = models.IntegerField(null=True)


# -*- coding: utf-8 -*-
# @Author  : ydf
# @Time    : 2022/8/20 0008 12:12

import os

if os.name == 'nt':
    """
    为了保险起见，这样做一下,设置一下path，否则anaconda安装的python可能出现 ImportError: DLL load failed while importing cimpl: 找不到指定的模块。
    多设置没事，少设置了才麻烦。
    """
    from pathlib import Path
    import sys

    # print(sys.executable)  #F:\minicondadir\Miniconda2\envs\py38\python.exe
    # print(os.getenv('path'))
    python_install_path = Path(sys.executable).parent.absolute()
    kafka_libs_path = python_install_path / Path(r'.\Lib\site-packages\confluent_kafka.libs')
    dlls_path = python_install_path / Path(r'.\DLLs')
    library_bin_path = python_install_path / Path(r'.\Library\bin')
    # print(library_bin_path)
    path_env = os.getenv('path')
    os.environ['path'] = f'''{path_env};{kafka_libs_path};{dlls_path};{library_bin_path};'''

import atexit
import time
# noinspection PyPackageRequirements
from kafka import KafkaProducer, KafkaAdminClient
# noinspection PyPackageRequirements
from kafka.admin import NewTopic
# noinspection PyPackageRequirements
from kafka.errors import TopicAlreadyExistsError

from funboost import funboost_config_deafult
from funboost.publishers.base_publisher import AbstractPublisher

class ConfluentKafkaPublisher(AbstractPublisher, ):
    """
    使用kafka作为中间件，这个confluent_kafka包的性能远强于 kafka-pyhton
    """

    # noinspection PyAttributeOutsideInit
    def custom_init(self):
        from confluent_kafka import Producer as ConfluentProducer  # 这个包不好安装，用户用这个中间件的时候自己再想办法安装。win用户需要安装c++ 14.0以上环境。
        # self._producer = KafkaProducer(bootstrap_servers=funboost_config_deafult.KAFKA_BOOTSTRAP_SERVERS)
        try:
            admin_client = KafkaAdminClient(bootstrap_servers=funboost_config_deafult.KAFKA_BOOTSTRAP_SERVERS)
            admin_client.create_topics([NewTopic(self._queue_name, 10, 1)])
            # admin_client.create_partitions({self._queue_name: NewPartitions(total_count=16)})
        except TopicAlreadyExistsError:
            pass
        except BaseException as e:
            self.logger.exception(e)
        atexit.register(self.close)  # 程序退出前不主动关闭，会报错。
        self._confluent_producer = ConfluentProducer({'bootstrap.servers': ','.join(funboost_config_deafult.KAFKA_BOOTSTRAP_SERVERS)})
        self._recent_produce_time = time.time()

    # noinspection PyAttributeOutsideInit
    def concrete_realization_of_publish(self, msg):
        # noinspection PyTypeChecker
        # self.logger.debug(msg)
        self._confluent_producer.produce(self._queue_name, msg.encode(), )
        if time.time() - self._recent_produce_time > 1:
            self._confluent_producer.flush()
            self._recent_produce_time = time.time()

    def clear(self):
        self.logger.warning('还没开始实现 kafka 清空 消息')
        # self._consumer.seek_to_end()
        # self.logger.warning(f'将kafka offset 重置到最后位置')

    def get_message_count(self):
        return -1  # 还没找到获取所有分区未消费数量的方法。

    def close(self):
        pass
        # self._confluent_producer.

    def _at_exit(self):
        # self._producer.flush()
        self._confluent_producer.flush()
        super()._at_exit()


import os
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score
from sklearn.metrics import accuracy_score,log_loss
import lightgbm as lgb
from config import *
from tfidfModel import *

staticFeatureFile = '/home/malware/malwaretc/file_features.csv'
def fileFeatures(inputfile):
    df = pd.read_csv(inputfile)
    print(df.columns)
    x = df[['max_thread_len','min_thread_len','mean_thread_len','thread_num']]
    fileIDs = df['file_id']

    return x,fileIDs

def fileLabels(inputfile):
    df = pd.read_csv(inputfile)
    return df['label']

def trainModel():
    x,fileIds = fileFeatures(staticFeatureFile)
    row,col = x.shape
    tfFeatures = loadTFFeatures()
    x= pd.concat([x,tfFeatures.iloc[0:row,:]],axis=1)
    y = fileLabels(staticFeatureFile)
    x_train,x_test,y_train,y_test=train_test_split(x,y,test_size=.3)
    train_data=lgb.Dataset(x_train,y_train)

    param = {'objective':'multiclass','learning_rate':.05,'num_class':6,'max_bin':300}
    param.update({'max_depth':10,'num_leaves':550 })
    param['metric'] = ['multi_logloss']

    num_round=1000
    # test_data = lgb.Dataset(x_test,y_test)
    test_data = lgb.Dataset(x_test,y_test, reference=train_data)

    # lgbm=lgb.train(param,train_data,num_round)
    lgbm=lgb.train(param,train_data,num_round,valid_sets=[test_data], early_stopping_rounds=10)
    checkModel(y_test,lgbm.predict(x_test))
    lgbm.save_model('model.txt', num_iteration=lgbm.best_iteration)

if __name__=='__main__':
    trainModel()


'''
Codé par : RONGERE Julien et Rodde Théophile
Date : 05/12/202
Version : 1.0
Description : Class des aliens desquelles vont découlés les autres aliens
'''
import tkinter as tk
import threading
import time
import missile
from pygame import mixer
class alien:
    def __init__(self,master,canvas,fenetre,pv,vitesse,position,temps_recharge,etat) -> None:

        self.master=master
        self.canvas=canvas
        self.fenetre=fenetre

        self.pv = pv
        self.vitesse = vitesse
        self.temps_recharge = temps_recharge
        self.etat = etat #mort ou vivant

        self.display_bbox_bool = False
        self.bbox_rect = None

        self.position = position

        self.missiles = []

        self.nom = self.__class__.__name__ #type de l'alien
        self.image = tk.PhotoImage(file=f"images/aliens/{self.nom}.png")
        self.bruitages = {"explosion":mixer.Sound("musiques/bruitages/explosion_alien.mp3")}

        self.a_spawn = False # l'alien a t-il spawn ou non

    def collision_missile(self):
        #fonction qui gere la collision entre les missiles et les aliens

        if self.etat == "mort":
            return

        for missile in self.master.vaisseau.missiles:
            missile.nettoyage(self.master.vaisseau) #on profite de la boucle pour nettoyer les missiles qui sont en dehors du canvas
            if missile.etat == "mort":
                continue
            if missile.collision(self):
                self.pv -= 1
                missile.supprimer(self.master.vaisseau)

                score = self.master.score.get()
                self.master.score.set(score+self.score*self.master.niveau.score_multiplieur)

                if self.pv == 0:
                    self.supprimer()
                    break

        self.fenetre.after(50,self.collision_missile)

    def get_bbox(self):
        #retourne les coordonnees de la bbox de l'alien (hitbox)

        return self.canvas.bbox(self.item)

    def supprimer(self):
        #supprime l'alien en activant les bruitages et les effets visuels

        self.fenetre.after(0,self.bruitages["explosion"].play) #joue le bruitage de l'explosion
        self.etat = "mort"
        self.master.aliens.remove(self)
        self.canvas.delete(self.item)
        self.exploser()

    def exploser(self):
        #explose l'alien

        image_explosion = tk.PhotoImage(file="images/aliens/explosion.png")
        image_explosion.img = image_explosion #garde une reference a l'image pour ne pas qu'elle soit supprimée par le garbage collector

        explosion = self.canvas.create_image(self.position[0],self.position[1],image=image_explosion)
        self.canvas.itemconfig(explosion)

        self.fenetre.after(300,lambda: self.canvas.delete(explosion))

    def display_bbox(self):
        #affiche la bbox de l'alien
        if self.display_bbox_bool == False:
            return
        if self.bbox_rectangle != None:
            self.canvas.delete(self.bbox_rectangle)
        bbox = self.get_bbox()
        self.bbox_rectangle = self.canvas.create_rectangle(bbox[0],bbox[1],bbox[2],bbox[3],outline="red")
        self.fenetre.after(50,self.display_bbox)


import logging
from pprint import pformat
from typing import Optional

import tensorflow as tf

from mlpp_lib.callbacks import TimeHistory, EnsembleMetrics
from mlpp_lib.datasets import DataLoader, DataModule

from mlpp_lib.utils import (
    get_callback,
    get_loss,
    get_metric,
    get_model,
    get_optimizer,
    process_out_bias_init,
)

LOGGER = logging.getLogger(__name__)

def get_log_params(param_run: dict) -> dict:
    """Extract a selection of parameters for pretty logging"""
    log_params = {}
    for dimension in param_run["data_partitioning"].keys():
        params = {
            f"partitioning_{dimension}_{name}": value
            for name, value in param_run["data_partitioning"][dimension].items()
        }
        log_params.update(params)
    # log_params["features_names"] = param_run["features"]
    # Note to future self: the list of features can easily exceed the maximum length for
    # a logged parameter, so we excluded it. Instead, this is logged as an artifact
    # together with the input run parameters.
    log_params["targets_names"] = param_run["targets"]
    log_params["sample_weights_names"] = param_run.get("sample_weights")
    log_params["event_dims"] = param_run["batching"]["event_dims"]
    log_params["thinning"] = param_run.get("thinning")
    log_params["model_name"] = list(param_run["model"])[0]
    log_params.update(param_run["model"][log_params["model_name"]])
    log_params["loss"] = param_run["loss"]
    return log_params

def train(
    cfg: dict,
    datamodule: DataModule,
    callbacks: Optional[list] = None,
) -> tuple:

    LOGGER.debug(f"run params:\n{pformat(cfg)}")

    datamodule.setup("fit")
    model_config = cfg["model"]
    loss_config = cfg["loss"]

    # prepare model
    event_dims = list(set(datamodule.x.dims) - set(datamodule.batch_dims))
    out_bias_init = process_out_bias_init(
        datamodule.train.x, cfg.get("out_bias_init", "zeros"), event_dims
    )
    model_config[list(model_config)[0]].update({"out_bias_init": out_bias_init})
    input_shape = datamodule.train.x.shape[1:]
    output_shape = datamodule.train.y.shape[1:]
    model = get_model(input_shape, output_shape, model_config)
    loss = get_loss(loss_config)
    metrics = [get_metric(metric) for metric in cfg.get("metrics", [])]
    optimizer = get_optimizer(cfg.get("optimizer", "Adam"))
    model.compile(optimizer=optimizer, loss=loss, metrics=metrics)
    model.summary(print_fn=LOGGER.info)

    # callbacks
    if callbacks is None:
        callbacks = []

    for callback in cfg.get("callbacks", []):
        callback_instance = get_callback(callback)

        if isinstance(callback_instance, EnsembleMetrics):
            callback_instance.add_validation_data((datamodule.val.x, datamodule.val.y))

        callbacks.append(callback_instance)

    time_callback = TimeHistory()
    callbacks.append(time_callback)

    batch_size = cfg.get("batch_size", 1024)
    if datamodule.group_samples is not None:
        block_size = list(datamodule.group_samples.values())[0]
    else:
        block_size = 1
    # modify batch_size so that it becomes a multiple of block_size
    batch_size = (batch_size // block_size) * block_size

    train_dataloader = DataLoader(
        datamodule.train,
        batch_size=batch_size,
        shuffle=cfg.get("shuffle", True),
        block_size=block_size,
    )
    val_dataloader = DataLoader(
        datamodule.val,
        batch_size=batch_size,
        shuffle=False,
        block_size=block_size,
    )

    LOGGER.info("Start training.")
    res = model.fit(
        train_dataloader,
        epochs=cfg.get("epochs", 1),
        validation_data=val_dataloader,
        callbacks=callbacks,
        steps_per_epoch=cfg.get("steps_per_epoch", None),
        verbose=2,
    )
    LOGGER.info("Done! \U0001F40D")

    # we don't need to export loss and metric functions for deployments
    model.compile(optimizer=optimizer, loss=None, metrics=None)

    custom_objects = tf.keras.layers.serialize(model)

    history = res.history
    # for some reasons, 'lr' is provided as float32
    # and needs to be casted in order to be serialized
    for k in history:
        history[k] = list(map(float, history[k]))

    return model, custom_objects, datamodule.standardizer, history


import functools
import time
import psutil

from rptest.services.redpanda import RedpandaService
from ducktape.mark.resource import ClusterUseMetadata
from ducktape.mark._mark import Mark

def cluster(log_allow_list=None,
            check_allowed_error_logs=True,
            check_for_storage_usage_inconsistencies=True,
            **kwargs):
    """
    Drop-in replacement for Ducktape `cluster` that imposes additional
    redpanda-specific checks and defaults.

    These go into a decorator rather than setUp/tearDown methods
    because they may raise errors that we would like to expose
    as test failures.
    """
    def all_redpandas(test):
        """
        Most tests have a single RedpandaService at self.redpanda, but
        it is legal to create multiple instances, e.g. for read replica tests.

        We find all replicas by traversing ducktape's internal service registry.
        """
        yield test.redpanda

        for svc in test.test_context.services:
            if isinstance(svc, RedpandaService) and svc is not test.redpanda:
                yield svc

    def log_local_load(test_name, logger, t_initial, initial_disk_stats):
        """
        Log indicators of system load on the machine running ducktape tests.  When
        running tests in a single-node docker environment, this is useful to help
        diagnose whether slow-running/flaky tests are the victim of an overloaded
        node.
        """
        load = psutil.getloadavg()
        memory = psutil.virtual_memory()
        swap = psutil.swap_memory()
        disk_stats = psutil.disk_io_counters()
        disk_deltas = {}
        disk_rates = {}
        runtime = time.time() - t_initial
        for f in disk_stats._fields:
            a = getattr(initial_disk_stats, f)
            b = getattr(disk_stats, f)
            disk_deltas[f] = b - a
            disk_rates[f] = (b - a) / runtime

        logger.info(f"Load average after {test_name}: {load}")
        logger.info(f"Memory after {test_name}: {memory}")
        logger.info(f"Swap after {test_name}: {swap}")
        logger.info(f"Disk activity during {test_name}: {disk_deltas}")
        logger.info(f"Disk rates during {test_name}: {disk_rates}")

    def cluster_use_metadata_adder(f):
        Mark.mark(f, ClusterUseMetadata(**kwargs))

        @functools.wraps(f)
        def wrapped(self, *args, **kwargs):
            # This decorator will only work on test classes that have a RedpandaService,
            # such as RedpandaTest subclasses
            assert hasattr(self, 'redpanda')

            t_initial = time.time()
            disk_stats_initial = psutil.disk_io_counters()
            try:
                r = f(self, *args, **kwargs)
            except:
                if not hasattr(self, 'redpanda') or self.redpanda is None:
                    # We failed so early there isn't even a RedpandaService instantiated
                    raise

                log_local_load(self.test_context.test_name,
                               self.redpanda.logger, t_initial,
                               disk_stats_initial)

                for redpanda in all_redpandas(self):
                    redpanda.logger.exception(
                        f"Test failed, doing failure checks on {redpanda.who_am_i()}..."
                    )

                    # Disabled to avoid addr2line hangs
                    # (https://github.com/redpanda-data/redpanda/issues/5004)
                    # self.redpanda.decode_backtraces()

                    redpanda.cloud_storage_diagnostics()

                    redpanda.raise_on_crash(log_allow_list=log_allow_list)

                raise
            else:
                if not hasattr(self, 'redpanda') or self.redpanda is None:
                    # We passed without instantiating a RedpandaService, for example
                    # in a skipped test
                    return r

                log_local_load(self.test_context.test_name,
                               self.redpanda.logger, t_initial,
                               disk_stats_initial)

                # In debug mode, any test writing too much traffic will impose too much
                # load on the system and destabilize other tests.  Detect this with a
                # post-test check for total bytes written.
                debug_mode_data_limit = 64 * 1024 * 1024
                if hasattr(self, 'debug_mode') and self.debug_mode is True:
                    bytes_written = self.redpanda.estimate_bytes_written()
                    if bytes_written is not None:
                        self.redpanda.logger.info(
                            f"Estimated bytes written: {bytes_written}")
                        if bytes_written > debug_mode_data_limit:
                            self.redpanda.logger.error(
                                f"Debug-mode test wrote too much data ({int(bytes_written) // (1024 * 1024)}MiB)"
                            )

                for redpanda in all_redpandas(self):
                    redpanda.logger.info(
                        f"Test passed, doing log checks on {redpanda.who_am_i()}..."
                    )
                    if check_allowed_error_logs:
                        # Only do log inspections on tests that are otherwise
                        # successful.  This executes *before* the end-of-test
                        # shutdown, thereby avoiding having to add the various
                        # gate_closed etc errors to our allow list.
                        # TODO: extend this to cover shutdown logging too, and
                        # clean up redpanda to not log so many errors on shutdown.
                        try:
                            redpanda.raise_on_bad_logs(
                                allow_list=log_allow_list)
                        except:
                            redpanda.cloud_storage_diagnostics()
                            raise

                    if check_for_storage_usage_inconsistencies:
                        try:
                            redpanda.raise_on_storage_usage_inconsistency()
                        except:
                            redpanda.cloud_storage_diagnostics()
                            raise

                self.redpanda.validate_controller_log()

                if self.redpanda.si_settings is not None:
                    try:
                        self.redpanda.stop_and_scrub_object_storage()
                    except:
                        self.redpanda.cloud_storage_diagnostics()
                        raise

                # Finally, if the test passed and all post-test checks
                # also passed, we may trim the logs to INFO level to
                # save space.
                for redpanda in all_redpandas(self):
                    redpanda.trim_logs()

                return r

        # Propagate ducktape markers (e.g. parameterize) to our function
        # wrapper
        wrapped.marks = f.marks
        wrapped.mark_names = f.mark_names

        return wrapped

    return cluster_use_metadata_adder


from pprint import pprint
import csv
import pymongo
import json
import urllib.request
import re
import sys
import datetime
from bson import ObjectId
import hashids
from hashlib import blake2b, blake2s
import time

client = pymongo.MongoClient('localhost', 27017)
db = client.USTdev

components = db.stock.aggregate([
    # {"$match": {"_id": ObjectId('5c70984612875079b91f89e9')}},
    {"$lookup": {"from": 'store_positions', "localField":'position.posid', "foreignField": '_id', "as": 'position_info'}},
])

errs = set()

for i, comp in enumerate(components):
    try:
        print("====================================")
        # print(comp)
        position = None
        cid = comp['_id']
        stocks = set()

        # prevodni tabulka ze skladu na nezarazenou pozici
        stock = {
            ObjectId("5c67444e7e875154440cc28f"): None,# Praha
            ObjectId("5c67445b7e875154440cc297"): None,# sob provozovna
            ObjectId("5c67446b7e875154440cc299"): None,# Hvezdarna
        }
        stock_usage = {
            ObjectId("5c67444e7e875154440cc28f"): False,# Praha
            ObjectId("5c67445b7e875154440cc297"): False,# sob provozovna
            ObjectId("5c67446b7e875154440cc299"): False,# Hvezdarna
        }

        ## Ziskat pozice ve skladu (uprednostnit primarni a pak posledni)
        for j, p in enumerate(comp.get('position', [])):
            info = comp['position_info'][j]
            if p['primary']:
                stock[info['warehouse']] = p['posid']

        for j, p in enumerate(comp.get('position', [])):
            info = comp['position_info'][j]
            if not stock[info['warehouse']] or stock[info['warehouse']] in [ObjectId('5c709fa012875079fd76e238')]:
                stock[info['warehouse']] = p['posid']

        if not stock[ObjectId("5c67444e7e875154440cc28f")]: stock[ObjectId("5c67444e7e875154440cc28f")] = ObjectId('5c709fa012875079fd76e238')

        db.stock.update(
            {
                "_id": cid
            },
            {
                "$set": {"packets": []},
                #"$set": {"parameters": {}}
            }
        )

        for j, h in enumerate(comp.get('history', [])):
            if not 'stock' in h: h['stock'] = ObjectId("5c67444e7e875154440cc28f")
            if not stock_usage[h['stock']]:
                pid = ObjectId()
                stock_usage[h['stock']] = pid
                print("Vytvorim novy sacek {}".format(pid))
                query = {
                    "$push": {
                        "packets": {
                            '_id': pid,
                            'type': 'zip_bag',
                            'supplier': 'null',
                            'position': stock[h['stock']],
                            'description': 'Vytvoreno pri aktualizaci databaze'
                        }
                    }
                }
                db.stock.update(
                    {
                        "_id": cid
                    },
                    query
                )

        for j, h in enumerate(comp.get('history', [])):
            print("* ",j, h)
            pos = stock[h['stock']]

            query = {
                    "pid": stock_usage[h['stock']],
                    "count" : h.get('bilance', 0),
                    "unit_price": h.get('price', 0),
                    "type" : h.get('operation', 'service'),
                    "date" : h['_id'].generation_time,
                    "user" : h.get('user', 'null'),
                    "invoice" : h.get('invoice'),
                    "supplier" : h.get('supplier'),
                    "description" : "upgrade, " + h.get('description', '')
                }
            print(query)
            db.stock_operation.insert(query)

            # db.stock.update({
            #     "_id": cid
            # }, {
            #     "$push": {
            #         "packets": {
            #             '_id': pid,
            #             'type': 'zip_bag',
            #             'supplier': 'null',
            #             'position': 'null',
            #             'description': 'Vytvoreno oddelenim z {}'.format(src)
            #         }
            #     }
            # })
        # print(positions)

    except Exception as e:
        print("chyba", e)
        errs.add((comp['_id'], comp['name'], e))

print(i, "celkem mame tolik polozek")
for e in errs:
    print(e)


#!/usr/bin/env python
""" Constellano .XML to .RECORD converter
Python script for .xml to .record (tensorflow) conversation.

Command format:
    py ConvertXmlToRecord.py --images <images_dir> --labels <path_to_label_map> --output <output_dir_name>
"""

import sys
import os
import cv2
import getopt
import numpy as np
from datetime import datetime

import os
import glob
import pandas as pd
import io
import xml.etree.ElementTree as ET
import argparse

#os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'    # Suppress TensorFlow logging (1)
import tensorflow as tf
tf = tf.compat.v1
from PIL import Image
from object_detection.utils import dataset_util, label_map_util
label_map_util.tf = tf.compat.v1
tf.gfile = tf.io.gfile
from collections import namedtuple

__author__ = "Marin Maslov"
__license__ = "MIT Licence"
__version__ = "2.1.2"
__maintainer__ = "Marin Maslov"
__email__ = "mmaslo00@fesb.hr"
__status__ = "Stable"

COMMAND_FORMAT = "Error! The command should be: py ConvertXmlToRecord.py --images <images_dir> --labels <path_to_label_map> --output <output_dir_name>"

def xml_to_csv(path):
    xml_list = []
    for xml_file in glob.glob(path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        filename = root.find('filename').text
        width = int(root.find('size').find('width').text)
        height = int(root.find('size').find('height').text)
        for member in root.findall('object'):
            bndbox = member.find('bndbox')
            value = (filename,
                    width,
                    height,
                    member.find('name').text,
                    int(bndbox.find('xmin').text),
                    int(bndbox.find('ymin').text),
                    int(bndbox.find('xmax').text),
                    int(bndbox.find('ymax').text),
                    )
            xml_list.append(value)
    column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
    xml_df = pd.DataFrame(xml_list, columns=column_name)
    return xml_df

def split(df, group):
    data = namedtuple('data', ['filename', 'object'])
    gb = df.groupby(group)
    return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)]

def class_text_to_int(row_label, label_map_dict):
    return label_map_dict[row_label]

def create_tf_example(group, path, label_map_dict):
    with tf.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid:
        encoded_jpg = fid.read()
    encoded_jpg_io = io.BytesIO(encoded_jpg)
    image = Image.open(encoded_jpg_io)
    width, height = image.size

    filename = group.filename.encode('utf8')
    image_format = b'jpg'
    xmins = []
    xmaxs = []
    ymins = []
    ymaxs = []
    classes_text = []
    classes = []

    for index, row in group.object.iterrows():
        xmins.append(row['xmin'] / width)
        xmaxs.append(row['xmax'] / width)
        ymins.append(row['ymin'] / height)
        ymaxs.append(row['ymax'] / height)
        classes_text.append(row['class'].encode('utf8'))
        classes.append(class_text_to_int(row['class'], label_map_dict))

    tf_example = tf.train.Example(features=tf.train.Features(feature={
        'image/height': dataset_util.int64_feature(height),
        'image/width': dataset_util.int64_feature(width),
        'image/filename': dataset_util.bytes_feature(filename),
        'image/source_id': dataset_util.bytes_feature(filename),
        'image/encoded': dataset_util.bytes_feature(encoded_jpg),
        'image/format': dataset_util.bytes_feature(image_format),
        'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
        'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
        'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
        'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
        'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
        'image/object/class/label': dataset_util.int64_list_feature(classes),
    }))
    return tf_example

def main(argv):
    images = ''
    labels = ''
    output = ''
    name = ''

    try:
        opts, args = getopt.getopt(argv, "h", ["images=", "labels=", "output=", "name="])
    except getopt.GetoptError:
        print(COMMAND_FORMAT)
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-h':
            print(COMMAND_FORMAT)
            sys.exit()
        elif opt in ("--images"):
            images = arg
        elif opt in ("--labels"):
            labels = arg
        elif opt in ("--output"):
            if arg[-1] == "/":
                output = arg
            else:
                output = arg + "/"
        elif opt in ("--name"):
            name = arg

    start_time = datetime.now()

    if not os.path.exists(output) != False:
        print("Creating directory: " + output)
        os.mkdir(output)

    #label_map = label_map_util.load_labelmap(labels)
    #label_map_dict = label_map_util.get_label_map_dict(label_map)

    label_map_dict = label_map_util.get_label_map_dict(labels)

    print(output)
    writer = tf.python_io.TFRecordWriter(output + name + ".record")
    examples = xml_to_csv(images)
    grouped = split(examples, 'filename')

    for group in grouped:
        tf_example = create_tf_example(group, images, label_map_dict)
        writer.write(tf_example.SerializeToString())
    writer.close()

    print('Successfully created the TFRecord file: {}'.format(output))
    print("------------------------------------")
    end_time = datetime.now()
    print("[INFO]\tTotal execution time: " + str(end_time - start_time) + ".")

if __name__ == "__main__":
    main(sys.argv[1:])


import logging
from pathlib import Path
from rich.console import Console
from rich.style import Style
from rich.theme import Theme
from rich.logging import RichHandler
import pandas as pd
import inspect
import sys
from rich.panel import Panel

console = Console(
    width=120,
    theme=Theme({
        "debug": "#cccccc",
        "main": "bold cyan",
        "info": "white",
        "warning": "yellow",
        "error": "bold red",
        "critical": "bold red",
        "success": "bold green",
        "repr.ellipsis": "medium_purple1"}),
)

def output_message(msg, msg_style="debug"):
    EMOJIS = {
        "debug": ":eyes:",
        "info": ":information:",
        "warning": ":warning:",
        "error": ":prohibited:",
        "critical": ":x:",
        "success": ":thumbs_up:"}

    last_line_no = inspect.currentframe().f_back.f_lineno
    func = inspect.currentframe().f_back.f_code
    log_msg = f"{msg}: [{func.co_name} in {func.co_filename}:{last_line_no}]"

    if msg_style == "debug":
        logging.debug(log_msg)
    elif msg_style == "info":
        logging.info(log_msg)
    elif msg_style == "warning":
        logging.warning("WARNING: " + log_msg)
    elif msg_style == "error":
        logging.error("ERROR:" + log_msg)
    elif msg_style == "critical":
        logging.error("CRITICAL:" + log_msg)
    elif msg_style == "success":
        logging.info(log_msg)

    if msg_style != "debug":
        if msg_style in EMOJIS:
            msg = f"{EMOJIS[msg_style]} {msg}"
        console.print(Panel(f"{msg}", style=msg_style))

    if msg_style == 'critical':
        sys.exit()

def comment(df: pd.DataFrame, *, msg: str, level: str = "INFO", **extra) -> pd.DataFrame:
    """
    Log a message.

    Parameters
    ----------
    df : pandas.DataFrame
      dataframe to inject into log messages

    msg : str
      string to write to log messages

    level : str = INFO
      logging levelname

    **extra
      locals to pass into the message format
    """
    extra["df"] = df
    lvl = logging.getLevelName(level)
    logging.log(lvl, msg.format(**extra))
    return df


from rest_framework import serializers
from cride.circles.models.memberships import Membership
from cride.circles.models.invitations import Invitation
from cride.users.serializers import UserModelSerializer
from django.utils import timezone

class MembershipModelSerializer(serializers.ModelSerializer):

    joined_at = serializers.DateTimeField(source='created', read_only=True)
    user = UserModelSerializer(read_only=True)
    invited_by = serializers.StringRelatedField()

    class Meta:

        model = Membership
        fields = (
            'user',
            'is_admin', 'is_active',
            'used_invitation', 'remaining_invitation',
            'invited_by', 'rides_taken', 'rides_offered',
            'joined_at'
        )
        read_only_fields = (
            'user',
            'used_invitation',
            'invited_by',
            'rides_taken', 'rides_offered'
        )

class AddMemberSerializer(serializers.Serializer):

    invitation_code = serializers.CharField(min_length=8)
    user = serializers.HiddenField(default=serializers.CurrentUserDefault())

    def validate_user(self, data):
        circle = self.context['circle']
        user = data
        query = Membership.objects.filter(circle=circle, user=user)
        if query.exists():
            raise serializers.ValidationError('User is already a register member.')

    def validate_invitation_code(self, data):
        try:
            invitation = Invitation.objects.get(
                    code=data,
                    circle=self.context['circle'],
                    used=False
                )
        except Invitation.DoesNotExists:
            raise serializers.ValidationError('Invalid invitation code')

    def validate(self, data):
        circle = self.context['circle']
        if circle.is_limited and circle.members.count() >= circle.members_limit:
            raise serializers.ValidationError('Circle is in its invitations limit')
        return data

    def create(self, data):
        circle = self.context['circle']
        invitation = self.context['invitation']
        user = data['user']
        now = timezone.now()

        member = Membership.objects.create(
            user=user,
            profile=user.profile,
            cicle=circle,
            invited_by=invitation.issued_by
        )

        invitation.used_by = user
        invitation.used = True
        invitation.used_at = now
        invitation.save()

        issuer = Membership.objects.get(
            user=invitation.issued_by,
            circle=circle
        )
        issuer.used_invitation +=1
        issuer.remaining_invitation -= 1

        return member


from flask import Flask, render_template
import os
app = Flask(__name__)

PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__))
STATIC_ROOT = os.path.join(PROJECT_ROOT,'static/')

URLS  = ['basic','selections','datajoins','enterexit','marybartlett'
        ,'d3noob','hierarchy', 'treemap','treeclustertransition','radialtree'
        ,'indentedtree','divmap','sidemap']

@app.route('/')
def index():
    return render_template('index.html', urls=URLS)

@app.route('/<url>')
def practice(url):
    if url == 'example.html':
        return render_template('example.html')

    #static
    elif url.endswith('.csv') or url.endswith('.json'):
        return app.send_static_file(url)

    return render_template('practice.html', url=url)

if __name__ == '__main__':
    app.debug = True;
    port = int(os.environ.get('PORT', 5000))
    app.run(host='0.0.0.0', port=port)


import math
import copy

import torch
import torch.nn.functional as F
from torch import nn

from util import segment_ops
from util.misc import (NestedTensor, nested_tensor_from_tensor_list,
                       accuracy, get_world_size,
                       is_dist_avail_and_initialized, inverse_sigmoid)
from models.matcher import build_matcher
from models.position_encoding import build_position_encoding
from .custom_loss import sigmoid_focal_loss
from .transformer import build_deformable_transformer
from opts import cfg

if not cfg.disable_cuda:
    from models.ops.roi_align import ROIAlign

def _get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

def get_norm(norm_type, dim, num_groups=None):
    if norm_type == 'gn':
        assert num_groups is not None, 'num_groups must be specified'
        return nn.GroupNorm(num_groups, dim)
    elif norm_type == 'bn':
        return nn.BatchNorm1d(dim)
    else:
        raise NotImplementedError

class TadTR(nn.Module):
    """ This is the TadTR module that performs temporal action detection """

    def __init__(self, position_embedding, transformer, num_classes, num_queries, aux_loss=True, with_segment_refine=True, with_act_reg=True):
        """ Initializes the model.
        Parameters:
            backbone: torch module of the backbone to be used. See backbone.py
            transformer: torch module of the transformer architecture. See deformable_transformer.py
            num_classes: number of action classes
            num_queries: number of action queries, ie detection slot. This is the maximal number of actions
                         TadTR can detect in a single video.
            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
            with_segment_refine: iterative segment refinement
        """
        super().__init__()
        self.num_queries = num_queries
        self.transformer = transformer
        hidden_dim = transformer.d_model
        self.class_embed = nn.Linear(hidden_dim, num_classes)
        self.segment_embed = MLP(hidden_dim, hidden_dim, 2, 3)
        self.query_embed = nn.Embedding(num_queries, hidden_dim*2)

        self.input_proj = nn.ModuleList([
            nn.Sequential(
                nn.Conv1d(2048, hidden_dim, kernel_size=1),
                nn.GroupNorm(32, hidden_dim),
            )])
        # self.backbone = backbone
        self.position_embedding = position_embedding
        self.aux_loss = aux_loss
        self.with_segment_refine = with_segment_refine
        self.with_act_reg = with_act_reg

        prior_prob = 0.01
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        self.class_embed.bias.data = torch.ones(num_classes) * bias_value
        nn.init.constant_(self.segment_embed.layers[-1].weight.data, 0)
        nn.init.constant_(self.segment_embed.layers[-1].bias.data, 0)
        for proj in self.input_proj:
            nn.init.xavier_uniform_(proj[0].weight, gain=1)
            nn.init.constant_(proj[0].bias, 0)

        num_pred = transformer.decoder.num_layers
        if with_segment_refine:
            self.class_embed = _get_clones(self.class_embed, num_pred)
            self.segment_embed = _get_clones(self.segment_embed, num_pred)
            nn.init.constant_(
                self.segment_embed[0].layers[-1].bias.data[1:], -2.0)
            # hack implementation for segment refinement
            self.transformer.decoder.segment_embed = self.segment_embed
        else:
            nn.init.constant_(
                self.segment_embed.layers[-1].bias.data[1:], -2.0)
            self.class_embed = nn.ModuleList(
                [self.class_embed for _ in range(num_pred)])
            self.segment_embed = nn.ModuleList(
                [self.segment_embed for _ in range(num_pred)])
            self.transformer.decoder.segment_embed = None

        if with_act_reg:
            # RoIAlign params
            self.roi_size = 16
            self.roi_scale = 0
            self.roi_extractor = ROIAlign(self.roi_size, self.roi_scale)
            self.actionness_pred = nn.Sequential(
                nn.Linear(self.roi_size * hidden_dim, hidden_dim),
                nn.ReLU(inplace=True),
                nn.Linear(hidden_dim, hidden_dim),
                nn.ReLU(inplace=True),
                nn.Linear(hidden_dim, 1),
                nn.Sigmoid()
            )

    def _to_roi_align_format(self, rois, T, scale_factor=1):
        '''Convert RoIs to RoIAlign format.
        Params:
            RoIs: normalized segments coordinates, shape (batch_size, num_segments, 4)
            T: length of the video feature sequence
        '''
        # transform to absolute axis
        B, N = rois.shape[:2]
        rois_center = rois[:, :, 0:1]
        rois_size = rois[:, :, 1:2] * scale_factor
        rois_abs = torch.cat(
            (rois_center - rois_size/2, rois_center + rois_size/2), dim=2) * T
        # expand the RoIs
        rois_abs = torch.clamp(rois_abs, min=0, max=T)  # (N, T, 2)
        # add batch index
        batch_ind = torch.arange(0, B).view((B, 1, 1)).to(rois_abs.device)
        batch_ind = batch_ind.repeat(1, N, 1)
        rois_abs = torch.cat((batch_ind, rois_abs), dim=2)
        # NOTE: stop gradient here to stablize training
        return rois_abs.view((B*N, 3)).detach()

    def forward(self, samples):
        """ The forward expects a NestedTensor, which consists of:
               - samples.tensors: batched images, of shape [batch_size x 3 x H x W]
               - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels
            or a tuple of tensors and mask

            It returns a dict with the following elements:
               - "pred_logits": the classification logits (including no-action) for all queries.
                                Shape= [batch_size x num_queries x (num_classes + 1)]
               - "pred_segments": The normalized segments coordinates for all queries, represented as
                               (center, width). These values are normalized in [0, 1],
                               relative to the size of each individual image (disregarding possible padding).
                               See PostProcess for information on how to retrieve the unnormalized segment.
               - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of
                                dictionnaries containing the two above keys for each decoder layer.
        """
        if not isinstance(samples, NestedTensor):
            if isinstance(samples, (list, tuple)):
                samples = NestedTensor(*samples)
            else:
                samples = nested_tensor_from_tensor_list(samples)  # (n, c, t)

        pos = [self.position_embedding(samples)]
        src, mask = samples.tensors, samples.mask
        srcs = [self.input_proj[0](src)]
        masks = [mask]

        query_embeds = self.query_embed.weight
        hs, init_reference, inter_references, memory = self.transformer(
            srcs, masks, pos, query_embeds)

        outputs_classes = []
        outputs_coords = []
        # gather outputs from each decoder layer
        for lvl in range(hs.shape[0]):
            if lvl == 0:
                reference = init_reference
            else:
                reference = inter_references[lvl - 1]

            reference = inverse_sigmoid(reference)
            outputs_class = self.class_embed[lvl](hs[lvl])
            tmp = self.segment_embed[lvl](hs[lvl])
            # the l-th layer (l >= 2)
            if reference.shape[-1] == 2:
                tmp += reference
            # the first layer
            else:
                assert reference.shape[-1] == 1
                tmp[..., 0] += reference[..., 0]
            outputs_coord = tmp.sigmoid()
            outputs_classes.append(outputs_class)
            outputs_coords.append(outputs_coord)
        outputs_class = torch.stack(outputs_classes)
        outputs_coord = torch.stack(outputs_coords)

        if not self.with_act_reg:
            out = {'pred_logits': outputs_class[-1],
                   'pred_segments': outputs_coord[-1]}
        else:
            # perform RoIAlign
            B, N = outputs_coord[-1].shape[:2]
            origin_feat = memory

            rois = self._to_roi_align_format(
                outputs_coord[-1], origin_feat.shape[2], scale_factor=1.5)
            roi_features = self.roi_extractor(origin_feat, rois)
            roi_features = roi_features.view((B, N, -1))
            pred_actionness = self.actionness_pred(roi_features)

            last_layer_cls = outputs_class[-1]
            last_layer_reg = outputs_coord[-1]

            out = {'pred_logits': last_layer_cls,
                   'pred_segments': last_layer_reg, 'pred_actionness': pred_actionness}

        if self.aux_loss:
            out['aux_outputs'] = self._set_aux_loss(
                outputs_class, outputs_coord)

        return out

    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord):
        # this is a workaround to make torchscript happy, as torchscript
        # doesn't support dictionary with non-homogeneous values, such
        # as a dict having both a Tensor and a list.
        return [{'pred_logits': a, 'pred_segments': b}
                for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]

class SetCriterion(nn.Module):
    """ This class computes the loss for TadTR.
    The process happens in two steps:
        1) we compute hungarian assignment between ground truth segments and the outputs of the model
        2) we supervise each pair of matched ground-truth / prediction (supervise class and segment)
    """

    def __init__(self, num_classes, matcher, weight_dict, losses, focal_alpha=0.25):
        """ Create the criterion.
        Parameters:
            num_classes: number of action categories, omitting the special no-action category
            matcher: module able to compute a matching between targets and proposals
            weight_dict: dict containing as key the names of the losses and as values their relative weight.
            losses: list of all the losses to be applied. See get_loss for list of available losses.
            focal_alpha: alpha in Focal Loss
        """
        super().__init__()
        self.num_classes = num_classes
        self.matcher = matcher
        self.weight_dict = weight_dict
        self.losses = losses
        self.focal_alpha = focal_alpha

    def loss_labels(self, outputs, targets, indices, num_segments, log=True):
        """Classification loss (NLL)
        targets dicts must contain the key "labels" containing a tensor of dim [nb_target_segments]
        """
        assert 'pred_logits' in outputs
        src_logits = outputs['pred_logits']

        idx = self._get_src_permutation_idx(indices)
        target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)])
        target_classes = torch.full(src_logits.shape[:2], self.num_classes,
                                    dtype=torch.int64, device=src_logits.device)
        target_classes[idx] = target_classes_o

        target_classes_onehot = torch.zeros([src_logits.shape[0], src_logits.shape[1], src_logits.shape[2] + 1],
                                            dtype=src_logits.dtype, layout=src_logits.layout, device=src_logits.device)
        target_classes_onehot.scatter_(2, target_classes.unsqueeze(-1), 1)

        target_classes_onehot = target_classes_onehot[:,:,:-1]
        loss_ce = sigmoid_focal_loss(src_logits, target_classes_onehot, num_segments, alpha=self.focal_alpha, gamma=2) * src_logits.shape[1]  # nq
        losses = {'loss_ce': loss_ce}

        if log:
            # TODO this should probably be a separate loss, not hacked in this one here
            losses['class_error'] = 100 - accuracy(src_logits[idx], target_classes_o)[0]

        return losses

    def loss_segments(self, outputs, targets, indices, num_segments):
        """Compute the losses related to the segmentes, the L1 regression loss and the IoU loss
           targets dicts must contain the key "segments" containing a tensor of dim [nb_target_segments, 2]
           The target segments are expected in format (center, width), normalized by the video length.
        """
        assert 'pred_segments' in outputs
        idx = self._get_src_permutation_idx(indices)
        src_segments = outputs['pred_segments'][idx]
        target_segments = torch.cat([t['segments'][i] for t, (_, i) in zip(targets, indices)], dim=0)

        loss_segment = F.l1_loss(src_segments, target_segments, reduction='none')

        losses = {}
        losses['loss_segments'] = loss_segment.sum() / num_segments

        loss_iou = 1 - torch.diag(segment_ops.segment_iou(
            segment_ops.segment_cw_to_t1t2(src_segments),
            segment_ops.segment_cw_to_t1t2(target_segments)))
        losses['loss_iou'] = loss_iou.sum() / num_segments
        return losses

    def loss_actionness(self, outputs, targets, indices, num_segments):
        """Compute the actionness regression loss
           targets dicts must contain the key "segments" containing a tensor of dim [nb_target_segments, 2]
           The target segments are expected in format (center, width), normalized by the video length.
        """
        assert 'pred_segments' in outputs
        assert 'pred_actionness' in outputs
        src_segments = outputs['pred_segments'].view((-1, 2))
        target_segments = torch.cat([t['segments'] for t in targets], dim=0)

        losses = {}

        iou_mat = segment_ops.segment_iou(
            segment_ops.segment_cw_to_t1t2(src_segments),
            segment_ops.segment_cw_to_t1t2(target_segments))

        gt_iou = iou_mat.max(dim=1)[0]
        pred_actionness = outputs['pred_actionness']
        loss_actionness = F.l1_loss(pred_actionness.view(-1), gt_iou.view(-1).detach())

        losses['loss_actionness'] = loss_actionness
        return losses

    def _get_src_permutation_idx(self, indices):
        # permute predictions following indices
        batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
        src_idx = torch.cat([src for (src, _) in indices])
        return batch_idx, src_idx

    def _get_tgt_permutation_idx(self, indices):
        # permute targets following indices
        batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
        tgt_idx = torch.cat([tgt for (_, tgt) in indices])
        return batch_idx, tgt_idx

    def get_loss(self, loss, outputs, targets, indices, num_segments, **kwargs):
        loss_map = {
            'labels': self.loss_labels,
            'segments': self.loss_segments,
            'actionness': self.loss_actionness,
        }

        assert loss in loss_map, f'do you really want to compute {loss} loss?'
        return loss_map[loss](outputs, targets, indices, num_segments, **kwargs)

    def forward(self, outputs, targets):
        """ This performs the loss computation.
        Parameters:
             outputs: dict of tensors, see the output specification of the model for the format
             targets: list of dicts, such that len(targets) == batch_size.
                      The expected keys in each dict depends on the losses applied, see each loss' doc
        """
        outputs_without_aux = {k: v for k, v in outputs.items() if k != 'aux_outputs'}

        # Retrieve the matching between the outputs of the last layer and the targets
        indices = self.matcher(outputs_without_aux, targets)

        # Compute the average number of target segments accross all nodes, for normalization purposes
        num_segments = sum(len(t["labels"]) for t in targets)
        num_segments = torch.as_tensor([num_segments], dtype=torch.float, device=next(iter(outputs.values())).device)
        if is_dist_avail_and_initialized():
            torch.distributed.all_reduce(num_segments)
        num_segments = torch.clamp(num_segments / get_world_size(), min=1).item()

        # Compute all the requested losses
        losses = {}
        for loss in self.losses:
            kwargs = {}
            losses.update(self.get_loss(loss, outputs, targets, indices, num_segments, **kwargs))

        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
        if 'aux_outputs' in outputs:
            for i, aux_outputs in enumerate(outputs['aux_outputs']):
                indices = self.matcher(aux_outputs, targets)
                for loss in self.losses:
                    # we do not compute actionness loss for aux outputs
                    if 'actionness' in loss:
                        continue

                    kwargs = {}
                    if loss == 'labels':
                        # Logging is enabled only for the last layer
                        kwargs['log'] = False
                    l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_segments, **kwargs)
                    l_dict = {k + f'_{i}': v for k, v in l_dict.items()}
                    losses.update(l_dict)

        self.indices = indices
        return losses

class PostProcess(nn.Module):
    """ This module converts the model's output into the format expected by the TADEvaluator"""

    @torch.no_grad()
    def forward(self, outputs, target_sizes, fuse_score=True):
        """ Perform the computation
        Parameters:
            outputs: raw outputs of the model
            target_sizes: tensor of dimension [batch_size] containing the duration of each video of the batch
        """
        out_logits, out_segments = outputs['pred_logits'], outputs['pred_segments']

        assert len(out_logits) == len(target_sizes)
        # assert target_sizes.shape[1] == 1

        prob = out_logits.sigmoid()   # [bs, nq, C]
        if fuse_score:
            prob *= outputs['pred_actionness']

        segments = segment_ops.segment_cw_to_t1t2(out_segments)   # bs, nq, 2

        if cfg.postproc_rank == 1:     # default
            # sort across different instances, pick top 100 at most
            topk_values, topk_indexes = torch.topk(prob.view(
                out_logits.shape[0], -1), min(cfg.postproc_ins_topk, prob.shape[1]*prob.shape[2]), dim=1)
            scores = topk_values
            topk_segments = topk_indexes // out_logits.shape[2]
            labels = topk_indexes % out_logits.shape[2]

            # bs, nq, 2; bs, num, 2
            segments = torch.gather(
                segments, 1, topk_segments.unsqueeze(-1).repeat(1, 1, 2))
            query_ids = topk_segments
        else:
            # pick topk classes for each query
            # pdb.set_trace()
            scores, labels = torch.topk(prob, cfg.postproc_cls_topk, dim=-1)
            scores, labels = scores.flatten(1), labels.flatten(1)
            # (bs, nq, 1, 2)
            segments = segments[:, [
                i//cfg.postproc_cls_topk for i in range(cfg.postproc_cls_topk*segments.shape[1])], :]
            query_ids = (torch.arange(0, cfg.postproc_cls_topk*segments.shape[1], 1, dtype=labels.dtype,
                         device=labels.device) // cfg.postproc_cls_topk)[None, :].repeat(labels.shape[0], 1)

        # from normalized [0, 1] to absolute [0, length] coordinates
        vid_length = target_sizes
        scale_fct = torch.stack([vid_length, vid_length], dim=1)
        segments = segments * scale_fct[:, None, :]

        results = [{'scores': s, 'labels': l, 'segments': b, 'query_ids': q}
                   for s, l, b, q in zip(scores, labels, segments, query_ids)]

        return results

class MLP(nn.Module):
    """ Very simple multi-layer perceptron (also called FFN)"""

    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super().__init__()
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        self.layers = nn.ModuleList(nn.Linear(n, k)
                                    for n, k in zip([input_dim] + h, h + [output_dim]))

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        return x

def build(args):
    if args.binary:
        num_classes = 1
    else:
        if args.dataset_name == 'thumos14':
            num_classes = 20
        elif args.dataset_name == 'muses':
            num_classes = 25
        elif args.dataset_name in ['activitynet', 'hacs']:
            num_classes = 200
        else:
            raise ValueError('unknown dataset {}'.format(args.dataset_name))

    pos_embed = build_position_encoding(args)
    transformer = build_deformable_transformer(args)

    model = TadTR(
        pos_embed,
        transformer,
        num_classes=num_classes,
        num_queries=args.num_queries,
        aux_loss=args.aux_loss,
        with_segment_refine=args.seg_refine,
        with_act_reg=args.act_reg
    )

    matcher = build_matcher(args)
    losses = ['labels', 'segments']

    weight_dict = {
        'loss_ce': args.cls_loss_coef,
        'loss_segments': args.seg_loss_coef,
        'loss_iou': args.iou_loss_coef}

    if args.act_reg:
        weight_dict['loss_actionness'] = args.act_loss_coef
        losses.append('actionness')

    if args.aux_loss:
        aux_weight_dict = {}
        for i in range(args.dec_layers - 1):
            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()})
        aux_weight_dict.update({k + f'_enc': v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)

    criterion = SetCriterion(num_classes, matcher,
        weight_dict, losses, focal_alpha=args.focal_alpha)

    postprocessor = PostProcess()

    return model, criterion, postprocessor


#!/usr/bin/env python3
# -*- coding:utf-8 -*-

# 导入webdriver API对象，可以调用浏览器和操作页面
from selenium import webdriver
# 导入Keys， 可以使用操作键盘、标签和鼠标
from selenium.webdriver.common.keys import Keys

# options = webdriver.FirefoxOptions()
# options.headless = True
# firefox = webdriver.Firefox(executable_path=r'L:\code\scrapy\Selenium\geckodriver.exe', options=options)
# url = 'https://bbs.level-plus.net/'
# firefox.get(url)
driver = webdriver.PhantomJS(executable_path=r"L:\code\scrapy\Selenium\phantomjs.exe")
driver.get('https://bbs.level-plus.net/')
driver.save_screenshot('soul_plus.png')  # 能截取整个网页，比Firefox好

# login before
driver.find_element_by_name('pwuser').send_keys('xxx')
driver.find_element_by_name('pwpwd').send_keys('xxx')
driver.save_screenshot('login_before.png')

# login after
driver.find_element_by_class_name('btn').click()
driver.save_screenshot('login_after.png')


import os
import numpy as np
import cv2
import csv

dir = "../data1/"
for image in os.listdir(dir):
    if image.endswith("jpg") or image.endswith("JPG"):
        if os.path.isfile(dir+image+".csv"):
            with open(dir+image+ ".csv", 'r') as csvfile:
                spamwriter = csv.reader(csvfile, delimiter=' ',
                                        quotechar='|', quoting=csv.QUOTE_MINIMAL)
                img = cv2.imread(dir + image)
                no = 0
                for row in spamwriter:
                    no+=1
                    print (row)
                    img = cv2.circle(img, (int(float(row[0])), int(float(row[1]))), 2,(255-no*60,no*60,0),90)
                img = cv2.resize(img, (300,300))
                cv2.imshow("a",img)
                cv2.waitKey(0)


from typing import Callable, Generator

import jax
import jax.numpy as jnp
from jaxtyping import Array

def make_quadrature_points_batches(
  grids: Array,
  weights: Array,
  batch_size: int,
  epochs: int,
  num_copies: int,
  seed: int = 137
) -> Generator:
  """
  Args:
    grids: an array of 3d electron coordinates
    weights: an array of weights for each point in the grid
  Return:
   a generator of batches
  """

  @jax.jit
  def shuffle_fn(key):
    return shuffle_quadrature_points(grids, weights, batch_size, key)

  rng = jax.random.PRNGKey(seed)
  key, rng = jax.random.split(rng)

  for _ in range(epochs):

    batches = []
    for _ in range(num_copies):
      key, rng = jax.random.split(rng)
      batches.append(shuffle_fn(key))

    for batch in zip(*batches):
      if num_copies == 1:
        yield batch[0]
      else:
        yield batch

def shuffle_quadrature_points(
  grids: Array,
  weights: Array,
  batch_size: int,
  key: Array,
):
  """Shuffle all

  Args:
    grids: an array of 3d electron coordinates
    weights: an array of weights for each point in the grid
    batch_size: batch size
    key: jax prng key
  Returns:
    shuffled quadrature points
  """
  # combine grids and weights
  weights = jnp.expand_dims(weights, 1)
  gw = jnp.concatenate((grids, weights), axis=1)

  # shuffle quadrature points
  num_total_points = grids.shape[0]
  # key = jax.random.PRNGKey(seed)
  gw = jax.random.permutation(key, gw)
  shuffled_grids = gw[:, :3]
  shuffled_weights = jnp.squeeze(gw[:, 3])

  # make batches
  batch_size = min(num_total_points, batch_size)
  num_batches = num_total_points // batch_size
  batched_g = jnp.split(shuffled_grids[:num_batches * batch_size], num_batches)
  rescaled_w = shuffled_weights * num_total_points / batch_size
  batched_w = jnp.split(rescaled_w[:num_batches * batch_size], num_batches)

  return list(zip(batched_g, batched_w))

def quadrature_integral(integrand: Callable, *coords_and_weights) -> Array:
  """Integrate a multivariate function with quadrature.

  - The integral can be computed with a randomly sampled batch of
  quadrature points, as described in the D4FT paper.

  Args:
    integrand: a multivariable function, with kwargs keepdims
    keepdims: whether to output the outer product matrix as explained above
    coords_and_weights: the quadrature points that the function will
      be evaluated and the weights of these points.

  Returns:
    A scalar value as the result of the integral.
  """
  # break down the coordinates and the weights
  coords = [coord for coord, _ in coords_and_weights]
  weights = [weight for _, weight in coords_and_weights]
  # vmap the integrand
  num_dims = len(coords_and_weights)
  fn = integrand
  for i in range(num_dims):
    in_axes = (None,) * (num_dims - i - 1) + (0,) + (None,) * (i)
    fn = jax.vmap(fn, in_axes=in_axes)
  out = fn(*coords)

  if len(weights) == 1:
    return jnp.sum((out.T * weights[0]).T, axis=0)

  # multi-dimensional quadrature
  for weight in reversed(weights):
    out = jnp.dot(out, weight)

  return out

def get_integrand(f: Callable, g: Callable, keepdims: bool = False) -> Callable:
  """
  If the integrand is the inner product of two vectors of basis functions
  F(r)=(f_1(r), ..., f_N(r)) and G(r)=(g_1(r), ..., g_N(r)), we can keep
  the basis information by computing the outer product matrix <f_i|g_j>.

  - Since the integral is computed with quadrature, i.e.

  .. math::
  <f_i|g_j>= \\sum_k w_k \\dot (f_i(r_k)g_j(r_k))

  where w_k is the quadrature weight and r_k are the grid points,
  the outer prodcut matrix can be computed as the outer product of the
  vector F(r) and G(r):

  .. math::
  (<f_i|g_j>)_{ij} = \\sum_k w_k \\dot (F(r_k)G(r_k)^T)
  """

  def integrand(r: Array) -> Array:
    F, G = f(r), g(r)
    if keepdims:
      if len(F.shape) == 1:
        return jnp.outer(F, G)
      else:  # map over spin axis
        return jax.vmap(jnp.outer)(F, G)
    return jnp.sum(F * G)

  return integrand

def wave2density(orbitals: Callable, polarized=False) -> Callable:
  """
  Transform the wave function into density function.
  Args:
    mo: a [3] -> [2, N] function, where N is the number of molecular orbitals.
      mo only takes one argment, which is the coordinate.
    polarized: if True will return a 1D array with two elements indicating
      the density of each spin
  Return:
    density function: [3] -> float or 1D array.
  """
  if polarized:
    return lambda r: jnp.sum((orbitals(r))**2, axis=1)
  else:
    return lambda r: jnp.sum((orbitals(r))**2)


'''
Given a list of words, return the shortest unique prefix of each word. For example, given the list:

    dog
    cat
    apple
    apricot
    fish

Return the list:

    d
    c
    app
    apr
    f
'''

from typing import List

class TrieNode:
    def __init__(self, letter="", parent=None):
        self.letter = letter
        self.word = None
        self.next_letters = {}
        self.parent = parent

    def __repr__(self):
        return self.letter

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def addWord(self, word: str):
        def rAddWord(trie_node: TrieNode, word: str, letter_index: int):
            if letter_index == len(word):
                trie_node.word = word
                return

            if word[letter_index] not in trie_node.next_letters:
                trie_node.next_letters[word[letter_index]] = TrieNode(word[letter_index], trie_node)

            rAddWord(trie_node.next_letters[word[letter_index]], word, letter_index+1)

            return

        rAddWord(self.root, word, 0)

        return

    def getUniquePrefix(self, word: str):
        def getTargetWordNode(cur_node: TrieNode, word, letter_index: int):
            if cur_node.word == word:
                return cur_node

            if letter_index >= len(word) or word[letter_index] not in cur_node.next_letters:
                raise ValueError(f"Word {word} Not Found")

            return getTargetWordNode(cur_node.next_letters[word[letter_index]], word, letter_index+1)

        if word == "":
            return ""

        node_ptr = getTargetWordNode(self.root, word, 0)

        # If there are any children, the unique prefix is the full word.
        if node_ptr.next_letters:
            return word

        # Backtrack until the parent node has no parent, has more than 1 child, or contains a word.
        prefix_length = len(word)
        while node_ptr != self.root:
            if (not node_ptr.parent.parent or
                len(node_ptr.parent.next_letters) > 1 or
                node_ptr.parent.word):
                return word[:prefix_length]

            prefix_length -= 1
            node_ptr = node_ptr.parent

        return ""

def getAllUniquePrefixes(word_list: List[str]):
    # Build Trie
    word_trie = Trie()
    for word in word_list:
        word_trie.addWord(word)

    unique_prefixes = []

    for word in word_list:
        unique_prefixes.append(word_trie.getUniquePrefix(word))

    return unique_prefixes

print(getAllUniquePrefixes(["dog", "cat", "apple", "applesauce", "apricot", "fish"]))


from __future__ import annotations

from typing import Iterator, List, Optional, Sequence, Tuple, Union

from bip_utils.bip.bip32.bip32_ex import Bip32PathError
from bip_utils.bip.bip32.bip32_key_data import Bip32KeyIndex

class Bip32PathConst:
    """Class container for BIP32 path constants."""

    # Hardened characters
    HARDENED_CHARS: Tuple[str, str, str] = ("'", "h", "p")
    # Master character
    MASTER_CHAR: str = "m"

class Bip32Path:
    """
    BIP32 path class.
    It represents a BIP-0032 path.
    """

    m_elems: List[Bip32KeyIndex]
    m_is_absolute: bool

    def __init__(self,
                 elems: Optional[Sequence[Union[int, Bip32KeyIndex]]] = None,
                 is_absolute: bool = True) -> None:
        """
        Construct class.

        Args:
            elems (list, optional)      : Path elements (default: empty)
            is_absolute (bool, optional): True if path is an absolute one, false otherwise (default: True)
        """
        try:
            self.m_elems = ([]
                            if elems is None
                            else [Bip32KeyIndex(elem) if isinstance(elem, int) else elem for elem in elems])
        except ValueError as ex:
            raise Bip32PathError("The path contains some invalid key indexes") from ex

        self.m_is_absolute = is_absolute

    def AddElem(self,
                elem: Union[int, Bip32KeyIndex]) -> Bip32Path:
        """
        Return a new path object with the specified element added.

        Args:
            elem (str or Bip32KeyIndex): Path element

        Returns:
            Bip32Path object: Bip32Path object

        Raises:
            Bip32PathError: If the path element is not valid
        """
        if isinstance(elem, int):
            elem = Bip32KeyIndex(elem)
        return Bip32Path(self.m_elems + [elem], self.m_is_absolute)

    def IsAbsolute(self) -> bool:
        """
        Get if absolute path.

        Returns:
            bool: True if absolute path, false otherwise
        """
        return self.m_is_absolute

    def Length(self) -> int:
        """
        Get the number of elements of the path.

        Returns:
            int: Number of elements
        """
        return len(self.m_elems)

    def ToList(self) -> List[int]:
        """
        Get the path as a list of integers.

        Returns:
            list[int]: Path as a list of integers
        """
        return [int(elem) for elem in self.m_elems]

    def ToStr(self) -> str:
        """
        Get the path as a string.

        Returns:
            str: Path as a string
        """
        path_str = "" if not self.m_is_absolute else f"{Bip32PathConst.MASTER_CHAR}/"
        for elem in self.m_elems:
            if not elem.IsHardened():
                path_str += f"{str(elem.ToInt())}/"
            else:
                path_str += f"{str(Bip32KeyIndex.UnhardenIndex(elem.ToInt()))}'/"

        return path_str[:-1]

    def __str__(self) -> str:
        """
        Get the path as a string.

        Returns:
            str: Path as a string
        """
        return self.ToStr()

    def __getitem__(self,
                    idx: int) -> Bip32KeyIndex:
        """
        Get the specified element index.

        Args:
            idx (int): Element index

        Returns:
            Bip32KeyIndex object: Bip32KeyIndex object
        """
        return self.m_elems[idx]

    def __iter__(self) -> Iterator[Bip32KeyIndex]:
        """
        Get the iterator to the current element.

        Returns:
            Iterator object: Iterator to the current element
        """
        yield from self.m_elems

class Bip32PathParser:
    """
    BIP32 path parser class.
    It parses a BIP-0032 path and returns a Bip32Path object.
    """

    @staticmethod
    def Parse(path: str) -> Bip32Path:
        """
        Parse a path and return a Bip32Path object.

        Args:
            path (str): Path

        Returns:
            Bip32Path object: Bip32Path object

        Raises:
            Bip32PathError: If the path is not valid
        """

        # Remove trailing "/" if any
        if path.endswith("/"):
            path = path[:-1]

        # Parse elements
        return Bip32PathParser.__ParseElements(
            list(filter(None, path.split("/")))
        )

    @staticmethod
    def __ParseElements(path_elems: List[str]) -> Bip32Path:
        """
        Parse path elements and return a Bip32Path object.

        Args:
            path_elems (list[str]): Path elements

        Returns:
            Bip32Path object: Bip32Path object

        Raises:
            Bip32PathError: If the path is not valid
        """

        # Remove the initial "m" character if any
        if len(path_elems) > 0 and path_elems[0] == Bip32PathConst.MASTER_CHAR:
            path_elems = path_elems[1:]
            is_absolute = True
        else:
            is_absolute = False

        # Parse elements
        parsed_elems = list(map(Bip32PathParser.__ParseElem, path_elems))
        return Bip32Path(parsed_elems, is_absolute)

    @staticmethod
    def __ParseElem(path_elem: str) -> int:
        """
        Parse path element and get the correspondent index.

        Args:
            path_elem (str): Path element

        Returns:
            int: Index of the element, None if the element is not a valid index

        Raises:
            Bip32PathError: If the path is not valid
        """

        # Strip spaces
        path_elem = path_elem.strip()

        # Get if hardened
        is_hardened = path_elem.endswith(Bip32PathConst.HARDENED_CHARS)

        # If hardened, remove the last character from the string
        if is_hardened:
            path_elem = path_elem[:-1]

        # The remaining string shall be numeric
        if not path_elem.isnumeric():
            raise Bip32PathError(f"Invalid path element ({path_elem})")

        return int(path_elem) if not is_hardened else Bip32KeyIndex.HardenIndex(int(path_elem))


# Converting image to NDVI colored scale

from PIL import ImageDraw, ImageFont, Image  # Pillow library

data = open('index_convert.csv', 'w')
data.write("Name: ;Longtitude: ; Latitude: ;Index name: ;Maximum: ;Average: ;Minimum: \n")
data.close()

def index_convert(image, lon, lat):
    dr = ''
    path = image.split("/")
    for p in path[:-1]:
        dr += p + '/'
    dr += '/indicies/'
    date = '.'.join(path[-1].split('.')[:-1])
    img = Image.open(image)
    # I get height and width of the image
    w, h = img.size

    # function below is useful, because it can be used to many various indexes defined after
    def operation(code, name, contrast):
        # ^^ kolejno: nazwa funkcji indeksu, nazwa do zapisania w pliku, funkcja kontrastu
        # name of index func, name to save in file, contrast func

        # for every single index, that's what happen:
        nonlocal img, date
        global data
        print('index: ' + name)
        index_col = img.copy().convert('HSV')  # creating hsv picture (h means color, i'll use it later)

        img_rgb = img.convert('RGB')
        index_bw = img_rgb.copy()  # creating copy for black-white index

        px_index_bw = index_bw.load()  # I get pixel information of b-w picture
        px_index_col = index_col.load()  # and hsv one
        index_sum = []  # list of every pixel's index

        # checking every single pixel

        for X in range(0, w):  # width
            for Y in range(0, h):  # height
                pixel_rgb = img_rgb.getpixel((X, Y))  # Get pixel's RGB values

                r, g, b = pixel_rgb

                # ~via Kuba Frączek
                brightness = sum(pixel_rgb) / 3  # 0 is dark (black) and 255 is bright (white)

                if 35 < brightness < 250:  # white and dark spots off
                    index = code(r, g, b)
                    index_sum.append(index)  # add index value to the list t.b.c.

                    index = contrast(index)  # a nice, suitable contrast

                    px_index_bw[X, Y] = (index, index, index)
                    # one picture in GRAYSCALE (r +b +g equals gray (or white (or black)))
                    px_index_col[X, Y] = (index, 200, 200)
                    # and one picture in HSV, in which H is index, so index value is one color of the spectrum

        # ####### scale ######

        s = int((h - 100) / 200)  # scale of the scale (I check how many times it could fit)

        if s > 3:
            s -= 3  # however its still to big, unless it would be too small to exist
        end = s * 200 + 100
        beg = 100

        # if average index value is too high or too low scale must adapt
        if sum(index_sum) / len(index_sum) > 0.4:
            end = h - 10
            if s > 1:
                s -= 1
        if sum(index_sum) / len(index_sum) < -0.4:
            beg = 10
            if s > 2:
                s -= 2
        try:
            font = ImageFont.truetype("DejaVuSansMono.ttf", 2 + s * 10)  # for linux
        except OSError:
            font = ImageFont.truetype("arial.ttf", 2 + s * 10)  # for windows

        for x in range(w - 10 - s * 10, w - 10):
            n = -0.2
            for y in range(beg, end):
                nv = contrast(n)  # the same contrast as at the picture
                px_index_bw[x, y] = (nv, nv, nv)
                px_index_col[x, y] = (nv, 200, 200)
                n += 0.003 / s

        ss = 32 * s  # odstępy między numerami
        hh = beg - 2  # początek skali

        # scale text on the hsv picture
        zo = ImageDraw.Draw(index_col)
        zo.text((w - 13 - s * 30, hh),          "-.2",  (0, 0, 200), font=font)
        zo.text((w - 13 - s * 30, hh + 1 * ss), "-.1",  (0, 0, 200), font=font)
        zo.text((w - 3 - s * 30, hh + 2 * ss),  "0",    (0, 0, 200), font=font)
        zo.text((w - 7 - s * 30, hh + 3 * ss),  ".1",   (0, 0, 200), font=font)
        zo.text((w - 7 - s * 30, hh + 4 * ss),  ".2",   (0, 0, 200), font=font)
        zo.text((w - 7 - s * 30, hh + 5 * ss),  ".3",   (0, 0, 200), font=font)
        zo.text((w - 7 - s * 30, hh + 6 * ss),  ".4",   (0, 0, 200), font=font)

        # scale on the grayscale picture
        za = ImageDraw.Draw(index_bw)
        za.text((w - 13 - s * 30, hh),          "-.2",  (255, 255, 255), font=font)
        za.text((w - 13 - s * 30, hh + 1 * ss), "-.1",  (255, 255, 255), font=font)
        za.text((w - 3 - s * 30, hh + 2 * ss),  "0",    (255, 255, 255), font=font)
        za.text((w - 7 - s * 30, hh + 3 * ss),  ".1",   (255, 255, 255), font=font)
        za.text((w - 7 - s * 30, hh + 4 * ss),  ".2",   (255, 255, 255), font=font)
        za.text((w - 7 - s * 30, hh + 5 * ss),  ".3",   (255, 255, 255), font=font)
        za.text((w - 7 - s * 30, hh + 6 * ss),  ".4",   (255, 255, 255), font=font)

        '''
        what scale does:
        - fits in any picture above 300 px height
        - change its value according to contrast on the picture
        - if the average index value is too high or too low, the scale can do things to fit
        - ...
        '''

        try:
            # works on raspberry pi:
            index_bw.save(dr + date + "_bw_" + name + ".jpg")  # saves picture in grayscale
            index_col = index_col.convert("RGB")
            index_col.save(dr + date + "_col_" + name + ".jpg")  # picture in color scale
        except (FileNotFoundError, OSError):
            index_bw.save(date + "_bw_" + name + ".jpg")  # saves picture in grayscale
            index_col = index_col.convert("RGB")
            index_col.save(date + "_col_" + name + ".jpg")  # picture in color scale

        # print('worked')

        # max value, average, and min
        data = open('index_convert.csv', 'a')
        data.write(str(date) + ';' +
                   str(lon) + ';' +
                   str(lat) + ';' +
                   str(name) + ';' +
                   str(min(index_sum)) + ';' +
                   str(sum(index_sum) / len(index_sum)) + ';' +
                   str(max(index_sum)) + ';' + '\n')

        data.close()

        dane = [name, (min(index_sum), sum(index_sum) / len(index_sum), max(index_sum))]
        return dane
    out = []

    class Code:
        # I define indices functions (mostly) from:
        # https://www.int-arch-photogramm-remote-sens-spatial-inf-sci.net/XLII-3/1215/2018/isprs-archives-XLII-3-1215-2018.pdf
        # plus once I read:
        # because of the way the dyes are coupled to these layers, reproduce:
        # infrared as red, red as green, and green as blue.
        # do zrobienia: sprawdzić wszystkie indeksy w odsłonach normalnych i po powyższej zmianie

        @staticmethod
        def rgbvi(r, g, b):
            if (r - g) / (r + g) < -0.2:
                ix = -2
            else:
                ix = (g * g - r * b) / (g * g + r * b + 0.1)
            return ix

        @staticmethod
        def ndvi(r, g, b):
            ix = (r - b) / (r + b + 0.1)
            return ix

        @staticmethod
        def ndwi(r, g, b):
            ix = (r - g) / (r + g + 0.1)
            return ix

        @staticmethod
        def gli(r, g, b):
            if (r - g) / (r + g + 0.1) < -0.2:
                ix = -2
            else:
                ix = (2 * g - r - b) / (2 * g + r + b + 0.1)
            return ix

        @staticmethod
        def vari(r, g, b):
            if (r - g) / (r + g) < -0.2:
                ix = 2
            else:
                ix = (g - r) / (g + r - b + 0.1)
            return - ix

        @staticmethod
        def rgi(r, g, b):
            ix = r / (g + 0.1)
            return ix

        @staticmethod
        def ergbve(r, g, b):
            ix = 3.14159 * (g * g - r * b) / (g * g + r * b + 0.1)
            return ix

    # contrast function
    # every value below 0 ends up as 0, water starts at about -0.1, so I add 0.15
    # then I multiply it by 500 to make a good contrast
    # it cannot be float, because hsv and rgb can't read that ;c
    # different indexes need different contrast
    # after all I changed it randomly to make it look good

    def con(add, multi):
        def con_new(v):
            return int((v + add) * multi)
        return con_new

    # IMPORTANT! comment if needn't:
    '''
    operation(Code.vari, 'vari', con(0.15, 500))
    operation(Code.ergbve, 'ergbve', con(0.2, 300))
    operation(Code.gli, 'gli', con(0.01, 2000))
    operation(Code.rgbvi, 'rgbvi', con(0.2, 450))

    '''
    out.append(operation(Code.ndwi, 'ndwi', con(0.2, 550)))
    out.append(operation(Code.ndvi, 'ndvi', con(0.35, 300)))
    out.append(operation(Code.rgi, 'rgi', con(-0.6, 200)))

    return out