#PIPELINE GALO VERSION JON

#!/usr/bin/env python

#   Standard library imports
import os
import sys
import re
import argparse
import subprocess
import pandas as pd
import numpy as np

def get_arguments():

    parser = argparse.ArgumentParser(prog = 'snp_covid19.py', description = 'Pipeline to do variant calling with Sars-CoV-2 samples')

    input_group = parser.add_argument_group('Input', 'Input parameters')

    input_group.add_argument('-i', '--input', dest="input_dir", type=str, required=True, help="Required. Input dir with the fastq files")
    input_group.add_argument('-r', '--reference_genome', metavar="reference", type=str, required=True, help='File to map against')
    input_group.add_argument('-p', '--primers_bed', dest="primers", required=True, help='File with the primers')
    input_group.add_argument('-a', '--annotation', dest="annotation", required=True, help='File with the annotation files')
    input_group.add_argument('-R', '--run_id', dest="run_id", required=True, help='Run ID')
    input_group.add_argument('-S', '--sample_list', type=str, required=False, help='Sample names to analyse only in the file supplied')

    output_group = parser.add_argument_group('Output', 'Output parameters')

    output_group.add_argument('-o', '--output', type=str, required=True, help="Required.Output directory to extract all results")

    params_group = parser.add_argument_group('Parameters', 'parameters for diferent stringent conditions')

    params_group.add_argument('-T', '--threads', type=str, dest = "threads", required=False, default=16, help='Threads to use')

    arguments = parser.parse_args()

    return arguments

args = get_arguments()

def extract_read_list(input_dir):
    """
    Search files in a directory sort by name and extract comon name of R1 and R2
    with extract_sample() function
    190615 - Limit only parent folder, not subdirectories
    """
    input_dir = os.path.abspath(input_dir)
    all_fasta = []
    r1_list = []
    r2_list = []
    for root, _, files in os.walk(input_dir):
        if root == input_dir: # This only apply to parent folder, not subdirectories
            for name in files:
                filename = os.path.join(root, name)
                is_fasta = re.match(r'.*\.f(ast)*[aq](\.gz)*',filename)
                if is_fasta:
                    all_fasta.append(filename)
    all_fasta = sorted(all_fasta)
    if len(all_fasta) % 2 == 0:
        for index, fasta_file in enumerate(all_fasta):
            if index % 2 == 0:
                r1_list.append(fasta_file)
            elif index % 1 == 0:
                r2_list.append(fasta_file)
    else:
        print('ERROR: The number of fastq sequence are not paired')

    r1_list = sorted(r1_list)
    r2_list = sorted(r2_list)

    return r1_list, r2_list

def extract_sample(R1_file, R2_file):
    """
    Extract sample from R1, R2 files.
    """
    basename_R1 = os.path.basename(R1_file)
    basename_R2 = os.path.basename(R2_file)

    sample_name_R = os.path.commonprefix([basename_R1, basename_R2])

    long_suffix = re.search('_S.*', sample_name_R)
    short_suffix = re.search('_R.*', sample_name_R)
    bar_suffix = re.search('_$', sample_name_R)

    if long_suffix:
        match = long_suffix.group()
        sample_name = sample_name_R.split(match)[0]
    elif short_suffix:
        match = short_suffix.group()
        sample_name = sample_name_R.split(match)[0]
    elif bar_suffix:
        match = bar_suffix.group()
        sample_name = sample_name_R.rstrip("_")
    else:
        sample_name = sample_name_R

    return sample_name

def file_to_list(file_name):
    list_F = []
    file_name_abs = os.path.abspath(file_name)
    with open(file_name_abs, "r") as f:
        for line in f:
            list_F.append(line.strip())
    return list_F

def check_create_dir(path):
    #exists = os.path.isfile(path)
    #exists = os.path.isdir(path)
    if os.path.exists(path):
        pass
    else:
        os.mkdir(path)

out_consensus_dir = os.path.join(args.output, "Consensus")
out_variants_dir = os.path.join(args.output, "Variants")
out_reports_dir = os.path.join(args.output, "Reports")
out_log_dir = os.path.join(args.output, "Log")

out_bam_dir = os.path.join(args.output, "Bam")
out_trim_dir = os.path.join(args.output, "Trimmed")
out_phyl_dir = os.path.join(args.output, "Phylogeny")

out_line_dir = os.path.join(args.output, "Lineages")
out_compare_dir = os.path.join(args.output, "Compare")

out_emptyvariants_dir = os.path.join(args.output, "Empty_variants")
out_onlysnps_variants = os.path.join(args.output, "Variants_snps")

#Obtain all R1 and R2 from folder
r1, r2 = extract_read_list(args.input_dir)

#Check if there are samples to filter
sample_list_F = []
if args.sample_list == None:
    print("\n" + "No samples to filter")
    for r1_file, r2_file in zip(r1, r2):
        sample = extract_sample(r1_file, r2_file)
        sample_list_F.append(sample)
else:
    print("samples will be filtered")
    sample_list_F = file_to_list(args.sample_list)
print("\n%d samples will be analysed: %s" % (len(sample_list_F), ",".join(sample_list_F)))

for r1_file, r2_file in zip(r1, r2):
    sample = extract_sample(r1_file, r2_file)
    args.sample = sample
    if sample in sample_list_F:

        sample_number = str(sample_list_F.index(sample) + 1)
        sample_total = str(len(sample_list_F))

        args.r1_file = r1_file
        args.r2_file = r2_file

        print("\n" + "STARTING SAMPLE: " + sample + " (" + sample_number + "/" + sample_total + ")")
        # previous trimming

        check_create_dir(out_trim_dir)

        r1 = os.path.abspath(args.r1_file)
        r2 = os.path.abspath(args.r2_file)

        r1_trim = sample + "_R1.trim.fastq.gz"
        out_r1_trim = os.path.join(out_trim_dir, r1_trim)
        r2_trim = sample + "_R2.trim.fastq.gz"
        out_r2_trim = os.path.join(out_trim_dir, r2_trim)

        json = sample + "_fastp_trim.json"
        json_path = os.path.join(out_reports_dir, json)
        html = sample + "_fastp_trim.html"
        html_path = os.path.join(out_reports_dir, html)

        adapter = "/home/laura/COVID/adaptors/IlluminaAdaptors.fasta"

        cmd_trim = ["fastp --in1 " + r1 + " --in2 " + r2 + " --out1 " + out_r1_trim + " --out2 " + out_r2_trim + " --detect_adapter_for_pe --adapter_fasta " + adapter + " --cut_tail --cut_window_size 10 --cut_mean_quality 30 --thread 4 --max_len1 250 --max_len2 250 --json " + json_path + " --html " + html_path]
        subprocess.run(cmd_trim, shell=True)

        # MAPPING

        #bwa mem  -Y -M -R '@RG\tID:.\tSM:.' -t ${THREADS} ${REF_GENOME} ${FQ1} ${FQ2} 2>>${RUN_ID}.log | samtools sort | samtools view -F 4 -b -@ ${THREADS} - > ${SAMPLE}.sort.bam

        check_create_dir(out_bam_dir)
        #os.mkdir(out_log_dir)

        out_map_name = sample + ".sort.bam"
        out_map_file = os.path.join(out_bam_dir, out_map_name)

        reference = os.path.abspath(args.reference_genome)
        r1_trimmed = sample + "_R1.trim.fastq.gz"
        out_r1_trimmed = os.path.join(out_trim_dir, r1_trimmed)
        r2_trimmed = sample + "_R2.trim.fastq.gz"
        out_r2_trimmed = os.path.join(out_trim_dir, r2_trimmed)
        #r1 = os.path.abspath(args.r1_file)
        #r2 = os.path.abspath(args.r2_file)

        run = args.run_id
        run_file = run + ".log"
        run_path = os.path.join(args.output, run_file)

        cmd_map = ['bwa mem -Y -M -R "@RG\\tID:.\\tSM:." -t 4 ' + reference + " " + out_r1_trimmed + " " + out_r2_trimmed + ' | samtools sort | samtools view -F 4 -b -@ 4 -o ' + out_map_file]
        subprocess.run(cmd_map, shell=True)

        # PRME MASKIN RIMMING

        #ivar trim -i ${SAMPLE}.sort.bam -p ${SAMPLE}.trim -m ${MINTRIM_LENGTH} -q ${MINQUALTRIM} -s ${WINDOW} -b ${PRIMERS_BED} &>>${RUN_ID}.log

        in_trim = sample + ".sort.bam"
        input_trim = os.path.join(out_bam_dir, in_trim)

        out_trim = sample + ".trim"
        output_trim = os.path.join(out_bam_dir, out_trim)

        primer = os.path.abspath(args.primers)

        cmd_trim = ["/home/laura/Downloads/ivar-1.2/src/ivar trim -i " + input_trim + " -p " + output_trim + " -m 30 -q 20 -s 4 -b " + primer + " -e"]
        subprocess.run(cmd_trim, shell=True)

        # resort bam, index and depth

        #samtools sort -o ${SAMPLE}.trim.sort.bam ${SAMPLE}.trim.bam 1>&2 2>>${RUN_ID}.log
        #samtools index ${SAMPLE}.trim.sort.bam &>>${RUN_ID}.log
        #samtools depth -aa --reference ${REF_GENOME} ${SAMPLE}.trim.sort.bam > ${SAMPLE}.depth 2>>${RUN_ID}.log

        in_sort = sample + ".trim.bam"
        input_sort = os.path.join(out_bam_dir, in_sort)

        out_sort = sample + ".trim.sort.bam"
        output_sort = os.path.join(out_bam_dir, out_sort)

        cmd_sort = ["samtools sort -o " + output_sort + " " + input_sort]
        subprocess.run(cmd_sort, shell=True)

        in_index = sample + ".trim.sort.bam"
        input_index = os.path.join(out_bam_dir, in_index)

        cmd_index = ["samtools index " + input_index]
        subprocess.run(cmd_index, shell=True)

        check_create_dir(out_reports_dir)

        in_depth = sample + ".trim.sort.bam"
        input_depth = os.path.join(out_bam_dir, in_depth)
        out_depth = sample + ".depth"
        output_depth = os.path.join(out_reports_dir, out_depth)
        reference = os.path.abspath(args.reference_genome)

        cmd_depth = ["samtools depth -aa --reference " + reference + " " + input_depth + " > " + output_depth]
        subprocess.run(cmd_depth, shell=True)

        ### IVAR CNSESUS

        #samtools mpileup -aa -A -d 0 -B -Q 0 ${SAMPLE}.trim.sort.bam | ivar consensus -p ${SAMPLE} -q ${MINQUAL_CONS} -t ${MINFREQ_CONS} -m ${MINDEPTH_CONS} -n ${AMBIGUOUS_CHAR} &>>${RUN_ID}.log

        check_create_dir(out_consensus_dir)

        in_consensus = sample + ".trim.sort.bam"
        input_consensus = os.path.join(out_bam_dir, in_consensus)

        cmd_consensus = ["samtools mpileup -aa -A -d 0 -B -Q 0 " + input_consensus + " | /home/laura/Downloads/ivar-1.2/src/ivar consensus -p " + sample + " -q 20 -t 0.8 -m 30 -n N"]
        subprocess.run(cmd_consensus, shell=True)

        out_consensus = sample + ".fa"
        output_consensus = os.path.join(args.output, out_consensus)
        cmd_mv_one = ["mv " + output_consensus + " " + out_consensus_dir]
        subprocess.run(cmd_mv_one, shell=True)

        out_qual_consensus = sample + ".qual.txt"
        output_qual_consensus = os.path.join(args.output, out_qual_consensus)
        cmd_mv_two = ["mv " + output_qual_consensus + " " + out_consensus_dir]
        subprocess.run(cmd_mv_two, shell=True)

        ### ivar variant calling

        check_create_dir(out_variants_dir)

        #samtools mpileup -aa -A -d 0 -B -Q 0 --reference ${REF_GENOME}  ${SAMPLE}.trim.sort.bam | ivar variants -p ${SAMPLE} -q ${MINQUAL_VAR} -t ${MINFREQ_VAR}

        in_variant = sample + ".trim.sort.bam"
        input_variant = os.path.join(out_bam_dir, in_variant)

        reference = os.path.abspath(args.reference_genome)
        annotation = os.path.abspath(args.annotation)

        cmd_variant = ["samtools mpileup -aa -A -d 0 -B -Q 0 --reference " + reference + " " + input_variant + " | /home/laura/Downloads/ivar-1.2/src/ivar variants -p " + sample + " -q 20 -t 0.05 -m 20 -r " + reference + " -g " + annotation]
        subprocess.run(cmd_variant, shell=True)

        out_var = sample + ".tsv"
        output_var = os.path.join(args.output, out_var)
        cmd_mv_three = ["mv " + output_var + " " + out_variants_dir]
        subprocess.run(cmd_mv_three, shell=True)

        ### calculate depth (.depth is done before)

        depth = sample + ".depth"
        in_depth = os.path.join(out_reports_dir, depth)

        wdepth = sample + ".wdepth"
        out_depth = os.path.join(out_reports_dir, wdepth)

        primer = os.path.abspath(args.primers)

        cmd_depth_one = ["python /home/laura/COVID/coverage_files/COVID-window-coverage.py -i " + in_depth + " -p " + primer + " -w 100 -H -s " + sample + " | sort -nk2,2 > " + out_depth]
        subprocess.run(cmd_depth_one, shell=True)

        first_variant = sample + ".trim.sort.bam"
        first_input_variant = os.path.join(out_bam_dir, first_variant)

        bamstat = sample + ".bamstats"
        bamstat_path = os.path.join(out_reports_dir, bamstat)
        cmd_flagstat = ["samtools flagstat " + first_input_variant + " > " + bamstat_path]
        subprocess.run(cmd_flagstat, shell=True)

        depth_two = sample + ".wdepth"
        out_depth_two = os.path.join(out_reports_dir, depth_two)

        consensus = sample + ".fa"
        consensus_out = os.path.join(out_consensus_dir, consensus)

        report_qc = sample + ".depthQC"
        out_report_qc = os.path.join(out_reports_dir, report_qc)

        global_summary = "QC_summary." + args.run_id + ".txt"
        global_summary_path = os.path.join(out_reports_dir, global_summary)

        cmd_depth_two = ["python /home/laura/COVID/coverage_files/COVID-depth-report.py -d " + out_depth_two + " -c " + consensus_out + " -H -s " + sample + " --min-depth 400 -o " + out_report_qc + " >> " + global_summary_path]
        subprocess.run(cmd_depth_two, shell=True)

        # This file does not exist anymore, but we keep it. Nevertheless it superposes the final one, we must solve it

        depth_two = args.run_id + ".depthQC.txt"
        depth_two_out = os.path.join(args.output, depth_two)

        cmd_one = ["cat " + out_report_qc +  " |  awk 'NR == 1 || $1 !=\"Sample\"' >> " +  depth_two_out]
        subprocess.run(cmd_one, shell=True)

        move = ["mv " + depth_two_out + " " + out_reports_dir]
        subprocess.run(move, shell=True)

## Move empty variants to another folder just to know, because it does not affect to the compare and other steps, as we initially thought

check_create_dir(out_emptyvariants_dir)

for r,d,f in os.walk(out_variants_dir):
    for file in f:
        if file.endswith(".tsv"):
            whole_name = os.path.join(r, file)
            read = pd.read_csv(whole_name, sep="\t")
            if read.shape == (0, 19):
                print("This file is empty and will be removed from the variants analysis")
                cmd_copy = ["cp " + whole_name + " " + out_emptyvariants_dir]
                subprocess.run(cmd_copy, shell=True)
            else:
                pass

## Group coverage stats

def calculate_cov_stats(file_cov):
    df = pd.read_csv(file_cov, sep="\t", names=["#CHROM", "POS", "COV" ])
    unmmaped_pos = len(df.POS[df.COV == 0].tolist())
    pos_0_10 = len(df.POS[(df.COV > 0) & (df.COV <= 10)].tolist())
    pos_10_20 = len(df.POS[(df.COV > 10) & (df.COV <= 20)].tolist())
    pos_high20 = len(df.POS[(df.COV > 20)].tolist())
    pos_high30 = len(df.POS[(df.COV > 30)].tolist())
    pos_high500 = len(df.POS[(df.COV >= 500)].tolist())
    pos_high800 = len(df.POS[(df.COV >= 800)].tolist())
    pos_high1200 = len(df.POS[(df.COV >= 1200)].tolist())
    total_pos = df.shape[0]
    unmmaped_prop = "%.2f" % ((unmmaped_pos/total_pos)*100)
    prop_0_10 = "%.2f" % ((pos_0_10/total_pos)*100)
    prop_10_20 = "%.2f" % ((pos_10_20/total_pos)*100)
    prop_high20 = "%.2f" % ((pos_high20/total_pos)*100)
    prop_high30 = "%.2f" % ((pos_high30/total_pos)*100)
    prop_high500 = "%.2f" % ((pos_high500/total_pos)*100)
    prop_high800 = "%.2f" % ((pos_high800/total_pos)*100)
    prop_high1200 = "%.2f" % ((pos_high1200/total_pos)*100)

    mean_cov = "%.2f" % (df.COV.mean())

    return mean_cov, unmmaped_prop, prop_0_10, prop_10_20, prop_high20, prop_high30, prop_high500, prop_high800, prop_high1200

directory_path = os.path.abspath(out_reports_dir)

output_file_name = args.run_id + ".coverage.tab"
output_file = os.path.join(directory_path,output_file_name)

with open(output_file, "w") as outfile:
        outfile.write("#SAMPLE" + "\t" + "MEAN_COV" + "\t" + "UNMMAPED_PROP" + "\t" + "COV1-10X" + "\t" + "COV10-20X" + "\t" + "COV>20X" + "\t" + "COV>30X" + "\t" + "COV>500X" + "\t" + "COV>800X" + "\t" + "COV>1200" + "\n")
        for root, _, files in os.walk(directory_path):
            for name in files:
                filename = os.path.join(root, name)
                file_name_cov = os.path.basename(filename)
                sample = file_name_cov.split(".")[0]
                if filename.endswith(".depth") and (os.path.getsize(filename) > 0):
                    coverage_stats = calculate_cov_stats(filename)
                    mean_cov = coverage_stats[0]
                    unmmaped_prop = coverage_stats[1]
                    outfile.write(sample + "\t" + ("\t").join(coverage_stats) + "\n")

## Mutation

mut_file = args.run_id + ".mutation.txt"
mutation_file = os.path.join(out_variants_dir, mut_file)

with open(mutation_file, "w") as fout:
    for r, d, f in os.walk(out_variants_dir):
        for file in f:
            if file.endswith(".tsv"):
                name = file.split(".")[0]
                path = os.path.join(r, file)
                prueba = pd.read_csv(path, sep="\t")
                for index, row in prueba.iterrows():
                    if (str(row["REF"]) + str(row["POS"]) + str(row["ALT"])) == "A23403G":
                        fout.write("There's D614G mutation in " + name + "\n")
                    else:
                        pass

## Lineage

# this works, but im not sure with the positions. i've 3 papers reference, but i must contrast.
# also, redirect the results to a txt

# legacy
"""
check_create_dir(out_line_dir)

res = args.run_id + ".lineages.txt"
res_file = os.path.join(out_line_dir, res)

with open(res_file, "w") as fout:
    for r, d, f in os.walk(out_variants_dir):
        for file in f:
            if file.endswith(".tsv"):
                name = file.split(".")[0]
                path = os.path.join(r, file)
                prueba = pd.read_csv(path, sep="\t")
                pos_prueba = list(prueba["POS"])
                if 28144 and 8782 and 14408 and 23403 and 26144 in pos_prueba:
                    print(name + " tiene SNPs marcadores de linaje V")
                    fout.write(name + " tiene SNPs marcadores de linaje V" + "\n")
                elif 14408 and 23403 and 28144 and 8782 in pos_prueba:
                    print(name + " tiene SNPs marcadores de linaje S")
                    fout.write(name + " tiene SNPs marcadores de linaje S" + "\n")
                elif 14408 and 23403 in pos_prueba:
                    print(name + " tiene SNPs marcadores de linaje G")
                    fout.write(name + " tiene SNPs marcadores de linaje G" + "\n")
                else:
                    print(name + " no tiene SNPs marcadores de linaje")
                    fout.write(name + " no tiene SNPs marcadores de linaje" + "\n")
"""

check_create_dir(out_line_dir)

res = args.run_id + ".lineages.txt"
res_file = os.path.join(out_line_dir, res)

with open(res_file, "w") as fout:
    for r, d, f in os.walk(out_variants_dir):
        for file in f:
            if file.endswith(".tsv"):
                name = file.split(".")[0]
                path = os.path.join(r, file)
                prueba = pd.read_csv(path, sep="\t")
                pos_prueba = list(prueba["POS"])
                if 3037 and 14408 and 23403 in pos_prueba: #quitamos 8782 a pesar de que pensamos que debe estar
                    print(name + " tiene SNPs marcadores de linaje 20A")
                    fout.write(name + " tiene SNPs marcadores de linaje 20A" + "\n")
                elif 28881 and 28882 and 28883 in pos_prueba:
                    print(name + " tiene SNPs marcadores de linaje 20B")
                    fout.write(name + " tiene SNPs marcadores de linaje 20B" + "\n")
                elif 1059 and 25563 in pos_prueba:
                    print(name + " tiene SNPs marcadores de linaje 20C")
                    fout.write(name + " tiene SNPs marcadores de linaje 20C" + "\n")
                elif 8782 and 28144 in pos_prueba:
                    print(name + " tiene SNPs marcadores de linaje 19A")
                    fout.write(name + " tiene SNPs marcadores de linaje 19A" + "\n")
                #elif 8782 in pos_prueba:
                    #print(name + " tiene SNPs marcadores de linaje 19A")
                    #fout.write(name + " tiene SNPs marcadores de linaje 19A" + "\n")
                elif 11083 and 26144 in pos_prueba:
                    print(name + " tiene SNPs marcadores de linaje 19B")
                    fout.write(name + " tiene SNPs marcadores de linaje 19B" + "\n")
                else:
                    print(name + " no tiene SNPs marcadores de linaje conocidos")
                    fout.write(name + " no tiene SNPs marcadores de linaje" + "\n")

## Compare

check_create_dir(out_compare_dir)

# 3 parts: presence matrix, pairwise and dendogram

#1 - Presence matrix

def blank_database():
    new_pandas_ddtb = pd.DataFrame(columns=['Position','N', 'Samples'])
    return new_pandas_ddtb

final_ddbb = blank_database()
sample_filter_list = []
all_samples = 0
new_samples = 0
for r,d,f in os.walk(out_variants_dir):
    for file in f:
        if file.endswith(".tsv"):
            all_samples = all_samples + 1
            positions_shared = []
            positions_added = []
            file_name = file.split(".")[0]
            print(file_name)
            file_path = os.path.join(r, file)
            print(file_path)
            read_file = pd.read_csv(file_path, sep="\t") #read the file
            file_pass = read_file[read_file["PASS"] == True] ###
            final_file = file_pass[~file_pass.ALT.str.startswith("-")&~file_pass.ALT.str.startswith("+")] ###
            for position in final_file["POS"].unique(): #access to the snps of each sample
                if position not in final_ddbb["Position"].values:
                    positions_added.append(int(position))

                    new_row = len(final_ddbb.index)
                    final_ddbb.loc[new_row, "Position"] = int(position)
                    final_ddbb.loc[new_row, "Samples"] = file_name
                    final_ddbb.loc[new_row, "N"] = int(1)
                    final_ddbb.loc[new_row, file_name] = str(1)
                else:
                    positions_shared.append(int(position))

                    index_position = final_ddbb.index[final_ddbb["Position"] == int(position)][0]

                    number_samples_with_position = final_ddbb.loc[index_position, "N"]
                    names_samples_with_position = final_ddbb.loc[index_position, "Samples"]
                    new_names_samples = names_samples_with_position + "," + file_name

                    final_ddbb.loc[index_position, "N"] = number_samples_with_position + 1
                    final_ddbb.loc[index_position, "Samples"] = new_names_samples
                    final_ddbb.loc[index_position, file_name] = str(1)

final_ddbb = final_ddbb.fillna(0)
final_ddbb["Position"] = final_ddbb["Position"].astype(int)
final_ddbb["N"] = final_ddbb["N"].astype(int)

file_presence = args.run_id + ".presence.tsv"
file_presence_def = os.path.join(out_compare_dir, file_presence)

final_ddbb.to_csv(file_presence_def, sep="\t", index=False)

#2 - Pairwise

from sklearn.metrics import pairwise_distances, accuracy_score

def compare_snp_columns(sample1, sample2, df):
    jaccard_similarity = accuracy_score(df[sample1], df[sample2]) #similarities between colums
    hamming_similarity = 1 - jaccard_similarity #disagreements between colums
    snp_distance = int(hamming_similarity * (len(df.index)+1))
    return snp_distance

def snp_distance_pairwise(dataframe, output_file):
    with open(output_file, "a") as f:
        for sample1 in dataframe.iloc[:,3:].columns:
            for sample2 in dataframe.iloc[:,3:].columns:
                if sample1 != sample2:
                    snp_distance = compare_snp_columns(sample1, sample2, dataframe)
                    line_distance = "%s\t%s\t%s\n" % (sample1, sample2, snp_distance)
                    f.write(line_distance)

def import_to_pandas(file_table, header=False, sep='\t'):
    if header == False:
        #exclude first line, exclusive for vcf outputted by PipelineTB
        dataframe = pd.read_csv(file_table, sep=sep, skiprows=[0], header=None)
    else:
        #Use first line as header
        dataframe = pd.read_csv(file_table, sep=sep, header=0)

    return dataframe

file_presence_again = args.run_id + ".presence.tsv"
file_presence_again_def = os.path.join(out_compare_dir, file_presence_again) # instead of using previous one, we upload again, as in jupyter there was an error if we don't upload de output file

presence_ddbb = import_to_pandas(file_presence_again_def, header=True)

file_pairwise = args.run_id + ".snps.pairwise.tsv"
file_pairwise_path = os.path.join(out_compare_dir, file_pairwise)

snp_distance_pairwise(presence_ddbb, file_pairwise_path)

#3 - Dendrogram

import scipy.cluster.hierarchy as shc
import matplotlib.pyplot as plt

def dendogram_dataframe(dataframe, output_file):
    dataframe_only_samples = dataframe.set_index(dataframe['Position'].astype(int)).drop(['Position','N','Samples'], axis=1) #extract three first colums and use 'Position' as index
    labelList = dataframe_only_samples.columns.tolist()
    Z = shc.linkage(dataframe_only_samples.T, method='average') #method='single'

    plt.rcParams['lines.linewidth'] = 8 #Dendrogram line with
    plt.rcParams['xtick.major.size'] = 10 #Only affect to tick (line) size
    plt.rcParams.update({'font.size': 30}) #Increase x tick label size
    #plt.tick_params(labelsize=30)
    plt.figure(figsize=(30, 50))
    plt.ylabel('samples', fontsize=30)
    plt.xlabel('snp distance', fontsize=30)

    shc.dendrogram(Z, labels=labelList, orientation='left', distance_sort='descending', show_leaf_counts=True, color_threshold=10, leaf_font_size=20)

    plt.savefig(output_file, format="png")

png_file = args.run_id + ".snp.dendrogram.png"
png_dend_file = os.path.join(out_compare_dir, png_file)

dendogram_dataframe(presence_ddbb, png_dend_file)

## Edit variant files to quit indels and put the final file in another folder

check_create_dir(out_onlysnps_variants)

for r, d, f in os.walk(out_variants_dir):
    for file in f:
        new = os.path.join(r, file)
        if new.endswith("tsv"):
            one = new.split("/")[-1]
            two = one.split(".")[0]
            #print(two)
            read_file = pd.read_csv(new, sep="\t")
            file_pass = read_file[read_file["PASS"] == True] ###
            final_file = file_pass[~file_pass.ALT.str.startswith("-")&~file_pass.ALT.str.startswith("+")]
            final_name = two + ".tsv"
            path_dest = out_onlysnps_variants
            final_path = os.path.join(path_dest, final_name)
            final_file.to_csv(final_path, sep="\t", index=False)

# We move it to the bottom because this step lasts more, so in the meantime we could analize the rest of the results
## Phylogeny and trees

# legacy

"""
for root, dirs, file in os.walk(out_consensus_dir):
    for f in file:
        if f.endswith(".fa"):
            new = os.path.join(root, f)
            with open(new, "r+") as f:
                for line in f:
                    if not line.startswith(">"):
                        new_line = line.rstrip()
                        #print(len(new_line))
                        content_n = "N" * 29903
                        #print(len(content_n))
                        if new_line == content_n:
                            name = new.split(".")[0]
                            print(name) # to see which is the name; it does work in juyter but not here
                            name_two = name.split("/")[-1]
                            print(name_two + " has not enough coverage and will be removed from the analysis")
                            f.truncate(0)
"""

"""
for root, dirs, file in os.walk(out_consensus_dir):
    for f in file:
        if f.endswith(".fa"):
            new = os.path.join(root, f)
            with open(new, "r+") as f:
                for line in f:
                    if not line.startswith(">"):
                        new_line = line.rstrip()
                        #print(len(new_line))
                        content_n = "N" * 29903
                        #print(len(content_n))
                        content_na = "N" * 29868 + "A" * 35
                        #print(len(content_n_a))
                        content_na_dif = "N" * 29870 + "A" * 33
                        content_nan = "N" * 29870 + "A" * 30 + "N" * 3
                        if new_line == content_n:
                            name = new.split(".")[0]
                            print(name) # to see which is the name; it does work in juyter but not here
                            name_two = name.split("/")[-1]
                            print(name_two + " has not enough coverage and will be removed from the analysis")
                            f.truncate(0)
                        elif new_line == content_na:
                            name = new.split(".")[0]
                            print(name) # to see which is the name; it does work in juyter but not here
                            name_two = name.split("/")[-1]
                            print(name_two + " has not enough coverage and will be removed from the analysis")
                            f.truncate(0)
                        elif new_line == content_nan:
                            name = new.split(".")[0]
                            print(name) # to see which is the name; it does work in juyter but not here
                            name_two = name.split("/")[-1]
                            print(name_two + " has not enough coverage and will be removed from the analysis")
                            f.truncate(0)
                        elif new_line == content_na_dif:
                            name = new.split(".")[0]
                            print(name) # to see which is the name; it does work in juyter but not here
                            name_two = name.split("/")[-1]
                            print(name_two + " has not enough coverage and will be removed from the analysis")
                            f.truncate(0)
"""
"""
for root, dirs, file in os.walk(out_consensus_dir):
    for f in file:
        if f.endswith(".fa"):
            new = os.path.join(root, f)
            with open(new, "r+") as f:
                for line in f:
                    if not line.startswith(">"):
                        new_line = line.rstrip()
                        #print(len(new_line))
                        content_n = "N" * 29866
                        #print(len(content_n))
                        if content_n in new_line:
                            name = new.split(".")[0]
                            print(name) # to see which is the name; it does work in juyter but not here
                            name_two = name.split("/")[-1]
                            print(name_two + " has not enough coverage and will be removed from the analysis")
                            f.truncate(0)
"""
for root, dirs, file in os.walk(out_consensus_dir):
    for f in file:
        if f.endswith(".fa"):
            new = os.path.join(root, f)
            with open(new, "r+") as f:
                for line in f:
                    if not line.startswith(">"):
                        new_line = line.rstrip()
                        count = 0
                        letter_count = 0
                        #print(len(new_line))
                        for char in new_line:
                            if char.isalpha():
                                count += 1
                        for e in new_line:
                            if e == "N":
                                letter_count += 1
                        p = float(letter_count) / float(count) * 100
                        #print(len(content_n))
                        if p > 50:
                            name = new.split(".")[0]
                            print(name) # to see which is the name; it does work in juyter but not here
                            name_two = name.split("/")[-1]
                            print(name_two + " has not enough coverage and will be removed from the analysis")
                            f.truncate(0)

# cat all the consensus files in just one and quit the consensus word from the fasta name in order not to haver problemas eventually

pre_converge = os.path.join(out_consensus_dir, "*.fa")

converge = args.run_id + ".fa"
converge_file = os.path.join(out_consensus_dir, converge)

cmd_join = ["cat " + pre_converge + " | sed 's/Consensus_//g' > " + converge_file]
subprocess.run(cmd_join, shell=True)

        # run clustalo to align the multifasta file (if not, iqtree will raise an error, even if you choose fasta format)

check_create_dir(out_phyl_dir)

converged = args.run_id + ".fa"
converged_file = os.path.join(out_consensus_dir, converged)

phy = args.run_id + ".phy"
phy_file = os.path.join(out_phyl_dir, phy)

cmd_clustal = ["clustalo -i " + converged_file + " -o " + phy_file + " --outfmt=phy --threads 8"]
subprocess.run(cmd_clustal, shell=True)

        # run iqtree to obtain tree file. with this file treefile we upload to itol

phy_in = args.run_id + ".phy"
phy_file_in = os.path.join(out_phyl_dir, phy)

cmd_iq = ["iqtree -s " + phy_file_in + " -nt AUTO"]
subprocess.run(cmd_iq, shell=True)

## PANGOLIN LINEAGE

"""
pangolin_on = ["conda activate pangolin"]
subprocess.run(pangolin_on, shell=True)

converged = args.run_id + ".fa"
converged_file = os.path.join(out_consensus_dir, converged)

pangolin_file = args.run_id + ".lineage.report.csv"
final_pangolin_file = os.path.join(out_line_dir, pangolin_file)

run_pangolin = ["pangolin " + converged_file + " -o " + final_pangolin_file]
subprocess.run(run_pangolin, shell=True)

pangolin_off = ["conda deactivate"]
subprocess.run(pangolin_off, shell=True)
"""

# As it is very difficult to activate a conda environment from a python script, we will create and independent
# BASH script to activate and run the pangolin program, and we call it from here.
# If everything goes ok, we will have a LIneage

cmd_pangolin = ["bash -i pangolin.sh"]
subprocess.run(cmd_pangolin, shell=True)


from torch.utils.data import Dataset, DataLoader
import torch
from torchtext.datasets import BABI20

def dataloader(batch_size, memory_size, task, joint, tenK):
    train_iter, valid_iter, test_iter = BABI20.iters(
        batch_size=batch_size, memory_size=memory_size, task=task, joint=joint, tenK=tenK, device=torch.device("cpu"))
    return train_iter, valid_iter, test_iter, train_iter.dataset.fields['query'].vocab

class makeBatch:
    def __init__(self,q,c,s_idx,e_idx,padding_idx=0):
        self.q=q
        self.c=c

        self.s_idx=s_idx
        self.e_idx=e_idx

    def __getitem__(self, index):
        return self.q[index], self.c[index],self.s_idx[index],self.e_idx[index]

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

def pad_sequence(batch):

    Q_batch, C_batch,s_ix,e_ix = zip(*batch)

    batch_size=len(Q_batch)

    Q_max_seq_length = max([len(x) for x in Q_batch])
    Q_res = []

    for seq in Q_batch:
        if len(seq) < Q_max_seq_length:
            pad_seq = torch.LongTensor(seq + [0] * (Q_max_seq_length - len(seq)))
            Q_res.append(pad_seq)
        else:
            Q_res.append(torch.LongTensor(seq))
    ##
    C_max_seq_length = max([len(x) for x in C_batch])
    C_res = []

    for seq in C_batch:
        if len(seq) < C_max_seq_length:
            pad_seq = torch.LongTensor(seq + [0] * (C_max_seq_length - len(seq)))
            C_res.append(pad_seq)
        else:
            C_res.append(torch.LongTensor(seq))

    return torch.cat(Q_res).reshape(batch_size, Q_max_seq_length), torch.cat(C_res).reshape(batch_size,C_max_seq_length),\
          torch.LongTensor(s_ix).reshape(batch_size),torch.LongTensor(e_ix).reshape(batch_size)

def num_parameter(model):
    sum=0
    for p in model.parameters():
        if p.requires_grad:
            sum+=p.numel()
    return sum


class Solution:
    def findUnsortedSubarray(self, nums: List[int]) -> int:
        cp = nums[:]
        nums.sort()
        n = len(nums)
        l,r = -1,0
        for i in range(n):
            if l==-1 and nums[i] != cp[i]:
                l = i
            elif l != -1 and nums[i] != cp[i]:
                r = i

        if l == -1:
            return 0
        return r - l + 1


from LogRegression import MyLogisticRegression
from utils import loadData, zScoreNormalization, plotData
from sklearn import linear_model
import numpy as np

def accuracy(real, computed):
    return sum([1 if real[i] == computed[i] else 0 for i in range(len(real))]) / len(real)

if __name__ == '__main__':
    inputs, outputs, labels, feature1, feature2, feature3, feature4 = loadData()

    numericalLabels = [i for i in range(len(labels))]
    numericalOutputs = [0 for i in range(len(outputs))]
    for i in range(len(labels)):
        for j in range(len(outputs)):
            if outputs[j] == labels[i]:
                numericalOutputs[j] = i

    np.random.seed(5)
    indexes = [i for i in range(len(inputs))]
    trainSample = np.random.choice(indexes, int(0.8 * len(inputs)), replace=False)  # 80% train data
    testSample = [i for i in indexes if i not in trainSample]  # 20% test data

    trainInputs = [inputs[i] for i in trainSample]
    trainOutputs = [numericalOutputs[i] for i in trainSample]

    testInputs = [inputs[i] for i in testSample]
    testOutputs = [numericalOutputs[i] for i in testSample]

    # print(trainInputs)
    # print(testInputs)
    trainInputs, testInputs = \
        zScoreNormalization(trainInputs, testInputs)
    # print(trainInputs)
    # print(testInputs)

    trainFeature1 = [sample[0] for sample in trainInputs]
    trainFeature2 = [sample[1] for sample in trainInputs]
    trainFeature3 = [sample[2] for sample in trainInputs]
    trainFeature4 = [sample[3] for sample in trainInputs]

    plotData(trainFeature1, trainFeature2, trainFeature3, trainFeature4, "Train data")

    # tool
    classifier = linear_model.LogisticRegression(max_iter=1000)
    classifier.fit(trainInputs, trainOutputs)
    print("-----------tool-----------")
    print("Real: " + str(testOutputs))
    print("Comp: " + str(list(classifier.predict(testInputs))))
    print("Accuracy: " + str(accuracy(testOutputs, list(classifier.predict(testInputs)))))

    # manual

    classifierIrisSetosa = MyLogisticRegression()
    classifierIrisVersicolor = MyLogisticRegression()
    classifierIrisVirginica = MyLogisticRegression()

    trainOutputSetosa = [0 if sample != 0 else 1 for sample in trainOutputs]
    trainOutputVersicolor = [0 if sample != 1 else 1 for sample in trainOutputs]
    trainOutputVirginica = [0 if sample != 2 else 1 for sample in trainOutputs]

    classifierIrisSetosa.fit(trainInputs, trainOutputSetosa)
    classifierIrisVersicolor.fit(trainInputs, trainOutputVersicolor)
    classifierIrisVirginica.fit(trainInputs, trainOutputVirginica)

    probabilitiesSetosa = classifierIrisSetosa.predict(testInputs)
    probabilitiesVersicolor = classifierIrisVersicolor.predict(testInputs)
    probabilitiesVirginica = classifierIrisVirginica.predict(testInputs)

    # print(probabilitiesSetosa)
    # print(probabilitiesVersicolor)
    # print(probabilitiesVirginica)
    computedProbabilities = [[probabilitiesSetosa[i], probabilitiesVersicolor[i], probabilitiesVirginica[i]] for i in range(len(testOutputs))]

    # print(computedProbabilities)
    computedOutputs = [numericalLabels[np.argmax(sample)] for sample in computedProbabilities]

    print("----------manual----------")
    print("Real: " + str(testOutputs))
    print("Comp: " + str(computedOutputs))
    print("Accuracy: " + str(accuracy(testOutputs, computedOutputs)))

Machine Learning/Logistic Regression algorithm/main.py


import numpy as np

base_to_int_map = {'a' : 0, 'A' : 0, 't': 1, 'T':1, 'g' : 2, 'G' : 2, 'c' : 3, 'C' : 3, 'n' : -1, 'N': -1}
complement_map = np.array([1,0,3,2])

def write_chr_size(chr_names,chr_seq,file_name):
	f=open(file_name, 'w')
	for chr_idx in range(len(chr_names)):
		f.write(chr_names[chr_idx] + '\t' + str(len(chr_seq[chr_idx]))+'\n')
	f.close()

#### Functions to read the data
def load_genome(fasta_file):
	f = open(fasta_file, "r")
	fl = f.readlines()
	f.close()
	sequences = []
	sequences_name = []
	chr_number = 0
	for line in fl:
		if line[0] == ">":
			sequences_name.append(line[1:-1])
			if (chr_number!=0): sequences.append(my_seq)
			my_seq = ''
			chr_number+=1
		else:
			my_seq+=line[:-1]
	sequences.append(my_seq)
	return sequences_name,sequences

def load_bed_file(bed_file,chr_names,chr_seq):
	#For now cannot deal with headers, chr_names must be a list of of strings where each entry is the name of a chromosome

	# seq_pos : N x 4 numpy array: where N is the number of reads.
	#			First column indicates the chromosome number in the chr_seq array
	#			Second column indicates the starting position of the read
	#			Third column indicates the ending position of the read
	#			Last column idicates the strand: 0=forward, 1=backward
	#			The sequence of the read i is therefore given by chr_seq[seq_pos[i,0]][seq_pos[i,1]: seq_pos[i,2]]
	# n_methyl: numpy array of length N+1: Cumulative number of methiylations per read
	#			It is used as a pointer for the methyl_rel_pos array
	#			The list of reltive positions of methylations of read i is given from position n_methyl[i] to n_methyl[i+1] in methyl_rel_pos
	# methyl_rel_pos: 1D numpy array with all the relative

	#Count number of lines
	f = open(bed_file, "r")
	n_lines = sum(1 for i in f)
	f.close()

	f = open(bed_file, "r")
	input_dat = np.zeros((2*n_lines,4), dtype=int)
	n_methyl = np.zeros(2*n_lines+1, dtype=int) #cumulative number of methylations per read
	counter_methyl = 0
	for i in range(n_lines):
		line = f.readline()
		line_split = line.strip().split()

		#counter_methyl += int(line_split[9])-2
		input_dat[(2*i):(2*(i+1)),0] = chr_names.index(line_split[0]) #chr
		input_dat[(2*i):(2*(i+1)),1] = int(line_split[1]) #start
		input_dat[(2*i):(2*(i+1)),2] = int(line_split[2]) #end
		input_dat[2*i,3] = 0 #srand
		input_dat[2*i+1,3] = 1
		#the [1:-1] is needed so that we remove the first and last relative mthylation positions which are artefacts
		methyl_rel_pos_read = np.array(list(map(int,line_split[11].split(","))),dtype=int)[1:-1]
		counter_A = 0
		counter_T = 0
		for pos in methyl_rel_pos_read:
			if (chr_seq[input_dat[2*i,0]][input_dat[2*i,1] + pos] in 'aA'):
				counter_A +=1
			elif(chr_seq[input_dat[2*i,0]][input_dat[2*i,1] + pos] in 'tT'):
				counter_T +=1
		n_methyl[2*i+1] = n_methyl[2*i] + counter_A
		n_methyl[2*(i+1)] = n_methyl[2*i+1] + counter_T
	f.close()

	f = open(bed_file, "r")
	methyl_rel_pos = np.zeros(n_methyl[-1], dtype=int) #n_methyl[-1] is the total number of methylated bases
	for i in range(n_lines):
		line = f.readline()
		line_split = line.strip().split()
		methyl_rel_pos_read = np.array(list(map(int,line_split[11].split(","))),dtype=int)[1:-1]

		if (len(methyl_rel_pos_read) == 0): continue

		pointer_forward_strand = n_methyl[2*i]
		pointer_backward_strand = n_methyl[2*i+1]
		for pos in methyl_rel_pos_read:
			if  chr_seq[input_dat[2*i,0]][input_dat[2*i,1] + pos] in 'aA':
				methyl_rel_pos[pointer_forward_strand] = pos
				pointer_forward_strand +=1
			elif chr_seq[input_dat[2*i,0]][input_dat[2*i,1] + pos] in 'tT':
				methyl_rel_pos[pointer_backward_strand] = pos
				pointer_backward_strand +=1
			else:
				raise("Base ", chr_seq[input_dat[2*i,0]][input_dat[2*i,1] + pos]," is methylated, only A or T can be methylated")
	f.close()
	return input_dat,n_methyl,methyl_rel_pos

def load_bed_file_peaks(bed_file,chr_names,chr_seq):
	f = open(bed_file, "r")
	n_lines = sum(1 for i in f)
	f.close()

	f = open(bed_file, "r")
	data = np.zeros((n_lines,3), dtype=int)
	for i in range(n_lines):
		line = f.readline()
		line_split = line.strip().split()
		data[i,0] = chr_names.index(line_split[0]) #chr
		data[i,1] = int(line_split[1]) #start
		data[i,2] = int(line_split[2]) #end
	return data
	f.close()

def histogram_fraction_methylated_A_vs_T_per_read(bed_file,chr_names,chr_seq):

	f = open(bed_file, "r")
	n_lines = sum(1 for i in f)
	f.close()

	fraction_methylated = np.zeros(n_lines)
	f = open(bed_file, "r")
	for i in range(n_lines):
		line = f.readline()
		line_split = line.strip().split()

		#counter_methyl += int(line_split[9])-2
		chr_idx = chr_names.index(line_split[0])
		s = int(line_split[1]) #start

		#the [1:-1] is needed so that we remove the first and last relative mthylation positions which are artefacts
		methyl_rel_pos_read = np.array(list(map(int,line_split[11].split(","))),dtype=int)[1:-1]
		counter_A = 0
		for pos in methyl_rel_pos_read:
			if (chr_seq[chr_idx][s + pos] in 'aA'):
				counter_A +=1
		if (counter_A>0):fraction_methylated[i] =  counter_A*1.0/len(methyl_rel_pos_read)

	f.close()

	plt.hist(fraction_methylated, bins=100)
	plt.title("Fraction of methylated bases per read that are Adenines")
	plt.xlabel('Fraction of Adenines')
	plt.show()

#### Functions to generate avergae prgile:
def get_average_profile(seq_pos,n_methyl,methyl_rel_pos, chr_seq):
	chr_length_cum = np.zeros(len(chr_seq)+1, dtype = int)
	for chr_idx in range(len(chr_seq)):
		chr_length_cum[chr_idx+1] = chr_length_cum[chr_idx] + len(chr_seq[chr_idx])

	Length_genome = chr_length_cum[-1]
	coverage_forward = np.zeros(Length_genome, dtype = int)
	coverage_backward = np.zeros(Length_genome, dtype = int)
	methyl_profile = np.zeros(Length_genome, dtype = float)

	for i in range(len(seq_pos)):
		absolute_start_pos = chr_length_cum[seq_pos[i,0]] + seq_pos[i,1]
		absolute_end_pos = chr_length_cum[seq_pos[i,0]] + seq_pos[i,2]
		if seq_pos[i,3] == 0:
			coverage_forward[absolute_start_pos:absolute_end_pos] +=1
		else:
			coverage_backward[absolute_start_pos:absolute_end_pos] +=1

		abs_pos_methyl = absolute_start_pos + methyl_rel_pos[n_methyl[i]:n_methyl[i+1]]
		methyl_profile[abs_pos_methyl] += 1

	i = 0
	for chr_idx in range(len(chr_seq)):
		for base_idx in range(len(chr_seq[chr_idx])):
			if (methyl_profile[i] != 0):
				if chr_seq[chr_idx][base_idx] in 'aA':
					methyl_profile[i] /= (1.0*coverage_forward[i])
				elif chr_seq[chr_idx][base_idx] in 'tT':
					methyl_profile[i] /= (1.0*coverage_backward[i])
			i+=1

	coverage = coverage_forward + coverage_backward

	return methyl_profile, coverage

def smooth_average_profile(profile,chr_names,chr_seq,w):
	#w is the number of A or T that will be used for the moving average on each side of the window
	#The G and C will be given an accesibility profile equal to the nearest A or T on their right

	#The first w and last w A/T bases will be given a smoothed profile of 0, I should change that
	w = int(w)
	smoothed_profile = np.zeros(len(profile))
	a = 'aAtT'
	start_pos_abs_chr = 0
	for chr_idx in range(len(chr_names)):
		counter_AT = 0
		pos = np.zeros(2*w+1, dtype = int)
		for j in range(len(chr_seq[chr_idx])):
			if chr_seq[chr_idx][j] in a:
				if (counter_AT < 2*w+1):
					pos[counter_AT] = start_pos_abs_chr + j
					counter_AT +=1
				else:
					pos = np.roll(pos, -1)
					pos[2*w] = start_pos_abs_chr + j
					counter_AT +=1

				if (counter_AT>= 2*w+1):
					smoothed_profile[pos[w]:pos[w+1]] = np.sum(profile[pos])/(2*w+1)
			if j>= len(profile) : break
		if j>= len(profile) : break
		start_pos_abs_chr += len(chr_seq[chr_idx])
	return smoothed_profile

def split_average_profile_with_peaks_to_fixed_size_inputs(profile, chr_seq,peaks,w):
	x = np.zeros((len(peaks),3), dtype = int)
	y = np.zeros((len(peaks),2*w+1))

	chr_length_cum = np.zeros(len(chr_seq)+1, dtype = int)
	for chr_idx in range(len(chr_seq)):
		chr_length_cum[chr_idx+1] = chr_length_cum[chr_idx] + len(chr_seq[chr_idx])

	x[:,0] = peaks[:,0]
	mid = np.round((peaks[:,2] + peaks[:,1])/2)
	x[:,1] = mid-w
	x[:,2] = mid+w+1

	for i in range(len(peaks)):
		#print((chr_length_cum[x[i,0]]+x[i,1]),(chr_length_cum[x[i,0]]+x[i,2]), 2*w+1)
		y[i,:] = profile[(chr_length_cum[x[i,0]]+x[i,1]):(chr_length_cum[x[i,0]]+x[i,2])]

	return x,y

def get_mu_and_std_average_profile_peaks(y):
	y = y.flatten()
	return np.mean(y), np.std(y)

#### Process data for NN
def seq_to_1_hot(seq,strand):
	# This is not vectorised!!!!
	idxs = [base_to_int_map[b] for b in seq]
	if(strand ==1):
		idxs = np.flip(complement_map[idxs])
	x = np.zeros((len(seq),4))
	for i in range(len(seq)):
		if (idxs[i]!=-1): x[i,idxs[i]] = 1
	return x

def seq_pos_to_1_hot(seq, start, end,strand):
	#For now, if the base is not known, all zeros.
	#If the segments goes after the end or before the beginning of the chr after augmentation,
	#those "virtual bases" are also set to all zeros
	subseq = seq[max(start,0):min(end,len(seq))]

	x = seq_to_1_hot(subseq,strand)
	if (start<0):
		if(strand == 0): x = np.concatenate((np.zeros((-start,4),dtype=int), x), axis=0)
		else:  x = np.concatenate((x, np.zeros((-start,4),dtype=int)), axis=0)
	if (end>len(seq)-1):
		if(strand == 0): x = np.concatenate((x,np.zeros((end-len(seq),4),dtype=int)), axis=0)
		else: x = np.concatenate((np.zeros((end-len(seq),4),dtype=int),x), axis=0)
	return x

def non_sparse_representation_methylation(n_methyl, methyl_rel_pos,read_number,output_length,strand):
	methyl = np.zeros((1,output_length), dtype=int)

	start_pointer_methyl = n_methyl[read_number]
	end_pointer_methyl = n_methyl[read_number+1]

	methyl_pos_read = methyl_rel_pos[start_pointer_methyl:end_pointer_methyl]

	methyl[0,methyl_pos_read] = 1
	if (strand == 1): methyl = np.flip(methyl,1)
	return methyl

def generator_data(seq_pos,n_methyl,methyl_rel_pos, chr_seq,augmentation_value,batch_size):
	# seq_pos : N x 4 numpy array: where N is the number of reads.
	#			First column indicates the chromosome number in the chr_seq array
	#			Second column indicates the starting position of the read
	#			Third column indicates the ending position of the read
	#			Last column idicates the strand: 0=forward, 1=backward
	#			The sequence of the read i is therefore given by chr_seq[seq_pos[i,0]][seq_pos[i,1]: seq_pos[i,2]]
	# n_methyl: numpy array of length N+1: Cumulative number of methiylations per read
	#			It is used as a pointer for the methyl_rel_pos array
	#			The list of reltive positions of methylations of read i is given from position n_methyl[i] to n_methyl[i+1] in methyl_rel_pos
	# methyl_rel_pos: 1D numpy array with all the relative
	# chr_seq: List of chromosome sequence given as a string.
	# augmentation_value: int : number of bases added to the input sequences of the both left and right
	# batch_size: int: number of samples given at each iteration of this generator function

	output_length = seq_pos[0,2] - seq_pos[0,1]
	input_length = output_length+2*augmentation_value
	batch_input = np.zeros((batch_size, input_length, 4))
	batch_output = np.zeros((batch_size,output_length))
	i = 0
	while True:
		for read_number in range(len(seq_pos)): #The reads should have been shuffled prior to this operation!
			chr_indx = seq_pos[read_number,0]
			read_s = seq_pos[read_number,1]
			read_e = seq_pos[read_number,2]
			strand = seq_pos[read_number,3]

			batch_input[i] = seq_pos_to_1_hot(chr_seq[chr_indx],read_s-augmentation_value, read_e + augmentation_value,strand)
			batch_output[i] = non_sparse_representation_methylation(n_methyl, methyl_rel_pos,read_number,output_length,strand)

			i+=1

			if(i%batch_size == 0):
				yield (batch_input, batch_output)
				i=0
				batch_input = np.zeros((batch_size, input_length, 4))
				batch_output = np.zeros((batch_size,output_length))

def generator_data_average_profile(x,y,chr_seq,augmentation_value,batch_size):
	output_length = x[0,2] - x[0,1]
	input_length = output_length+2*augmentation_value
	batch_input = np.zeros((batch_size, input_length, 4))
	batch_output = np.zeros((batch_size,output_length))
	i = 0
	while True:
		for read_number in range(len(x)): #The reads should have been shuffled prior to this operation!
			chr_indx = x[read_number,0]
			read_s = x[read_number,1]
			read_e = x[read_number,2]

			batch_input[i] = seq_pos_to_1_hot(chr_seq[chr_indx],read_s-augmentation_value, read_e + augmentation_value,0)
			batch_output[i] = y[read_number]

			i+=1

			if(i%batch_size == 0):
				yield (batch_input, batch_output)
				i=0
				batch_input = np.zeros((batch_size, input_length, 4))
				batch_output = np.zeros((batch_size,output_length))

def generator_input_average_profile(x,chr_seq,augmentation_value,batch_size):
	output_length = x[0,2] - x[0,1]
	input_length = output_length+2*augmentation_value
	batch_input = np.zeros((batch_size, input_length, 4))
	i = 0
	while True:
		for read_number in range(len(x)): #The reads should have been shuffled prior to this operation!
			chr_indx = x[read_number,0]
			read_s = x[read_number,1]
			read_e = x[read_number,2]
			batch_input[i] = seq_pos_to_1_hot(chr_seq[chr_indx],read_s-augmentation_value, read_e + augmentation_value,0)
			i+=1
			if(i%batch_size == 0):
				yield batch_input
				i=0
				batch_input = np.zeros((batch_size, input_length, 4))

def break_reads_to_fixed_size(seq_pos,n_methyl,methyl_rel_pos,fixed_size):
	n_samples = np.sum(np.floor((seq_pos[:,2]-seq_pos[:,1])/fixed_size).astype('int32'))
	seq_pos_new = np.zeros((n_samples,4), dtype=int)
	n_methyl_new =  np.zeros(n_samples+1, dtype=int)
	methyl_rel_pos_new = np.zeros(len(methyl_rel_pos), dtype=int)

	counter_input = 0
	tot_methyls = 0
	for read_number in range(len(seq_pos)):
		n_methyls_observed_read = 0
		counter_input_start_read = counter_input
		for i in range((np.floor((seq_pos[read_number,2]-seq_pos[read_number,1])/fixed_size)).astype('int32')):

			seq_pos_new[counter_input,0] = seq_pos[read_number,0]
			seq_pos_new[counter_input,1] = seq_pos[read_number,1] + i*fixed_size
			seq_pos_new[counter_input,2] = seq_pos[read_number,1] + (i+1)*fixed_size
			seq_pos_new[counter_input,3] = seq_pos[read_number,3]
			#print(n_methyl[read_number]+n_methyls_observed, n_methyl[read_number+1])
			for j in range(n_methyl[read_number]+n_methyls_observed_read, n_methyl[read_number+1]):
				if (methyl_rel_pos[j] >= ((i+1)*fixed_size)): break
				methyl_rel_pos_new[tot_methyls] = methyl_rel_pos[j] - i*fixed_size
				n_methyls_observed_read += 1
				tot_methyls+=1

			n_methyl_new[counter_input+1] = tot_methyls

			counter_input+=1

	methyl_rel_pos_new = methyl_rel_pos_new[0:tot_methyls]
	return seq_pos_new,n_methyl_new,methyl_rel_pos_new

def shuffle_data(seq_pos,n_methyl,methyl_rel_pos):
	permutation = np.random.permutation(len(seq_pos))

	seq_pos_new = seq_pos[permutation,:]
	n_methyl_new = np.zeros(len(n_methyl), dtype=int)
	methyl_rel_pos_new = np.zeros(len(methyl_rel_pos), dtype=int)

	for i in range(len(seq_pos)):
		n_methyl_new[i+1] = n_methyl_new[i]+n_methyl[permutation[i]+1]- n_methyl[permutation[i]]
		methyl_rel_pos_new[n_methyl_new[i]:n_methyl_new[i+1]] = methyl_rel_pos[n_methyl[permutation[i]]:n_methyl[permutation[i]+1]]

	return seq_pos_new,n_methyl_new,methyl_rel_pos_new

def split_training_testing(seq_pos,n_methyl,methyl_rel_pos,fraction_training):
	n_training = round(fraction_training*len(seq_pos))

	seq_pos_training = seq_pos[:n_training,:]
	seq_pos_testing = seq_pos[n_training:,:]

	n_methyl_training = n_methyl[:n_training+1]
	n_methyl_testing = n_methyl[n_training:]-n_methyl[n_training]

	methyl_rel_pos_training = methyl_rel_pos[:n_methyl[n_training]]
	methyl_rel_pos_testing = methyl_rel_pos[n_methyl[n_training]:]

	return seq_pos_training,n_methyl_training,methyl_rel_pos_training, seq_pos_testing,n_methyl_testing,methyl_rel_pos_testing

### Write the data to external files
def average_profile_to_bedgraph_file(profile,chr_names,chr_seq, file_name):
	header = 'track type=bedGraph name="BedGraph Format" description="BedGraph format" visibility=full color=200,100,0 altColor=0,100,200 priority=20\n'
	f=open(file_name, 'w')
	f.write(header)
	counter = 0
	for chr_idx in range(len(chr_names)):
		for j in range(len(chr_seq[chr_idx])):
			f.write(chr_names[chr_idx] + ' ' + str(j)+' '+str(j+1)+ ' '+str(profile[counter])+'\n')
			counter +=1
	f.close()

def average_profile_to_wig_file(profile,chr_names,chr_seq, file_name):
	f=open(file_name, 'w')
	counter = 0
	for chr_idx in range(len(chr_names)):
		f.write('fixedStep chrom='+chr_names[chr_idx] +' start=1 step=1\n')
		for j in range(len(chr_seq[chr_idx])):
			f.write(str(profile[counter])+'\n')
			counter +=1
	f.close()

def predictions_NN_to_wig_file(x,y,chr_names,file_name):
	f=open(file_name, 'w')
	for i in range(len(y)):
		f.write('fixedStep chrom='+chr_names[x[i,0]] +' start='+ str(x[i,1])+' step=1\n')
		for j in range(x[i,2] - x[i,1]):
			f.write(str(y[i,j])+'\n')
	f.close()


#!/usr/bin/env pybricks-micropython
from pybricks.hubs import EV3Brick
from pybricks.ev3devices import (Motor, TouchSensor, ColorSensor,
                                 InfraredSensor, UltrasonicSensor, GyroSensor)
from pybricks.parameters import Port, Stop, Direction, Button, Color
from pybricks.tools import wait, StopWatch, DataLog
from pybricks.robotics import DriveBase
from pybricks.media.ev3dev import SoundFile, ImageFile
import simulation_cube
import os
import sys

ev3 = EV3Brick()

motorA = Motor(Port.A) #A馬達 正值手臂抬起
motorB = Motor(Port.B) #B馬達 正值底盤順時針轉
motorC = Motor(Port.C) #C馬達 正值顏色手臂後退
ColorSensorS4 = ColorSensor(Port.S4)

cube_x, cube_y = 6, 9
cube3d = [[0 for j in range(cube_y)] for k in range(cube_x)] #儲存掃描顏色的2維陣列
#print(cube3d) #顯示空的陣列(檢查用)
print()

#U藍、R紅、F黃、D綠、L橘、B白
#U上、R右、F前、D下、L左、B後
x, y = 0, 0
def scan_color(): #讀取單顆顏色
    global x, y
    RGB = [0, 0, 0]
    RGB = ColorSensorS4.rgb()
    wait(50)
    #cube3d[x][y] = ColorSensorS4.rgb() #RGB測試用
    strcolor = ColorSensorS4.color()

    #掃描方案B 範圍判斷
    '''
    if RGB[0] < 19 and RGB[1] < 46 and RGB[2] > 42:
        cube3d[x][y] = 'U'
    if RGB[0] > 41 and RGB[1] < 41 and RGB[2] < 28:
        cube3d[x][y] = 'L'
    if RGB[0] > 44 and RGB[1] > 58 and RGB[2] < 36:
        cube3d[x][y] = 'F'
    if RGB[0] < 41 and RGB[1] < 31 and RGB[2] < 27:
        cube3d[x][y] = 'R'
    if RGB[0] > 40 and RGB[1] > 80 and RGB[2] > 80:
        cube3d[x][y] = 'B'
    if RGB[0] < 23 and RGB[1] > 40 and RGB[2] < 29:
        cube3d[x][y] = 'D'
    '''
    '''
    if RGB[0] > RGB[1] and RGB[0] > 47:
        cube3d[x][y] = 'L' #橘
    if RGB[0] > RGB[1] and RGB[0] <= 47:
        cube3d[x][y] = 'R' #紅
    '''
    if RGB[0] > RGB[1] and RGB[0] + RGB[1] >= 80:
        cube3d[x][y] = 'L'
    if RGB[0] > RGB[1] and RGB[0] + RGB[1] < 80:
        cube3d[x][y] = 'R'
    if RGB[0] < RGB[1] and RGB[1] < RGB[2]:
        cube3d[x][y] = 'U' #藍
    if RGB[0] < RGB[1] and RGB[1] > 60 and RGB[2] > 60:
        cube3d[x][y] = 'B' #白
    if RGB[0] < RGB[1] and RGB[1] > RGB[2] and RGB[0] > RGB[2]:
        cube3d[x][y] = 'F' #黃
    if RGB[0] < RGB[1] and RGB[1] > RGB[2] and RGB[0] <= RGB[2]:
        cube3d[x][y] = 'D' #綠

    '''
    #掃描方案C 例外處理
    if strcolor == Color.BLUE:
        cube3d[x][y] = 'U'
    elif strcolor == Color.YELLOW:
        cube3d[x][y] = 'F'
    elif strcolor == Color.RED:
        cube3d[x][y] = 'R'
    elif strcolor == Color.WHITE:
        cube3d[x][y] = 'B'
    elif strcolor == Color.GREEN:
        cube3d[x][y] = 'D'
    #lego掃不到橘色，會把橘當成紅色，所以先把可辨認的顏色排除，留下紅色，再用RGB光反射作區別
    if cube3d[x][y] == 'R':
        if RGB[0] > 41: #隨環境不同數字需修正
            cube3d[x][y] = 'L'
    #橘色誤判成黃色 顏色校正
    if cube3d[x][y] == 'F':
        if RGB[0] > 41 and RGB[1] < 41 and RGB[2] < 28:
            cube3d[x][y] = 'L'
    #紅色誤判成黃色 顏色校正
    if cube3d[x][y] == 'F':
        if RGB[0] < 41 and RGB[1] < 31 and RGB[2] < 27:
            cube3d[x][y] = 'R'
    '''

    '''
    #黃or綠誤判成藍色 顏色校正
    if cube3d[x][y] == 'U':
        if RGB[2] > 55:
            cube3d[x][y] = 'U'
        elif RGB[0] > 60 and RGB[2] < 45:
            cube3d[x][y] = 'F'
        elif RGB[0] < 30 and RGB[2] < 45:
            cube3d[x][y] = 'D'
    '''
    y += 1
    if y >= 9:
        y = 0
        x += 1
    wait(50)

def scan_plane(): #讀取一面的各個顏色
    motorC.run_until_stalled(400,then=Stop.HOLD, duty_limit=None) #後退到最底 1
    motorC.run_angle(400, -239*3, then=Stop.HOLD, wait=True) #往前到中間 2
    scan_color()
    motorB.run_angle(400, -45*3, then=Stop.HOLD, wait=True) #底盤逆時針轉45度
    motorC.run_angle(400, 67*3, then=Stop.HOLD, wait=True) #後退掃描角落的 3
    scan_color()
    for r in range (0, 3):
        motorB.run_angle(400, -45*3, then=Stop.HOLD, wait=True)
        motorC.run_angle(400, -31*3, then=Stop.HOLD, wait=True) #前進掃描邊邊中間4、6、8
        scan_color()
        motorB.run_angle(400, -45*3, then=Stop.HOLD, wait=True)
        motorC.run_angle(400, 31*3, then=Stop.HOLD, wait=True) #後退掃描角落5、7、9
        scan_color()
    motorB.run_angle(400, -45*3, then=Stop.HOLD, wait=True)
    motorC.run_angle(400, -32*3, then=Stop.HOLD, wait=True) #前進掃描邊邊的中間 10
    scan_color()
    motorC.run_until_stalled(400,then=Stop.HOLD, duty_limit=None) #後退到最底 11

    '''
    motorC.run_until_stalled(400,then=Stop.COAST, duty_limit=None) #後退到最底 1
    motorC.reset_angle(0)
    motorC.run_angle(400, -239*3, then=Stop.HOLD, wait=True) #往前到中間 2
    scan_color()
    motorB.run_angle(400, -50*3, then=Stop.HOLD, wait=True) #底盤逆時針轉50度
    motorC.run_angle(400, 64*3, then=Stop.HOLD, wait=True) #後退掃描角落的 3
    scan_color()
    motorB.run_angle(400, -40*3, then=Stop.HOLD, wait=True)
    motorC.run_angle(400, -30*3, then=Stop.HOLD, wait=True) #前進掃描邊邊中間
    scan_color()
    motorB.run_angle(400, -50*3, then=Stop.HOLD, wait=True)
    motorC.run_angle(400, 30*3, then=Stop.HOLD, wait=True) #後退掃描角落
    scan_color()
    for r in range (0, 2):
        motorB.run_angle(400, -40*3, then=Stop.HOLD, wait=True)
        motorC.run_angle(400, -28*3, then=Stop.HOLD, wait=True) #前進掃描邊邊中間6、8
        scan_color()
        motorB.run_angle(400, -50*3, then=Stop.HOLD, wait=True)
        motorC.run_angle(400, 28*3, then=Stop.HOLD, wait=True) #後退掃描角落7、9
        scan_color()
    motorB.run_angle(400, -40*3, then=Stop.HOLD, wait=True)
    motorC.run_angle(400, -32*3, then=Stop.HOLD, wait=True) #前進掃描邊邊的中間 10
    scan_color()
    motorC.run_until_stalled(400,then=Stop.COAST, duty_limit=None) #後退到最底 11
    motorC.reset_angle(0)
    '''

global arm_state
arm_state = 1 #手臂狀態bool 0=放下 1=抬起
def roboticarm_up(): #機器手臂抬起
    global arm_state
    wait(50)
    motorA.run_until_stalled(-300,then=Stop.HOLD, duty_limit=None)
    arm_state = 1

def roboticarm_to_cube(): #機器手臂放到方塊上，初始狀態:手臂抬起
    global arm_state
    wait(50)
    motorA.run_angle(250, 110)
    arm_state = 0

def roll_cube(): #馬達A手臂翻轉方塊，初始、結束狀態:手臂放在方塊上
    wait(50)
    motorA.run_angle(250, 110)
    motorA.run_angle(250, -110)
    arm_state = 0

def roll_cube_2times(): #馬達A手臂翻轉方塊2次，初始、結束狀態:手臂放在方塊上
    wait(50)
    for i in range(2):
        motorA.run_angle(250, 110)
        motorA.run_angle(250, -110)
        wait(50)
    arm_state = 0

def scan_roll_cube(): #馬達A手臂翻轉方塊，初始、結束狀態:手臂抬起
    roboticarm_to_cube()
    roll_cube()
    roboticarm_up()

def basemove(): #底盤順轉90度
    wait(50)
    motorB.run_angle(300, 111*3, then=Stop.HOLD, wait=True) #多的21度是先轉超過在回正
    wait(50)
    motorB.run_angle(300, -21*3, then=Stop.HOLD, wait=True)

def basemove_2times(): #底盤順轉180度
    wait(50)
    motorB.run_angle(300, 201*3, then=Stop.HOLD, wait=True) #多的21度是先轉超過在回正
    wait(50)
    motorB.run_angle(300, -21*3, then=Stop.HOLD, wait=True)

def basemove_reverse(): #底盤逆轉90度
    wait(50)
    motorB.run_angle(300, -111*3, then=Stop.HOLD, wait=True) #多的21度是先轉超過在回正
    wait(50)
    motorB.run_angle(300, 21*3, then=Stop.HOLD, wait=True)

def basemove_reverse_2times(): #底盤逆轉180度
    wait(50)
    motorB.run_angle(300, -201*3, then=Stop.HOLD, wait=True) #多的21度是先轉超過在回正
    wait(50)
    motorB.run_angle(300, 21*3, then=Stop.HOLD, wait=True)

#motorC.run_until_stalled(300,then=Stop.HOLD, duty_limit=None) #測試用 顏色感應器後退到最底

#掃描6面 初始、結束狀態:手臂抬起
scan_plane() #掃U
basemove_2times()
scan_roll_cube()

scan_plane() #掃L
basemove()
scan_roll_cube()

scan_plane() #掃F
scan_roll_cube()

scan_plane() #掃R
scan_roll_cube()

scan_plane() #掃B
basemove_reverse()
scan_roll_cube()

scan_plane() #掃D
wait(50)

#方塊位置回到掃描前狀態
roboticarm_to_cube()
roll_cube_2times()
roboticarm_up()
basemove_reverse()

simulation_cube3d = cube3d.copy() #複製一個新陣列 模擬方塊當下的狀態
kociemba_cube3d = cube3d.copy() #複製一個新陣列 儲存符合kociemba排列順序

#掃描完的方塊狀態
'''
simulation_cube3d = \
 [str(cube3d[0][5]) + str(cube3d[0][6]) + str(cube3d[0][7]) + str(cube3d[0][4]) + str(cube3d[0][0]) + str(cube3d[0][8]) + str(cube3d[0][3]) + str(cube3d[0][2]) + str(cube3d[0][1])]\
+[str(cube3d[1][3]) + str(cube3d[1][4]) + str(cube3d[1][5]) + str(cube3d[1][2]) + str(cube3d[1][0]) + str(cube3d[1][6]) + str(cube3d[1][1]) + str(cube3d[1][8]) + str(cube3d[1][7])]\
+[str(cube3d[2][5]) + str(cube3d[2][6]) + str(cube3d[2][7]) + str(cube3d[2][4]) + str(cube3d[2][0]) + str(cube3d[2][8]) + str(cube3d[2][3]) + str(cube3d[2][2]) + str(cube3d[2][1])]\
+[str(cube3d[3][5]) + str(cube3d[3][6]) + str(cube3d[3][7]) + str(cube3d[3][4]) + str(cube3d[3][0]) + str(cube3d[3][8]) + str(cube3d[3][3]) + str(cube3d[3][2]) + str(cube3d[3][1])]\
+[str(cube3d[4][5]) + str(cube3d[4][6]) + str(cube3d[4][7]) + str(cube3d[4][4]) + str(cube3d[4][0]) + str(cube3d[4][8]) + str(cube3d[4][3]) + str(cube3d[4][2]) + str(cube3d[4][1])]\
+[str(cube3d[5][7]) + str(cube3d[5][8]) + str(cube3d[5][1]) + str(cube3d[5][6]) + str(cube3d[5][0]) + str(cube3d[5][2]) + str(cube3d[5][5]) + str(cube3d[5][4]) + str(cube3d[5][3])]
'''

simulation_cube3d = \
[cube3d[0][5], cube3d[0][6], cube3d[0][7], cube3d[0][4], 'U', cube3d[0][8], cube3d[0][3], cube3d[0][2], cube3d[0][1]],\
[cube3d[1][3], cube3d[1][4], cube3d[1][5], cube3d[1][2], 'L', cube3d[1][6], cube3d[1][1], cube3d[1][8], cube3d[1][7]],\
[cube3d[2][5], cube3d[2][6], cube3d[2][7], cube3d[2][4], 'F', cube3d[2][8], cube3d[2][3],cube3d[2][2], cube3d[2][1]],\
[cube3d[3][5], cube3d[3][6], cube3d[3][7], cube3d[3][4], 'R', cube3d[3][8], cube3d[3][3], cube3d[3][2], cube3d[3][1]],\
[cube3d[4][5], cube3d[4][6], cube3d[4][7], cube3d[4][4], 'B', cube3d[4][8], cube3d[4][3], cube3d[4][2], cube3d[4][1]],\
[cube3d[5][7], cube3d[5][8], cube3d[5][1], cube3d[5][6], 'D', cube3d[5][2], cube3d[5][5], cube3d[5][4], cube3d[5][3]],

#符合kociemba演算法排序的方塊狀態
kociemba_cube3d = \
 str(cube3d[0][5]) + str(cube3d[0][6]) + str(cube3d[0][7]) + str(cube3d[0][4]) + 'U' + str(cube3d[0][8]) + str(cube3d[0][3]) + str(cube3d[0][2]) + str(cube3d[0][1])\
+str(cube3d[3][5]) + str(cube3d[3][6]) + str(cube3d[3][7]) + str(cube3d[3][4]) + 'R' + str(cube3d[3][8]) + str(cube3d[3][3]) + str(cube3d[3][2]) + str(cube3d[3][1])\
+str(cube3d[2][5]) + str(cube3d[2][6]) + str(cube3d[2][7]) + str(cube3d[2][4]) + 'F' + str(cube3d[2][8]) + str(cube3d[2][3]) + str(cube3d[2][2]) + str(cube3d[2][1])\
+str(cube3d[5][7]) + str(cube3d[5][8]) + str(cube3d[5][1]) + str(cube3d[5][6]) + 'D' + str(cube3d[5][2]) + str(cube3d[5][5]) + str(cube3d[5][4]) + str(cube3d[5][3])\
+str(cube3d[1][3]) + str(cube3d[1][4]) + str(cube3d[1][5]) + str(cube3d[1][2]) + 'L' + str(cube3d[1][6]) + str(cube3d[1][1]) + str(cube3d[1][8]) + str(cube3d[1][7])\
+str(cube3d[4][5]) + str(cube3d[4][6]) + str(cube3d[4][7]) + str(cube3d[4][4]) + 'B' + str(cube3d[4][8]) + str(cube3d[4][3]) + str(cube3d[4][2]) + str(cube3d[4][1])

scube = simulation_cube.Cube(simulation_cube3d)#建立物件 模擬旋轉方塊
#檢查用(中文方便辨識)
#cube3d = [['藍', '藍', '藍', '藍', '藍', '藍', '藍', '藍', '藍'], ['橘', '橘', '橘', '橘', '橘', '橘', '橘', '橘', '橘'], ['黃', '黃', '黃', '黃', '黃', '黃', '黃', '黃', '黃'], ['紅', '紅', '紅', '紅', '紅', '紅', '紅', '紅', '紅'], ['白', '白', '白', '白', '白', '白', '白', '白', '白'], ['綠', '綠', '綠', '綠', '綠', '綠', '綠', '綠', '綠']]

print()
scube.coutcube() #顯示模擬方塊狀態

print('kociemba演算法格式方塊狀態: ',kociemba_cube3d)

if kociemba_cube3d == 'UUUUUUUUURRRRRRRRRFFFFFFFFFDDDDDDDDDLLLLLLLLLBBBBBBBBB':
    print("已還原魔術方塊", file=sys.stderr)
else:

    rubikstring = "kociemba" + " " + kociemba_cube3d
    solvestep = os.popen(rubikstring).read() #透過終端機取得演算法解法
    print('kociemba演算法解法: ' + solvestep, file=sys.stderr)

    solvestep2 = solvestep.split(' ') #分割解法字串

    #機器讀取kociemba解法還原方塊，初始狀態:手臂抬起
    for i in range (len(solvestep2)):
        if i > 0: #方塊狀態定位
            for j in range (6):
                if solvestep2[i][0] == simulation_cube3d[j][4]:
                    if j == 0:
                        new_solvestep2 = 'U' + solvestep2[i][1:]
                        solvestep2[i] = new_solvestep2
                        break
                    elif j == 1:
                        new_solvestep2 = 'L' + solvestep2[i][1:]
                        solvestep2[i] = new_solvestep2
                        break
                    elif j == 2:
                        new_solvestep2 = 'F' + solvestep2[i][1:]
                        solvestep2[i] = new_solvestep2
                        break
                    elif j == 3:
                        new_solvestep2 = 'R' + solvestep2[i][1:]
                        solvestep2[i] = new_solvestep2
                        break
                    elif j == 4:
                        new_solvestep2 = 'B' + solvestep2[i][1:]
                        solvestep2[i] = new_solvestep2
                        break
                    elif j == 5:
                        new_solvestep2 = 'D' + solvestep2[i][1:]
                        solvestep2[i] = new_solvestep2
                        break
                    break

        if solvestep2[i] == "U" or solvestep2[i] == "L" or solvestep2[i] == "D" or solvestep2[i] == "U'" or solvestep2[i] == "L'" or solvestep2[i] == "D'" or solvestep2[i] == "U2" or solvestep2[i] == "L2" or solvestep2[i] == "D2":
            if arm_state == 1:
                    roboticarm_to_cube()

        if solvestep2[i] == "R" or solvestep2[i] == "F" or solvestep2[i] == "B" or solvestep2[i] == "R'" or solvestep2[i] == "F'" or solvestep2[i] == "B'" or solvestep2[i] == "R2" or solvestep2[i] == "F2" or solvestep2[i] == "B2":
            if arm_state == 0:
                    roboticarm_up()

        if solvestep2[i] == "U":
            simulation_cube.move(scube,'U',0) #模擬方塊狀態
            # A是手臂翻轉 B是底盤旋轉
            #(A)(A)(B -90)
            roll_cube_2times() #硬體運作
            basemove_reverse()

        elif solvestep2[i] == "R":
            simulation_cube.move(scube,'R',3)
            #(B +180)(A)(B -90)
            basemove_2times()
            roboticarm_to_cube()
            roll_cube()
            basemove_reverse()

        elif solvestep2[i] == "F":
            simulation_cube.move(scube,'F',2)
            #(B +90)(A)(B -90)
            basemove()
            roboticarm_to_cube()
            roll_cube()
            basemove_reverse()

        elif solvestep2[i] == "D":
            simulation_cube.move(scube,'D',5)
            #(B -90)
            basemove_reverse()

        elif solvestep2[i] == "L":
            simulation_cube.move(scube,'L',1)
            #(A)(B -90)
            roll_cube()
            basemove_reverse()

        elif solvestep2[i] == "B":
            simulation_cube.move(scube,'B',4)
            #(B -90)(A)(B -90)
            basemove_reverse()
            roboticarm_to_cube()
            roll_cube()
            basemove_reverse()

        elif solvestep2[i] == "U'":
            simulation_cube.move(scube,"U'",0)#模擬方塊狀態
            #(A)(A)(B +90)
            roll_cube_2times() #硬體運作
            basemove()

        elif solvestep2[i] == "R'":
            simulation_cube.move(scube,"R'",3)
            #(B +180)(A)(B +90)
            basemove_2times()
            roboticarm_to_cube()
            roll_cube()
            basemove()

        elif solvestep2[i] == "F'":
            simulation_cube.move(scube,"F'",2)
            #(B +90)(A)(B +90)
            basemove()
            roboticarm_to_cube()
            roll_cube()
            basemove()

        elif solvestep2[i] == "D'":
            simulation_cube.move(scube,"D'",5)
            #(B +90)
            basemove()

        elif solvestep2[i] == "L'":
            simulation_cube.move(scube,"L'",1)
            #(A)(B +90)
            roll_cube()
            basemove()

        elif solvestep2[i] == "B'":
            simulation_cube.move(scube,"B'",4)
            #(B -90)(A)(B +90)
            basemove_reverse()
            roboticarm_to_cube()
            roll_cube()
            basemove()

        elif solvestep2[i] == "U2":
            simulation_cube.move(scube,'U2',0)#模擬方塊狀態
            roll_cube_2times() #硬體運作
            basemove_reverse_2times()

        elif solvestep2[i] == "R2":
            simulation_cube.move(scube,'R2',3)
            basemove_2times()
            roboticarm_to_cube()
            roll_cube()
            basemove_reverse_2times()

        elif solvestep2[i] == "F2":
            simulation_cube.move(scube,'F2',2)
            basemove()
            roboticarm_to_cube()
            roll_cube()
            basemove_reverse_2times()

        elif solvestep2[i] == "D2":
            simulation_cube.move(scube,'D2',5)
            basemove_reverse_2times()

        elif solvestep2[i] == "L2":
            simulation_cube.move(scube,'L2',1)
            roll_cube()
            basemove_reverse_2times()

        elif solvestep2[i] == "B2":
            simulation_cube.move(scube,'B2',4)
            basemove_reverse()
            roboticarm_to_cube()
            roll_cube()
            basemove_reverse_2times()

        #print(i,solvestep2[i])

    scube.coutcube()
    roboticarm_up() #最後手臂抬起，展示還原方塊


import sys
input = sys.stdin.readline
from collections import deque

n, m = map(int, input().split())
d = {x+1:deque([0 for _ in range(20)]) for x in range(n)}
for _ in range(m):
    o = list(map(int, input().split()))
    if len(o) == 3:
        d[o[1]][o[2]-1] = (o[0]-1)^1
    elif o[0] == 3:
        d[o[1]].rotate()
        d[o[1]][0] = 0
    else:
        d[o[1]][0] = 0
        d[o[1]].rotate(-1)
log = set()
for k in d.keys():
    log.add(int(''.join(map(str, d[k])),2))
print(len(log))

백준/Silver/15787. 기차가 어둠을 헤치고 은하수를/기차가 어둠을 헤치고 은하수를.py

기차가 어둠을 헤치고 은하수를.py


import wx
import os
import builtins
import csv
import DSV
import gettext

ID_IMPORT = wx.NewIdRef()
ID_EXPORT = wx.NewIdRef()
ID_ADD = wx.NewIdRef()
ID_REMOVE = wx.NewIdRef()
ID_HELP = wx.NewIdRef()

_ = gettext.gettext

## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ##
#
# CLASS DEFIINTION
#
## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ##

class DiagramConstantsDialog(wx.Dialog):
	""" Dialogue to define constantes in a diagram.
	"""
	def __init__(self, *args, **kw):
		""" Constructor.
		"""
		super(DiagramConstantsDialog, self).__init__(*args, **kw)

		### local copy
		self.label = args[2]

		### all of the constants
		self.data = {}

		self.InitUI()

	def InitUI(self):
		""" Init the user interface.
		"""

		self.SetTitle(_("%s - Constants Manager")%(self.label))

		icon = wx.EmptyIcon() if wx.VERSION_STRING < '4.0' else wx.Icon()
		icon.CopyFromBitmap(wx.Bitmap(os.path.join(ICON_PATH_16_16, "properties.png"), wx.BITMAP_TYPE_ANY))
		self.SetIcon(icon)

		panel = wx.Panel(self)
		vbox = wx.BoxSizer(wx.VERTICAL)

		### Add a toolbar in order to add, delete, import, export and have help on the user of constant.
		tb = wx.ToolBar(panel)
		vbox.Add(tb, 0, wx.EXPAND)

		tb.SetToolBitmapSize((16,16))

		tb.AddTool(ID_ADD, "", wx.Bitmap(os.path.join(ICON_PATH,'comment_add.png')), shortHelp=_('New constant'))
		tb.AddTool(ID_REMOVE, "", wx.Bitmap(os.path.join(ICON_PATH,'comment_remove.png')), shortHelp=_('Delete constant'))
		tb.AddTool(ID_EXPORT, "", wx.Bitmap(os.path.join(ICON_PATH,'export.png')), shortHelp=_('Export constants into file'))
		tb.AddTool(ID_IMPORT, "", wx.Bitmap(os.path.join(ICON_PATH,'import.png')), wx.NullBitmap, shortHelp=_('Import constants from file'))
		tb.AddTool(ID_HELP, "", wx.Bitmap(os.path.join(ICON_PATH,'info.png')), wx.NullBitmap, shortHelp=_('Help'))

		tb.Realize()

		self._grid = wx.grid.Grid(panel, size=(300,200))
		self._grid.AutoSizeColumns(True)
		self._grid.CreateGrid(1, 2)
		self._grid.SetColLabelValue(0, _("Name"))
		self._grid.SetColSize(0, 150)
		self._grid.SetColLabelValue(1, _("Value"))
		self._grid.SetColSize(1, 150)

		### label column is not visible
		self._grid.SetRowLabelSize(0)

		vbox.Add(self._grid, 1, wx.EXPAND|wx.ALL,0)

		vbox.Add((-1, 5))

		self._button_ok = wx.Button(panel, wx.ID_OK, size=(70, 30))
		self._button_cancel = wx.Button(panel, wx.ID_CANCEL, size=(70, 30))

		hbox = wx.BoxSizer(wx.HORIZONTAL)
		hbox.Add(self._button_cancel)
		hbox.Add(self._button_ok, flag=wx.LEFT|wx.BOTTOM, border=5)

		vbox.Add(hbox, 0, flag=wx.CENTER, border=5)

#	From http://docs.wxwidgets.org/trunk/overview_windowsizing.html, wxWidgets provides two main methods for sizing:
#
#   Fit() sets the size of a window to fit around its children.
#	The size of each children is added and then this parent window changes its size to fit them all.
#   Layout() the opposite. The children will change their size, according to sizer rules, so they can fit into available space of their parent.
#	[...] is what is called by the default EVT_SIZE handler for container windows

#	Because a grid can have thousands of rows/cols its size can be huge.
#	Don't try to tell the parent to fit around it.
#	You better set max and min sizes for the grid (or its sizer) and then use Fit() or Layout() each time you change number of rows/cols or their sizes.

		panel.SetSizerAndFit(vbox)
		self.SetAutoLayout(True)
		self.Fit()

		self.__set_events()

		### just for windows
#		e = wx.SizeEvent(self.GetSize())
#		self.ProcessEvent(e)

		self.Center()

	def __set_events(self):
		""" Binding
		"""
		self.Bind(wx.EVT_TOOL, self.OnAdd, id=ID_ADD)
		self.Bind(wx.EVT_TOOL, self.OnRemove, id=ID_REMOVE)
		self.Bind(wx.EVT_TOOL, self.OnExport, id=ID_EXPORT)
		self.Bind(wx.EVT_TOOL, self.OnImport, id=ID_IMPORT)
		self.Bind(wx.EVT_TOOL, self.OnHelp, id=ID_HELP)
		self.Bind(wx.EVT_BUTTON, self.OnOk, self._button_ok)
		self.Bind(wx.EVT_BUTTON, self.OnCancel, self._button_cancel)

	def Populate(self, data):
		""" Populate the grid from a dictionary data.
		"""

		### populate the grid
		if data != {}:
			self._grid.DeleteRows(0)
			for i,key in enumerate(data):
				self._grid.AppendRows()
				self._grid.SetCellValue(i, 0, key)
				self._grid.SetCellValue(i, 1, str(data[key]))

	def OnAdd(self,evt):
		"""	Add line.
		"""
		self._grid.AppendRows()

	def OnRemove(self, evt):
		"""	Delete selected lines.
		"""

		### only if the grid has rows
		if self._grid.GetNumberRows():
			try:
				### only possible solution to have correct selected rows!
				i = self._grid.GetSelectionBlockTopLeft()[0][0]
				j = self._grid.GetSelectionBlockBottomRight()[0][0]

				if i == j:
					self._grid.DeleteRows(i)
				else:
					self._grid.DeleteRows(i,j)

			### no cells have been selected. We delete the current row according to the current position of the cursor.
			except:
				row = self._grid.GetGridCursorRow()
				self._grid.DeleteRows(row)

	def OnImport(self, event):
		""" csv file importing.
		"""

		dlg = wx.FileDialog(self, _("Choose a file"), os.getenv('USERPROFILE') or os.getenv('HOME') or HOME_PATH, "",
                                   _("CSV files (*.csv)|*.csv|Text files (*.txt)|*.txt|All files (*.*)|*.*"),
                                   wx.OPEN)
		if dlg.ShowModal() == wx.ID_OK:
			path = dlg.GetPath()
			dlg.Destroy()

			errorLog = open('import_error.log', 'a+')
			def logErrors(oldrow, newrow, expectedColumns, maxColumns, file = errorLog):
				# log the bad row to a file
				file.write(oldrow + '\n')

			dlg = DSV.ImportWizardDialog(self, wx.NewIdRef(), _('CSV Import Wizard'), path)
			if dlg.ShowModal() == wx.ID_OK:
				results = dlg.ImportData(errorHandler = logErrors)
				dlg.Destroy()
				errorLog.close()

			if results != None:

				dial = wx.MessageDialog(self, _('Do you want to clear the current values before importing?'), _('Import Manager'), wx.YES_NO | wx.NO_DEFAULT | wx.ICON_QUESTION)
				if dial.ShowModal() == wx.ID_YES:
					# delete rows
					self._grid.DeleteRows(0,self._grid.GetNumberRows())
				dial.Destroy()

				nbRows=self._grid.GetNumberRows()
				# import data to the grid
				for (row,data) in zip(list(range(len(results[1]))),results[1]):
					self._grid.AppendRows()
					self._grid.SetCellValue(row+nbRows,0,data[0])
					self._grid.SetCellValue(row+nbRows,1,str(data[1]))
				dial = wx.MessageDialog(self, _('Import completed!'), _('Import Manager'), wx.OK|wx.ICON_INFORMATION)
				dial.ShowModal()
			else:
	 			dlg.Destroy()

		else:
			dlg.Destroy()

	def OnExport(self,evt):
		"""	csv file exporting.
		"""

		wcd = _("CSV files (*.csv)|*.csv|Text files (*.txt)|*.txt|All files (*.*)|*.*")
		home = os.getenv('USERPROFILE') or os.getenv('HOME') or HOME_PATH
		export_dlg = wx.FileDialog(self, message=_('Choose a file'), defaultDir=home, defaultFile='data.csv', wildcard=wcd, style=wx.SAVE|wx.OVERWRITE_PROMPT)
		if export_dlg.ShowModal() == wx.ID_OK:
			fileName = export_dlg.GetPath()
			try:
				spamWriter = csv.writer(open(fileName, 'w'), delimiter=' ', quotechar='|', lineterminator='\n', quoting=csv.QUOTE_MINIMAL)
				for row in range(self._grid.GetNumberRows()):
					spamWriter.writerow([self._grid.GetCellValue(row,0),self._grid.GetCellValue(row,1)])

			except Exception as info:
				dlg = wx.MessageDialog(self, _('Error exporting data: %s\n'%info), _('Export Manager'), wx.OK|wx.ICON_ERROR)
				dlg.ShowModal()

			dial = wx.MessageDialog(self, _('Export completed'), _('Export Manager'), wx.OK|wx.ICON_ERROR)
			dial.ShowModal()
			export_dlg.Destroy()

	def OnOk(self,evt):
		"""	Defintion of constantes in builtin and close dialog
		"""

		for row in range(self._grid.GetNumberRows()):
			const=self._grid.GetCellValue(row,0)
			val= self._grid.GetCellValue(row,1)
			if val != '':
				if type(val) in [float, int]:
					self.data[const]=float(val)
				elif type(val) in [str, str]:
					self.data[const]=val
				else:
					pass

		if self.data != {}:
			builtins.__dict__[os.path.splitext(self.label)[0]] = self.data
		elif os.path.splitext(self.label)[0] in builtins.__dict__:
			del builtins.__dict__[os.path.splitext(self.label)[0]]

		evt.Skip()

	def GetData(self):
		""" Return the constants stored in the grid as a dictionnary.
		"""
		return self.data

	def OnCancel(self,evt):
		"""	Close dialog.
		"""
		self.Destroy()

	def OnHelp(self, event):
		""" Help message dialogue.
		"""
		dial = wx.MessageDialog(self, _("In order to use constante:\n\nCall constante by using \"Name of Diagram\"[\'Name of constante\']\n"), _('Help Manager'), wx.OK|wx.ICON_INFORMATION)
		dial.ShowModal()

## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ##
#
# MAIN PROGRAM
#
## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ## ##

### ------------------------------------------------------------
class TestApp(wx.App):
	""" Testing application
	"""

	def OnInit(self):

		builtins.__dict__['ICON_PATH_16_16']=os.path.join('icons','16x16')
		builtins.__dict__['ICON_PATH']=os.path.join('icons')
		builtins.__dict__['_'] = gettext.gettext

		frame = DiagramConstantsDialog(None, -1, "Model")
		frame.Populate({'a':1, 'b':2, 'c':3})
		frame.Show()
		return True

	def OnQuit(self, event):
		self.Close()

if __name__ == '__main__':
	app = TestApp(0)
	app.MainLoop()


prime = [True for _ in range(10000)]
for i in range(2, 100):
    if prime[i]:
        for j in range(i*2, 10000, i):
            prime[j] = False

def bfs():
    queue = [(start, 0)]
    visited[start] = True
    while queue:
        x, cnt = queue.pop(0)
        if x == target:
            return cnt
        for i in range(4):
            for j in range(10):
                dx = int(str(x)[:i] + str(j) + str(x)[i+1:])
                if not visited[dx] and prime[dx] and dx >= 1000:
                    visited[dx] = True
                    queue.append((dx, cnt+1))
    return -1

for _ in range(int(input())):
    start, target = map(int, input().split())
    visited = [False for _ in range(10000)]
    answer = bfs()
    print(answer if answer != -1 else "Impossible")


from app import db
from app.models.entities.usuario_entity import Usuario
from app.controllers.interfaces.usuario_interface import UsuarioInterface
from app.models.services.usuario_service import UsuarioService
from type import Type

class UsuarioDAO(UsuarioInterface):
    """Classe de gerenciamento UsuarioDAO"""

    @classmethod
    def cadastrar(
        cls, usuario: Type[Usuario], usuario_service: Type[UsuarioService]
    ) -> None:
        """Insere os dados em Usuário
        :param - usuario
        """

        try:
            if usuario_service.validar_usuario(usuario):
                db.session.add(usuario)
            raise Exception("Não foi possível cadastrar o usuário!")

        except Exception as e:
            db.session.rollback()

        finally:
            db.session.commit()


import argparse
import copy
import os
import os.path as osp
import time
import warnings

import mmcv
import torch
import torch.distributed as dist
from mmcv import Config, DictAction
from mmcv.runner import get_dist_info, init_dist, set_random_seed
from mmcv.utils import get_git_hash

from mmaction import __version__
from mmaction.apis import init_random_seed, train_model
from mmaction.datasets import build_dataset
from mmaction.models import build_model
from mmaction.utils import (collect_env, get_root_logger,
                            register_module_hooks, setup_multi_processes)

def parse_args():
    # 構建影片理解的輸入參數
    parser = argparse.ArgumentParser(description='Train a recognizer')
    # 指定的訓練模型config文件路徑
    parser.add_argument('config', help='train config file path')
    # 訓練過程資料保存資料夾路徑
    parser.add_argument('--work-dir', help='the dir to save logs and models')
    # 從上次訓練到的地方繼續訓練，這裡會將優化器以及學習率都進行設定
    parser.add_argument(
        '--resume-from', help='the checkpoint file to resume from')
    # 在訓練過程中是否需要進行驗證
    parser.add_argument(
        '--validate',
        action='store_true',
        help='whether to evaluate the checkpoint during training')
    # 在最後一個訓練結束後是否進行測試
    parser.add_argument(
        '--test-last',
        action='store_true',
        help='whether to test the checkpoint after training')
    # 在結束訓練後是否需要將訓練過程中最好的權重進行測試
    parser.add_argument(
        '--test-best',
        action='store_true',
        help=('whether to test the best checkpoint (if applicable) after '
              'training'))
    # 構建一個群組，在這群組當中的參數只能設定一個
    group_gpus = parser.add_mutually_exclusive_group()
    # 使用多少個gpu進行訓練
    group_gpus.add_argument(
        '--gpus',
        type=int,
        help='number of gpus to use '
        '(only applicable to non-distributed training)')
    # 可以使用的gpu在系統上是哪些index
    group_gpus.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='ids of gpus to use '
        '(only applicable to non-distributed training)')
    # 隨機種子碼，如果固定種子碼可以讓結果復現
    parser.add_argument('--seed', type=int, default=None, help='random seed')
    # 在不同的gpu上是否需要使用不同的種子碼
    parser.add_argument(
        '--diff-seed',
        action='store_true',
        help='Whether or not set different seeds for different ranks')
    # 對於CUDNN的設定，設定為True時可以讓每次模型訓練結果相同，會將每次返回的卷積算法固定
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')
    # 添加額外的config資料
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        default={},
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. For example, '
        "'--cfg-options model.backbone.depth=18 model.backbone.with_cp=True'")
    # 底層的框架選擇
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    # 分布式訓練會用到的參數
    parser.add_argument('--local_rank', type=int, default=0)
    # 將參數打包成args
    args = parser.parse_args()
    # 分布式訓練相關東西
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    # 將打包後的參數回傳
    return args

def main():
    # 獲取啟動時傳入的參數資訊
    args = parse_args()

    # 解析config文件內容
    cfg = Config.fromfile(args.config)

    # 如果有額外添加的config資訊就會到這裡進行融合
    cfg.merge_from_dict(args.cfg_options)

    # set multi-process settings
    # 設定多線程相關資料
    setup_multi_processes(cfg)

    # set cudnn_benchmark
    # 如果有需要開啟cudnn_benchmark就會開啟，會在訓練前找到比較好計算卷積的方式，如果輸入的圖像大小都固定可以提升訓練速度
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    # work_dir is determined in this priority:
    # CLI > config file > default (base filename)
    if args.work_dir is not None:
        # update configs according to CLI args if args.work_dir is not None
        # 如果有傳入work_dir就會到這裡指定訓練過程輸出
        cfg.work_dir = args.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        # 如果沒有指定輸出位址且config當中也沒有指定就會到這裡，使用默認的位置
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])
    if args.resume_from is not None:
        # 如果有傳入上次訓練到一半的保存權重檔案就會到這裡
        cfg.resume_from = args.resume_from

    if args.gpu_ids is not None or args.gpus is not None:
        # 如果有設定gpu_ids或是gpus就會到這裡
        # 官方強烈推薦使用分布式訓練，如果使用分布式訓練就不會有gpu_ids與gpus參數
        warnings.warn(
            'The Args `gpu_ids` and `gpus` are only used in non-distributed '
            'mode and we highly encourage you to use distributed mode, i.e., '
            'launch training with dist_train.sh. The two args will be '
            'deperacted.')
        if args.gpu_ids is not None:
            # 如果非分布式訓練就只能使用單gpu或是cpu進行訓練
            warnings.warn(
                'Non-distributed training can only use 1 gpu now. We will '
                'use the 1st one in gpu_ids. ')
            cfg.gpu_ids = [args.gpu_ids[0]]
        elif args.gpus is not None:
            # 如果非分布式訓練就只能使用單gpu或是cpu進行訓練
            warnings.warn('Non-distributed training can only use 1 gpu now. ')
            cfg.gpu_ids = range(1)

    # init distributed env first, since logger depends on the dist info.
    # 配置分布式訓練的環境設置
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
        _, world_size = get_dist_info()
        cfg.gpu_ids = range(world_size)

    # The flag is used to determine whether it is omnisource training
    # 設定是否使用omnisource訓練
    cfg.setdefault('omnisource', False)

    # The flag is used to register module's hooks
    # 用來註冊模型鉤子函數
    cfg.setdefault('module_hooks', [])

    # create work_dir
    # 檢查指定的work_dir是否存在，如果不存在就會進行創建
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    # dump config
    # 將設定資料與config資料進行保存，這裡就會保存到work_dir當中
    cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
    # init logger before other steps
    # 獲取當前時間戳，這裡主要是用在檔案命名上面使用
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    # 構建log的檔案名稱
    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
    # 創建logger實例對象
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    # init the meta dict to record some important information such as
    # environment info and seed, which will be logged
    # 創建meta字典，meta主要是保存一些訓練過程參數資料
    meta = dict()
    # log env info
    # 獲取當前環境相關資訊
    env_info_dict = collect_env()
    # 將環境相關資訊變成str格式
    env_info = '\n'.join([f'{k}: {v}' for k, v in env_info_dict.items()])
    # 排版用的
    dash_line = '-' * 60 + '\n'
    # 保存到logger當中
    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                dash_line)
    # 也保存到meta當中
    meta['env_info'] = env_info

    # log some basic info
    # 一些基礎的log資訊
    logger.info(f'Distributed training: {distributed}')
    logger.info(f'Config: {cfg.pretty_text}')

    # set random seeds
    # 進行種子碼設定
    seed = init_random_seed(args.seed, distributed=distributed)
    # 如果希望在不同gpu上有不同的種子碼就會在設定的種子碼上添加gpu的rank
    seed = seed + dist.get_rank() if args.diff_seed else seed
    logger.info(f'Set random seed to {seed}, '
                f'deterministic: {args.deterministic}')
    # 設定種子碼
    set_random_seed(seed, deterministic=args.deterministic)

    # 將種子碼資訊保存到config當中
    cfg.seed = seed
    # 也保存到meta當中
    meta['seed'] = seed
    meta['config_name'] = osp.basename(args.config)
    meta['work_dir'] = osp.basename(cfg.work_dir.rstrip('/\\'))

    # 構建模型實例化對象
    model = build_model(
        cfg.model,
        train_cfg=cfg.get('train_cfg'),
        test_cfg=cfg.get('test_cfg'))

    if len(cfg.module_hooks) > 0:
        # 如果有設定模型鉤子就會到這裡進行註冊
        register_module_hooks(model, cfg.module_hooks)

    if cfg.omnisource:
        # If omnisource flag is set, cfg.data.train should be a list
        # 如果有使用omnisource就會到這裡表示傳入的訓練dataset是由多個dataset組合成的
        assert isinstance(cfg.data.train, list)
        datasets = [build_dataset(dataset) for dataset in cfg.data.train]
    else:
        # 單一個訓練dataset
        datasets = [build_dataset(cfg.data.train)]

    if len(cfg.workflow) == 2:
        # For simplicity, omnisource is not compatible with val workflow,
        # we recommend you to use `--validate`
        # 如果workflow當中有指定val就會到這裡，不過官方建議使用--validate進行訓練過程中的驗證，這裡疑似是因為有部分內容沒有實作完成
        assert not cfg.omnisource
        if args.validate:
            warnings.warn('val workflow is duplicated with `--validate`, '
                          'it is recommended to use `--validate`. see '
                          'https://github.com/open-mmlab/mmaction2/pull/123')
        # 複製一份訓練的dataset資料
        val_dataset = copy.deepcopy(cfg.data.val)
        # 構建val用的dataset
        datasets.append(build_dataset(val_dataset))
    if cfg.checkpoint_config is not None:
        # save mmaction version, config file content and class names in
        # checkpoints as meta data
        # 設定checkpoint相關資訊
        cfg.checkpoint_config.meta = dict(
            mmaction_version=__version__ + get_git_hash(digits=7),
            config=cfg.pretty_text)

    # 設定測試相關的option，這裡會有是否在最後進行測試以及最後測試訓練過程中最佳的權重
    test_option = dict(test_last=args.test_last, test_best=args.test_best)
    # 開始進行訓練
    train_model(
        model,
        datasets,
        cfg,
        distributed=distributed,
        validate=args.validate,
        test=test_option,
        timestamp=timestamp,
        meta=meta)

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
from django.contrib import admin

from .models import Team, Project

@admin.register(Team)
class TeamAdmin(admin.ModelAdmin):

    list_display = (

        'title',
        'description',
        'date_created',
        'id',
        'slug',
        'published',
        'reviewed',
        'token',
    )
    list_filter = ('date_created', 'published', 'reviewed')
    raw_id_fields = ('members',)
    search_fields = ('slug',)

@admin.register(Project)
class ProjectAdmin(admin.ModelAdmin):

    list_display = (
        'title',
        'date_created',
        'slug',

        'access_token',
        'id',

        'image',
        'logo',
        'published',
        'reviewed',
    )
    list_filter = ('date_created', 'published', 'reviewed')
    search_fields = ('slug',)


from selenium import webdriver

import sys
sys.path.append('./config')
import config as config

from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.action_chains import ActionChains
from selenium.common.exceptions import TimeoutException
import time

# Replace the path with the path to your Chrome driver
driver = webdriver.Chrome('/drivers/chromedriver')

# maximize the browser window
driver.maximize_window()

# Open the website
driver.get(config.url_superadmin)

# Find the email and password fields and enter your credentials
email_field = WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.CSS_SELECTOR, "input[type='text']")))
email_field.send_keys("superadmin@gmail.com")

password_field = WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.CSS_SELECTOR, "input[type='password']")))
password_field.send_keys('Password@1234')

button = driver.find_element(By.CSS_SELECTOR, ".MuiButton-containedPrimary")
button.click()

# Wait for the page to load
time.sleep(5)

# get the current URL of the web page
current_url = driver.current_url

# check if the URL contains the expected text (in this case, "/dashboard")
if "dashboard" in current_url:
    print("Successfully redirected to dashboard page.")
else:
    print("Error: failed to redirect to dashboard page.")

# clicking  menus
wait = WebDriverWait(driver, 10)
routes_menu = wait.until(EC.element_to_be_clickable((By.XPATH, "//span[contains(text(), 'Users')]")))
routes_menu.click()

time.sleep(2)

# Find the button element by class name and text
button = driver.find_element(By.XPATH, "//button[contains(@class, 'MuiButton-root') and span[text()='+']]")
button.click()

#fill the form
input_element = driver.find_element(By.CSS_SELECTOR, "input[placeholder='Enter First Name']")
input_element.send_keys("Robin")

input_element = driver.find_element(By.CSS_SELECTOR, "input[placeholder='Enter Last Name']")
input_element.send_keys("Charm")

input_element = driver.find_element(By.CSS_SELECTOR, "input[placeholder='Enter Your Email']")
input_element.send_keys("sachintendulkar@usa.com")

input_element = driver.find_element(By.CSS_SELECTOR, "input[placeholder='1 (702) 123-4567']")
input_element.send_keys("987658742")

# #checkbox selection.
# checkbox = driver.find_element(By.CLASS_NAME, "jss1320")
# is_checked = checkbox.is_selected()

# if not is_checked:
#     checkbox.click()

#  Keep the browser open
input("Press Enter to close the browser...")

# Close the browser
driver.quit()


from collections import deque # 덱 불러옴(BFS에 이용)
import sys
input = sys.stdin.readline

# DFS
def dfs(graph, v, visited): # graph(이차원리스트), 시작점, 방문처리여부리스트

    visited[v] = True # 현재 노드를 방문처리
    print(v, end=' ') # v를 출력하는데, 띄어쓰기를 해주면서 출력한다.

    # 현재 노드와 연결된 다른 노드를 재귀적으로 방문
    for i in graph[v]:
        if not visited[i]: # 방문하지않은 노드에대해서
            dfs(graph, i, visited) # dfs를 재귀적으로 실행시킨다.

def bfs(graph, v, visited): # graph(이차원리스트), 시작점, 방문처리여부리스트
    queue = deque([v]) # 덱으로 큐를 구현한다.
    visited[v] = True

    while queue: # 큐가 빌 때까지 반복
        now = queue.popleft() # 큐에서 하나의 원소를 뽑아서 출력
        print(now, end=' ')

        for i in graph[now]: # 해당 원소와 연결된것들중에서
            if not visited[i]: # 방문하지 않은 원소들을
                queue.append(i) # 큐에 삽입
                visited[i] = True # 방문처리

n, m, v = map(int, input().split()) # 정점의 개수, 간선의 개수, 탐색시작번호
graph = [[] for _ in range(n+1)] # graph를 정점의 개수만큼 초기화(이차원리스트로)
for _ in range(m): # 간선의 개수만큼반복
    a, b = map(int, input().split()) # 정점 연결관계 받아줌
    # 양방향 간선이니깐 양방향으로 다 append
    graph[a].append(b)
    graph[b].append(a)

# 정점 낮은번호부터 먼저 방문해야 하니깐 오름차순으로 정렬
for i in range(1, n+1):
    graph[i].sort()

visited = [False] * (n+1) # 방문을 다 false로 초기화
dfs(graph, v, visited) # 그래프, v(처음 출발지점), visited(다 false로 초기화 되어있는)
print() # 중간에 개행해줘야해서
visited = [False] * (n+1) # 방문을 다 false로 초기화
bfs(graph, v, visited) # 그래프, v(처음 출발지점), visited(다 false로 초기화 되어있는)