#!/usr/bin/env python
__author__ = 'FAN Gong'
__version__ = '25 October 2022, created'
__coding__ = 'utf-8'

import os
import pandas as pd
import numpy as np
import glob
from multiprocessing import Pool, Process
import argparse

parser = argparse.ArgumentParser(usage='python3 <PATH/to/this/script>/edna.py [OPTIONS]',
                                 description='eDNA pipeline')
parser.add_argument('--raw', help='Directory contains the zipped raw reads', required=False)
parser.add_argument('--primer', help='the primers, samples txt', required=False, default='primer.txt')
parser.add_argument('--qc', help='Quality control', required=False, action='store_true')
parser.add_argument('--trim', help='trimming', required=False, action='store_true')
parser.add_argument('--assemble', help='if need to assemble', required=False, action='store_true')
parser.add_argument('--mapping', help='if need to map', required=False, action='store_true')
parser.add_argument('--zotu', help='denoise, chimera filtration and zotu clustering', required=False, action='store_true')
parser.add_argument('--otu', help='chimera filtration and otu clustering', required=False, action='store_true')
parser.add_argument('--ref', help='Path of reference for mapping', required=False, default=0)
parser.add_argument('--zblast', help='blastn zotu to database', required=False, action='store_true')
parser.add_argument('--blast', help='blastn otu to database', required=False, action='store_true')
parser.add_argument('--zfeature', help='output final zotu result', required=False, action='store_true')
parser.add_argument('--feature', help='output final otu result', required=False, action='store_true')
args = parser.parse_args()

def unzip(gzs):
    """
    Unzipping the raw reads
    :param gzs: zipped raw reads
    :return: 0
    """
    os.system(f"gunzip -c {gzs} > ../unzipped/{gzs.split('.gz')[0]}")
    return 0

def qc(sample):
    """
    Quality control
    :param sample: unzipped raw reads
    :return: 0
    """
    os.system(f"{fastp} -i {sample}.R1.fq -I {sample}.R2.fq "
              f"-o ../cleaned/{sample}.R1.fq -O ../cleaned/{sample}.R2.fq "
              f"-j {sample}.json "
              f"-h {sample}.html "
              f"-A "
              f"-q 30 "
              f"-n 5 "
              f"-x 10 "
              f"-l 15 "
              f"-w 4")
    return 0

def assembling(fq):
    """
    Merging trimmed reads
    :param fq: qced
    :return: 0
    """
    os.system(f'{flash} -M 80 -t 4 -o {fq} '
              f'../cleaned/{fq}.R1.fq '
              f'../cleaned/{fq}.R2.fq')
    os.system(f'mv {fq}.extendedFrags.fastq {fq}.fq')
    return 0

def trimming():
    """
    Trimming primer
    :return: 0
    """
    os.chdir('merged')
    os.system('for i in `ls *.fq`;do ln -s %s/merged/${i} %s/trimmed/${i};done' % (submit, submit))
    os.chdir(f'{submit}/trimmed')

    # def trim1(sample, pf):
    #     os.system(f'perl {tagcleaner} -64 '
    #               f'-fastq {sample}.R1.fq --out_format 3 -minlen 30 '
    #               f'-tag5 {pf} --verbose -out {sample}_tagcleaned.R1')
    #     return 0
    #
    # def trim2(sample, pr):
    #     os.system(f'perl {tagcleaner} -64 '
    #               f'-fastq {sample}.R2.fq --out_format 3 -minlen 30 '
    #               f'-tag5 {pr} --verbose -out {sample}_tagcleaned.R2')
    #     return 0
    base_dict = {'A': 'T',
                 'T': 'A',
                 'C': 'G',
                 'G': 'C',
                 'N': 'N'}

    if primer['primer'].str.contains('mix').sum() >= 1:
        mix = primer[primer['primer'].str.contains('mix')]
        mix.reset_index(drop=True, inplace=True)
        for i in mix.index:
            samples = mix.iat[i, 1].split(',')
            primersF = mix.iat[i, 2].split(',')
            primersR = mix.iat[i, 3].split(',')
            for j in range(len(samples)):
                for k in range(len(primersF)):
                    pR = ''.join([base_dict[i] for i in primersR[k][::-1]])
                    if k == 0:
                        os.system(f'perl {tagcleaner} -64 '
                                  f'-fastq {samples[j]}.fq --out_format 3 -minlen 30 '
                                  f'-mm5 2 -mm3 2 '
                                  f'-tag5 {primersF[k]} '
                                  f'-tag3 {pR} '
                                  f'--verbose -out {samples[j]}_{k}')
                    else:
                        os.system(f'perl {tagcleaner} -64 '
                                  f'-fastq {samples[j]}_{k-1}.fastq --out_format 3 -minlen 30 '
                                  f'-mm5 2 -mm3 2 '
                                  f'-tag5 {primersF[k]} '
                                  f'-tag3 {pR} '
                                  f'--verbose -out {samples[j]}_{k}')
                        os.system(f'rm {samples[j]}_{k - 1}.fastq')
                        if k == (len(primersF) - 1):
                            os.system(f'mv {samples[j]}_{k}.fastq {samples[j]}_tagcleaned.fastq')
                    # if k == 0:
                    #     os.system(f'perl {tagcleaner} -64 '
                    #               f'-fastq {samples[j]}.R1.fq --out_format 3 -minlen 30 '
                    #               f'-tag5 {primersF[k]} --verbose -out {samples[j]}_{k}.R1')
                    #     os.system(f'perl {tagcleaner} -64 '
                    #               f'-fastq {samples[j]}.R2.fq --out_format 3 -minlen 30 '
                    #               f'-tag5 {primersR[k]} --verbose -out {samples[j]}_{k}.R2')
                    # else:
                    #     os.system(f'perl {tagcleaner} -64 '
                    #               f'-fastq {samples[j]}_{k-1}.R1.fastq --out_format 3 -minlen 30 '
                    #               f'-tag5 {primersF[k]} --verbose -out {samples[j]}_{k}.R1')
                    #     os.system(f'perl {tagcleaner} -64 '
                    #               f'-fastq {samples[j]}_{k-1}.R2.fastq --out_format 3 -minlen 30 '
                    #               f'-tag5 {primersR[k]} --verbose -out {samples[j]}_{k}.R2')
                    #     os.system(f'rm {samples[j]}_{k - 1}.R*.fastq')
                    #     if k == (len(primersF) - 1):
                    #         os.system(f'mv {samples[j]}_{k}.R1.fastq {samples[j]}_tagcleaned.R1.fastq')
                    #         os.system(f'mv {samples[j]}_{k}.R2.fastq {samples[j]}_tagcleaned.R2.fastq')
        unmix = primer[~primer['primer'].str.contains('mix')]
        unmix.index = range(len(unmix))
        for i in unmix.index:
            samples = unmix.iat[i, 1].split(',')
            primersF = unmix.iat[i, 2]
            primersR = unmix.iat[i, 3]
            for j in samples:
                pR = ''.join([base_dict[i] for i in primersR[::-1]])
                os.system(f'perl {tagcleaner} -64 '
                          f'-fastq {j}.fq --out_format 3 -minlen 30 '
                          f'-mm5 2 -mm3 2 '
                          f'-tag5 {primersF} '
                          f'-tag3 {pR} '
                          f'--verbose -out {j}_tagcleaned')
               # did not filter the reads not matching the tag at 5'-end
               # os.popen(f'perl {tagcleaner} -64'
               #        f'-fastq {j}_R1.fq --out_format 3 -minlen 30 '
               #        f'-tag5 {primersF} --verbose -out {j}_tagcleaned_R1')
               # filter the reads not matching the tag at 5'-end
               # os.popen(f'perl {tagcleaner} -64'
               #        f'-fastq {j}_R1.fq --out_format 3 -minlen 30 '
               #        f'-tag5 {primersF} -nomatch 1 --verbose -out {j}_tagcleaned_R1')
               #  t1 = Process(target=trim1, args=(j, primersF))
               #  t2 = Process(target=trim2, args=(j, primersR))
               #  t1.start()
               #  t2.start()
               #  t1.join()
               #  t2.join()
    else:
        for i in primer.index:
            samples = primer.iat[i, 1].split(',')
            primersF = primer.iat[i, 2]
            primersR = primer.iat[i, 3]
            for j in samples:
                pR = ''.join([base_dict[i] for i in primersR[::-1]])
                os.system(f'perl {tagcleaner} -64 '
                          f'-fastq {j}.fq --out_format 3 -minlen 30 '
                          f'-mm5 2 -mm3 2 '
                          f'-tag5 {primersF} '
                          f'-tag3 {pR} '
                          f'--verbose -out {j}_tagcleaned')
                # t1 = Process(target=trim1, args=(j, primersF))
                # t2 = Process(target=trim2, args=(j, primersR))
                # t1.start()
                # t2.start()
                # t1.join()
                # t2.join()
    os.system('for i in `find ./ -type l`;do rm ${i};done')
    os.chdir(submit)
    return 0

def mapping(fq):
    """
    Trimmed reads map to a reference
    :param fq: trimmed reads
    :return: 0
    """
    os.system(f'bwa mem {ref} {fq}_tagcleaned.fastq > {fq}.sam')
    # Extracting reads which mapped to the ref
    # os.system(f'samtools view -bF 4 {fq}.sam > {fq}.bam')
    # Converting all reads
    os.system(f'samtools view -bS {fq}.sam > {fq}.bam')
    os.system(f'samtools sort {fq}.bam > {fq}.sorted.bam')
    os.system(f'samtools index {fq}.sorted.bam')
    # Statistics
    os.system(f'samtools flagstat {fq}.sorted.bam > {fq}_stat.txt')
    return 0

def usearching(fq):
    """
    Denoise, chimera filtration and zotu clustering using unoise3 algorithm.
    :param fq: trimmed reads
    :return: 0
    """
    # without pooling samples
    # Quality filter
    os.system(f'{usearch} '
              f'-fastq_filter ../trimmed/{fq}_tagcleaned.fastq '
              f'-fastq_maxee 1.0 '
              f'-fastaout {fq}_filtered.fa '
              f'-threads 4')
    # Find unique read sequences and annotate abundances
    os.system(f'{usearch} '
              f'-fastx_uniques {fq}_filtered.fa '
              f'-sizeout '
              f'-relabel Uniq '
              f'-fastaout {fq}_uniques.fa')
    if args.zotu:
        # Denoise: predict biological sequences and filter chimeras
        os.system(f'{usearch} '
                  f'-unoise3 {fq}_uniques.fa '
                  f'-zotus {fq}_zotus.fa')
        os.system(f'{usearch} '
                  f'-otutab ../trimmed/{fq}_tagcleaned.fastq '
                  f'-zotus {fq}_zotus.fa '
                  f'-otutabout {fq}_zotu.xls '
                  f'-threads 4')
                 # -mapout zmap.txt
        os.system(f"sed -i '1d' {fq}_zotu.xls && sed -i '1i #OTU_ID\t{fq}' {fq}_zotu.xls")
    if args.otu:
        # Make 97% OTUs and filter chimeras
        os.system(f'{usearch} '
                  f'-cluster_otus {fq}_uniques.fa '
                  f'-otus {fq}_otus.fa '
                  f'-relabel Otu')
        os.system(f'{usearch} '
                  f'-otutab ../trimmed/{fq}_tagcleaned.fastq '
                  f'-otus {fq}_otus.fa '
                  f'-otutabout {fq}_otu.xls '
                  f'-threads 4')
        os.system(f"sed -i '1d' {fq}_otu.xls && sed -i '1i #OTU_ID\t{fq}' {fq}_otu.xls")
    return 0

def blastn(fq):
    """
    Blast
    :param fq: samples names
    :return: 0
    """
    print(f'###### start {fq} blasting ######')
    if args.zblast:
        os.system(f'blastn '
                  f'-query {fq}_zotus.fa '
                  f'-db {mitofish} '
                  f'-task blastn '
                  f'-outfmt 6 '
                  f'-num_alignments 1 '
                  f'-num_threads 4 '
                  f'-out {fq}_aln.xls')
    if args.blast:
        os.system(f'blastn '
                  f'-query {fq}_otus.fa '
                  f'-db {mitofish} '
                  f'-task blastn '
                  f'-outfmt 6 '
                  f'-num_alignments 1 '
                  f'-num_threads 4 '
                  f'-out {fq}_aln.xls')
    os.system(f"sed -i '1i query id\tsubject id\t% identity\talignment length\t"
              f"mismatches\tgap opens\tq. start\tq. end\ts. start\ts. end\tevalue"
              f"\tbit score' {fq}_aln.xls")
    return 0

def featuring(fq):
    """
    featuring the result
    :param fq: samples names
    :return: 0
    """
    print(f'###### start {fq} featuring ######')
    try:
        if args.zfeature:
            otu_table = pd.read_table(f'{fq}_zotu.xls', sep='\t')
        if args.feature:
            otu_table = pd.read_table(f'{fq}_otu.xls', sep='\t')
       # otu_table.rename(columns={'#OTU_ID': 'first'}, inplace=True)
       # otu_table.reset_index(inplace=True)
       # otu_table.rename(columns={'index': '#OTU_ID'}, inplace=True)
        if len(otu_table.columns) > 2:
            otu_table['sum'] = otu_table.iloc[:, 1:].sum(1)
            otu_table = otu_table[['#OTU_ID', 'sum']]
            otu_table.rename(columns={'sum': f'{fq}'}, inplace=True)
        otu_aln = pd.read_table(f'{fq}_aln.xls', sep='\t')
        otu_aln = otu_aln[otu_aln['evalue'] < 1e-10]
        otu_aln = otu_aln[otu_aln['alignment length'] > 150]
        otu_aln = otu_aln.iloc[:, : 3]
        otu_aln.rename(columns={'query id': '#OTU_ID', 'subject id': 'taxonomy'}, inplace=True)
        otu_aln.sort_values(by='% identity', ascending=False, inplace=True)
        otu_aln.drop_duplicates(subset='#OTU_ID', keep='first', ignore_index=True, inplace=True)
        otu_t_a = pd.merge(otu_table, otu_aln, on='#OTU_ID')
        otu_t_a.insert(3, 'Taxonomy', np.ones(len(otu_t_a.index)))
        otu_t_a.insert(4, 'Order', np.ones(len(otu_t_a.index)))
        otu_t_a.insert(5, 'Family', np.ones(len(otu_t_a.index)))
        for i in otu_t_a.index:
            otu_t_a.iat[i, 3] = tax[tax['taxonomy'].str.contains(otu_t_a.iat[i, 2].split('|')[1])].iat[0, 1]
            otu_t_a.iat[i, 4] = tax[tax['taxonomy'].str.contains(otu_t_a.iat[i, 2].split('|')[1])].iat[0, 2]
            otu_t_a.iat[i, 5] = tax[tax['taxonomy'].str.contains(otu_t_a.iat[i, 2].split('|')[1])].iat[0, 3]
        otu_t_a.drop(labels='taxonomy', axis=1, inplace=True)
        otu_t_a.to_csv(f'{fq}_taxonomy.xls', sep='\t', index=False)
    except FileNotFoundError:
        print(f'-!- {fq} have no zotu/otu -!-')
    return 0

if __name__ == '__main__':
    fastp = '/mnt/disk2/Lab_Users/fangong/edna_programs/fastp/fastp'
    tagcleaner = '/mnt/disk2/Lab_Users/fangong/edna_programs/tagcleaner-standalone-0.16/tagcleaner.pl'
    flash = '/mnt/disk2/Lab_Users/fangong/edna_programs/FLASH-1.2.11/flash'
    usearch = '/mnt/disk2/Lab_Users/fangong/edna_programs/usearch11/usearch11'
    # makeblastdb -in complete_partial_mitogenomes.fa -dbtype nucl -parse_seqids -out mitofish
    mitofish = '/mnt/disk2/Lab_Users/fangong/edna_programs/mitofish_v3.87/mitofish_v3_87'
    submit = os.getcwd()
    primer = pd.read_table(args.primer, sep='\t', header=0, comment='#', converters={'primer': str})
    samplefq = [j for i in primer['sample'] for j in i.split(',')]
    tax = pd.read_table('/mnt/disk2/Lab_Users/fangong/edna_programs/mitofish_v3.87/taxonomy.txt')

    if args.raw:
        os.popen('mkdir unzipped')
        os.chdir(args.raw)
        gz = glob.glob('*.gz')

        gunzip = Pool(20)
        gunzip.map(unzip, gz)
        gunzip.close()
        gunzip.join()

        os.chdir(submit)

    if args.qc:
        os.popen('mkdir cleaned')
        os.chdir('unzipped')
        fqs = [i.strip() for i in os.popen("ls *.fq | sed 's/.R..fq//' | uniq").readlines()]

        fp = Pool(10)
        fp.map(qc, fqs)
        fp.close()
        fp.join()

        os.chdir(submit)
        os.system('mkdir -p intermediate/fastp/html_jason')
        os.system('mv unzipped/*.html unzipped/*.json intermediate/fastp/html_jason')

    if args.assemble:
        os.system('mkdir merged')
        os.chdir('merged')
        assembly = Pool(10)
        assembly.map(assembling, samplefq)
        assembly.close()
        assembly.join()
        os.chdir(submit)

        os.system('mkdir intermediate/flash')
        os.system('mv merged/*.fastq merged/*.his* intermediate/flash')

    if args.trim:
        os.system('mkdir trimmed')
        trimming()

    if args.mapping:
        if args.ref == 0:
            raise Exception('Reference did not specified')
        else:
            os.chdir('trimmed')
            ref = os.path.basename(args.ref)
            os.system(f'ln -s {args.ref} {ref}')
            os.system(f'bwa index {ref}')
            mapp = Pool(10)
            mapp.map(mapping, samplefq)
            mapp.close()
            mapp.join()
            if os.path.exists('mapping_result'):
                os.remove('mapping_result')
            os.mkdir('mapping_result')
            os.system('mv *.txt *.sam *bam *fasta *fasta.* *.bai mapping_result/')
            os.system(f'mv mapping_result/ {submit}')
            os.chdir(submit)

    if args.zotu or args.otu:
        if args.zotu:
            os.system('mkdir zotu')
            os.chdir('zotu')
        if args.otu:
            os.system('mkdir otu')
            os.chdir('otu')

        otu = Pool(10)
        otu.map(usearching, samplefq)
        otu.close()
        otu.join()

        os.chdir(submit)
        os.system('mkdir intermediate/usearch')
        if args.zotu:
            os.system('mv zotu/*filtered.fa zotu/*uniques.fa intermediate/usearch')
        if args.otu:
            os.system('mv otu/*filtered.fa otu/*uniques.fa intermediate/usearch')

    if args.zblast or args.blast:
        if args.zblast:
            os.chdir('zotu')
        if args.blast:
            os.chdir('otu')

        bl = Pool(10)
        bl.map(blastn, samplefq)
        bl.close()
        bl.join()

        os.chdir(submit)

    if args.zfeature or args.feature:
        if args.zfeature:
            os.chdir('zotu')
        if args.feature:
            os.chdir('otu')

        feat = Pool(10)
        feat.map(featuring, samplefq)
        feat.close()
        feat.join()

        os.system('mkdir feature')
        os.system('mv *taxonomy.xls feature/')
        os.system(f'mv feature/ {submit}')
        os.chdir(submit)


from flask import Flask, render_template, request
from src import spotify_analysis
import os

app = Flask(__name__)
spotify = spotify_analysis.SpotifyAnalysis()

@app.route("/")
def index():
    dataset = {}
    return render_template("index.html", dataset=dataset)

@app.route("/grow")
def grow():
    query = request.args["track"]
    spotify.search(query)

    uri = spotify.search_uri()
    dataset = spotify.get_results(uri)

    track = spotify.search_track()
    artist = "by " + spotify.search_artist()

    description_sections = f"""
    This arrangement consists of {spotify.number_sections} flowers, each representing a section
    in the music with its stem positioned in order of when it plays.
    The size of the flower corresponds to the loudness of the section.
    """

    description_bars = f"""
    Each of the {spotify.number_bars} petals is a musical bar, which progresses clockwise within each flower.
    The colors in each petal represent different pitches moving from the center to the tips.
    """

    return render_template(
        "index.html",
        dataset=dataset,
        track=track,
        artist=artist,
        description_sections=description_sections,
        description_bars=description_bars,
    )

if __name__ == "__main__":
    app.run(ssl_context="adhoc")


import numpy as np
import pandas as pd
from joblib import load

def load_data(filename):
    '''
    Function to load data from csv file into pandas DataFrame

    Parameters:
    filename - path to csv file

    Returns:
    DataFrame

    '''
    return pd.read_csv(filename)

def get_sex(string):
    '''
    Function to extract the sex from a string located in a pandas series.
    To be used in the pandas apply function.

    Parameters:
    string - string to be parsed

    Returns:
    string, either 'male' or 'female'

    '''
    string = str(string)
    if string.find('Male') >= 0:
        return 'male'
    if string.find('Female') >= 0:
        return 'female'
    return 'unknown'

def get_neutered(string):
    '''
    Function to extract whether an animal has been neutered or spayed
    from a string located in a pandas series.
    To be used in the pandas apply function.

    Parameters:
    string - string to be parsed

    Returns:
    string

    '''
    string = str(string)
    if string.find('Spayed') >= 0:
        return 'neutered'
    if string.find('Neutered') >= 0:
        return 'neutered'
    if string.find('Intact') >= 0:
        return 'intact'
    return 'unknown'

def calculate_age_years(age_string):
    '''
    Function to extract the age from a string located in a pandas series.
    To be used in the pandas apply function.

    Parameters:
    age_string - string to be parsed

    Returns:
    age - float value of age in years

    '''

    age_string = str(age_string)
    if age_string == 'nan':
        return 0

    age = int(age_string.split()[0])

    if age_string.find('year') > -1:
        return age
    elif age_string.find('month') > -1:
        return age / 12
    elif age_string.find('week') > -1:
        return age / 52
    elif age_string.find('day') > -1:
        return age / 365

def age_category(age):
    '''
    Function to place an animal into a certain predefined age category,
    based on its age in years.
    To be used with pandas apply function.

    Parameters:
    age - series containing ages in years

    Returns:
    string - string represting age category

    '''
    if age < 3:
        return 'young'
    elif age >= 3 and age < 5:
        return 'middle'
    elif age >= 5 and age < 10:
        return 'adult'
    elif age >= 10:
        return 'old'

def clean_data(dataframe, drop_columns):
    '''
    Function to clean data from dataframe

    Parameters:
    dataframe - dataframe to clean
    drop_columns - columns to remove (drop)

    Returns:
    None (dataframe modified in place)

    '''
    dataframe['sex'] = dataframe['SexuponOutcome'].apply(get_sex)
    dataframe['neutered'] = dataframe['SexuponOutcome'].apply(get_neutered)

    dataframe['age_years'] = dataframe['AgeuponOutcome'].apply(
        calculate_age_years)

    dataframe['age_category'] = dataframe['age_years'].apply(age_category)

    dataframe['DateTime'] = pd.to_datetime(dataframe['DateTime'])

    dataframe.drop(drop_columns, axis=1, inplace=True)

def calculate_metrics(dataframe):
    '''
    Function to calculate key metrics/KPI's related to the animal shelter.
    Used prior to being displayed with the streamlit metric function.

    Parameters:
    dataframe - contains data that metrics will be calculated on

    Returns:
    kpi_positive, kpi_returned, kpi_adopted,
        kpi_transferred, kpi_euthanized, kpi_died - metric values
    '''

    num_positive = len(
        dataframe[dataframe['OutcomeType'].isin(['Return_to_owner', 'Adoption'])])
    kpi_positive = round(num_positive / len(dataframe) * 100, 1)

    num_returned = len(
        dataframe[dataframe['OutcomeType'].isin(['Return_to_owner'])])
    kpi_returned = round(num_returned / len(dataframe) * 100, 1)

    num_adopted = len(dataframe[dataframe['OutcomeType'].isin(['Adoption'])])
    kpi_adopted = round(num_adopted / len(dataframe) * 100, 1)

    num_transferred = len(
        dataframe[dataframe['OutcomeType'].isin(['Transfer'])])
    kpi_transferred = round(num_transferred / len(dataframe) * 100, 1)

    num_euthanized = len(
        dataframe[dataframe['OutcomeType'].isin(['Euthanasia'])])
    kpi_euthanized = round(num_euthanized / len(dataframe) * 100, 1)

    num_died = len(dataframe[dataframe['OutcomeType'].isin(['Died'])])
    kpi_died = round(num_died / len(dataframe) * 100, 1)

    return kpi_positive, kpi_returned, kpi_adopted, kpi_transferred, kpi_euthanized, kpi_died

def get_month(timestamp):
    '''
    Function to get the month from a pandas timestamp.
    Used within the pandas apply function.

    Parameters:
    timestamp - pandas TimeStamp

    Returns:
    integer for the month

    '''
    return timestamp.month

def get_breed_mix(string):
    '''
    Function to determine whether the animal is a mixed breed.
    Used within the pandas apply function.

    Parameters:
    string - string within pandas series representing the breed

    Returns:
    integer - 1 for Mix, 0 for pure breed

    '''
    string = str(string)
    if string.find('Mix') >= 0 or string.find('/') >= 0:
        return 1
    else:
        return 0

def get_single_color(string):
    '''
    Function to determine whether the animal is a single color.
    Used within the pandas apply function.

    Parameters:
    string - string within pandas series representing the color

    Returns:
    integer - 1 for single color, 0 for mixed color

    '''
    string = str(string)
    if string.find('/') >= 0:
        return 0
    else:
        return 1

def data_preparation(df, drop_extra_columns):
    '''
    Function to prepare data for machine learning.

    Parameters:
    df - dataframe containing prepared data
    drop_extra_columns - list of columns to be dropped

    Returns:
    dataframe - dataframe with prepared data

    '''

    dataframe = df.copy()

    dataframe['month'] = dataframe['DateTime'].apply(get_month)

    dataframe['breed_mix'] = dataframe['Breed'].apply(get_breed_mix)

    dataframe['color_single'] = dataframe['Color'].apply(get_single_color)

    dataframe.drop(drop_extra_columns, axis=1, inplace=True)

    dataframe = pd.get_dummies(
        dataframe, columns=['AnimalType', 'sex', 'neutered'], drop_first=True)

    dataframe.drop(['sex_unknown', 'neutered_unknown'], axis=1, inplace=True)

    return dataframe

def single_animal_prediction(feature_values, model_filename):
    '''
    Function to predict the outcome for a single animal, using
    a previously trained machine learning model. Used within
    a streamlit form which gathers the feature values.

    Parameters:
    feature_values - list containing feature values
    model_filename - path to joblib model file

    Returns:
    prediction - numpy array containing prediction for the outcome
    probs - numpy array with probabilities for each possible outcome

    '''
    feature_array = np.array(feature_values).reshape(1, -1)

    model = load(model_filename)

    prediction = model.predict(feature_array)

    probs = model.predict_proba(feature_array)

    return prediction, probs


from __future__ import print_function
#%matplotlib inline
import os
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML

from discriminator import Discriminator
from generator import Generator
from minimizer import Minimizer

def load_images(plot_sample=False, device='cpu'):
    dataset = dset.ImageFolder(root=dataroot,
                               transform=transforms.Compose([
                                   transforms.Resize(image_size),
                                   transforms.CenterCrop(image_size),
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]))
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=workers)

    if plot_sample:
        real_batch = next(iter(dataloader))
        plt.figure(figsize=(8, 8))
        plt.axis("off")
        plt.title("Training Images")
        plt.imshow(np.transpose(
            vutils.make_grid(real_batch[0].to(device)[:64], padding=2, normalize=True).cpu(), (1, 2, 0)))

    return dataloader

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)

def create_generator():
    netG = Generator(ngpu, nz, ngf, nc).to(device)

    # Handle multi-gpu if desired
    if (device.type == 'cuda') and (ngpu > 1):
        netG = nn.DataParallel(netG, list(range(ngpu)))

    netG.apply(weights_init)
    return netG

def create_discriminator():
    netD = Discriminator(ngpu, nc, ndf).to(device)

    # Handle multi-gpu if desired
    if (device.type == 'cuda') and (ngpu > 1):
        netD = nn.DataParallel(netD, list(range(ngpu)))

    netD.apply(weights_init)
    return netD

def create_optimizers(netD, netG):
    optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))
    optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))
    return optimizerD, optimizerG

if __name__ == '__main__':
    workers = 2  # Number of workers for data-loader
    batch_size = 128  # Batch size during training
    image_size = 64  # Spatial size of training images. All images will be resized to this size using a transformer.
    nc = 3  # Number of channels in the training images. For color images this is 3
    nz = 100  # Size of z latent vector (i.e. size of generator input)
    ngf = 64  # Size of feature maps in generator
    ndf = 64  # Size of feature maps in discriminator
    num_epochs = 5  # Number of training epochs
    lr = 0.0002  # Learning rate for optimizers
    beta1 = 0.5  # Beta1 hyperparam for Adam optimizers
    ngpu = 1
    manualSeed = 999  # random.randint(1, 10000) # use if you want new results
    print("Random Seed: ", manualSeed)
    random.seed(manualSeed)
    torch.manual_seed(manualSeed)

    dataroot = os.path.join(os.path.dirname(os.getcwd()), 'celeba')  # Storing the data-set outside of pycharm's reach.
    device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
    dataloader = load_images(True)

    netG = create_generator()
    netD = create_discriminator()
    optimizerD, optimizerG = create_optimizers(netD, netG)

    criterion = nn.BCELoss()
    fixed_noise = torch.randn(64, nz, 1, 1, device=device)  # Latent vectors to visualize the progress of the generator

    real_label = 1
    fake_label = 0

    D_losses, G_losses, img_list = Minimizer.train(
        dataloader,
        netD,
        netG,
        optimizerD,
        optimizerG,
        criterion,
        num_epochs,
        device,
        nz,
        fixed_noise)

    plt.figure(figsize=(10, 5))
    plt.title("Generator and Discriminator Loss During Training")
    plt.plot(G_losses, label="G")
    plt.plot(D_losses, label="D")
    plt.xlabel("iterations")
    plt.ylabel("Loss")
    plt.legend()
    plt.show()

    # %%capture
    fig = plt.figure(figsize=(8, 8))
    plt.axis("off")
    ims = [[plt.imshow(np.transpose(i, (1, 2, 0)), animated=True)] for i in img_list]
    ani = animation.ArtistAnimation(fig, ims, interval=1000, repeat_delay=1000, blit=True)

    HTML(ani.to_jshtml())


from pydantic import BaseModel, Extra, Field

class Metric(BaseModel):
    label: str = Field(description="name of the judging parameter")
    score: int = Field(description="assigned score of the judging parameter")
    reason: str = Field(
        description="justification of your assigned score for the judging parameter"
    )

    class Config:
        extra = Extra.ignore

class MetricList(BaseModel):
    metrics: list[Metric] = Field(description="List of all the judging parameters")

    class Config:
        extra = Extra.ignore

class SimilarTools(BaseModel):
    toolname: str = Field(description="name of the tool that is related to the project")
    reason: str = Field(
        description="explain how the tool is similar and dissimilar to the project"
    )

    class Config:
        extra = Extra.ignore

class SimilarToolsList(BaseModel):
    listoftools: list[SimilarTools] = Field(description="list of all the similar tools")

    class Config:
        extra = Extra.ignore


# Copyright 2010 United States Government as represented by the
# Administrator of the National Aeronautics and Space Administration.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

from oslo_log import log as logging

from conveyor import compute
from conveyor import exception
from conveyor import network
from conveyor.resource.driver import base
from conveyor.resource.driver import instances
from conveyor.resource.driver import networks
from conveyor.resource import resource

LOG = logging.getLogger(__name__)

class LoadbalanceVip(base.Resource):

    def __init__(self, context, collected_resources=None,
                 collected_parameters=None, collected_dependencies=None):
        self.context = context
        self.neutron_api = network.API()
        self._collected_resources = collected_resources or {}
        self._collected_parameters = collected_parameters or {}
        self._collected_dependencies = collected_dependencies or {}

    def extract_loadbalanceVips(self, vip_ids):
        if not vip_ids:
            _msg = 'Create LB vip resource error: id is null.'
            LOG.error(_msg)
            raise exception.InvalidInput(reason=_msg)

        try:
            for vip_id in vip_ids:
                self.extract_loadbalanceVip(vip_id)
        except exception.ResourceExtractFailed:
            raise
        except Exception as e:
            _msg = 'Create LB vip resource error: %s' % e
            LOG.error(_msg)
            raise exception.ResourceExtractFailed(_msg)

    def extract_loadbalanceVip(self, vip_id, pool_name):

        # if vip resource exist in collect resource
        vip_col = self._collected_resources.get(vip_id)
        if vip_col:
            return vip_col
        dependences = []
        properties = {}

        # 1. query vip info
        try:
            vip_info = self.neutron_api.get_vip(self.context, vip_id)
        except Exception as e:
            _msg = 'Create LB vip resource error: %s' % e
            LOG.error(_msg)
            raise exception.ResourceExtractFailed(reason=_msg)

        properties['protocol_port'] = vip_info.get('protocol_port')
        properties['session_persistence'] = vip_info.\
            get('session_persistence')
        # properties['protocol'] = vip_info.get('protocol')
        vip_name = vip_info.get('name')
        vip_address = vip_info.get('address')
        vip_admin = vip_info.get('admin_state_up')
        vip_connection = vip_info.get('connection_limit')
        vip_des = vip_info.get('description')
        if vip_name:
            properties['name'] = vip_name

        if vip_address:
            properties['address'] = vip_address

        if vip_admin:
            properties['admin_state_up'] = vip_admin
        if vip_connection:
            properties['connection_limit'] = vip_connection
        if vip_des:
            properties['description'] = vip_des

        properties['pool_id'] = {'get_resource': pool_name}
        dependences.append(pool_name)

        # build subnet resource and build dependence relation
        subnet_id = vip_info.get('subnet_id')

        if subnet_id:
            newtwork_driver = \
                networks.NetworkResource(
                    self.context,
                    collected_resources=self._collected_resources,
                    collected_parameters=self._collected_parameters,
                    collected_dependencies=self._collected_dependencies)
            subnet_ids = []
            subnet_ids.append(subnet_id)
            subnet_res = newtwork_driver.extract_subnets(subnet_ids)
            dependences.append(subnet_res[0].name)
            properties['subnet'] = {'get_resource': subnet_res[0].name}

            # update collect resources and dependences
            self._collected_resources = \
                newtwork_driver.get_collected_resources()

            self._collected_dependencies = \
                newtwork_driver.get_collected_dependencies()

        vip_type = "OS::Neutron::Vip"
        vip_name = 'loadbalanceVip_%d' % self._get_resource_num(vip_type)

        # build listener resource and build dependence relation
        listener_extras = vip_info.get('extra_listeners')
        listener_ids = None
        if listener_extras:
            listener_ids = []
            for listener_extra in listener_extras:
                listener_id = listener_extra.get('id')
                listener_ids.append(listener_id)

        if listener_ids:
            listener_driver = \
                LoadbalanceListener(
                    self.context,
                    collected_resources=self._collected_resources,
                    collected_parameters=self._collected_parameters,
                    collected_dependencies=self._collected_dependencies)
            listener_driver.extract_loadbalanceListeners(listener_ids,
                                                         vip_id,
                                                         vip_name)

            # update collect resources and dependences
            self._collected_resources = \
                listener_driver.get_collected_resources()

            self._collected_dependencies = \
                listener_driver.get_collected_dependencies()

        vip_res = resource.Resource(vip_name, vip_type,
                                    vip_id, properties=properties)

        # remove duplicate dependencies
        dependencies = {}.fromkeys(dependences).keys()
        vip_dep = resource.ResourceDependency(vip_id, vip_name, '',
                                              vip_type,
                                              dependencies=dependencies)

        self._collected_resources[vip_id] = vip_res
        self._collected_dependencies[vip_id] = vip_dep

        return vip_res

class LoadbalancePool(base.Resource):

    def __init__(self, context, collected_resources=None,
                 collected_parameters=None, collected_dependencies=None):
        self.context = context
        self.neutron_api = network.API()
        self._collected_resources = collected_resources or {}
        self._collected_parameters = collected_parameters or {}
        self._collected_dependencies = collected_dependencies or {}

    def extract_loadbalancePools(self, pool_ids):

        if not pool_ids:
            _msg = 'Create LB pool resource error: id is null.'
            LOG.error(_msg)
            raise exception.InvalidInput(reason=_msg)

        try:
            for pool_id in pool_ids:
                self.extract_loadbalancePool(pool_id)
        except exception.ResourceExtractFailed:
            raise
        except Exception as e:
            _msg = 'Create LB pool resource error: %s' % e
            LOG.error(_msg)
            raise exception.ResourceExtractFailed(_msg)

    def extract_loadbalancePool(self, pool_id):

        pool_col = self._collected_resources.get(pool_id)

        if pool_col:
            return pool_col

        dependencies = []
        properties = {}

        # 1. query pool info
        try:
            pool = self.neutron_api.show_pool(self.context, pool_id)
        except Exception as e:
            LOG.error('Create LB pool %(pool)s resource error %(error)s',
                      {'pool': pool_id, 'error': e})
            _msg = 'Create LB pool resource error: %s' % e
            raise exception.ResourceExtractFailed(_msg)

        pool_info = pool.get('pool')

        properties['lb_method'] = pool_info.get('lb_method')
        properties['protocol'] = pool_info.get('protocol')
        pool_name = pool_info.get('name')
        pool_des = pool_info.get('description')
        pool_admin = pool_info.get('admin_state_up')
        if pool_name:
            properties['name'] = pool_name
        if pool_des:
            properties['description'] = pool_des
        if pool_admin:
            properties['admin_state_up'] = pool_admin

        # properties['provider'] = pool_info.get('provider')
        pool_type = "OS::Neutron::Pool"
        pool_name = 'loadbalancePool_%d' % self._get_resource_num(pool_type)

        # 2. build vip of pool and build dependence relation
        vip_id = pool_info.get('vip_id')
        if vip_id:
            vip_driver = \
                LoadbalanceVip(
                    self.context,
                    collected_resources=self._collected_resources,
                    collected_parameters=self._collected_parameters,
                    collected_dependencies=self._collected_dependencies)
            vip_res = vip_driver.extract_loadbalanceVip(vip_id, pool_name)
            properties['vip'] = {'get_resource': vip_res.name}
            # vip resource in lb resource list
            self._collected_resources = vip_driver.get_collected_resources()
            self._collected_dependencies = \
                vip_driver.get_collected_dependencies()

        subnet_id = pool_info.get('subnet_id')

        if subnet_id:
            # 3. build subnet resource and build dependence relation
            newtwork_driver = \
                networks.NetworkResource(
                    self.context,
                    collected_resources=self._collected_resources,
                    collected_parameters=self._collected_parameters,
                    collected_dependencies=self._collected_dependencies)
            subnet_ids = []
            subnet_ids.append(subnet_id)
            subnet_res = newtwork_driver.extract_subnets(subnet_ids)
            dependencies.append(subnet_res[0].name)
            properties['subnet'] = {'get_resource': subnet_res[0].name}

            # 3.2 add subnet resource in lb resource list
            self._collected_resources = \
                newtwork_driver.get_collected_resources()
            self._collected_dependencies = \
                newtwork_driver.get_collected_dependencies()

        # 4. build members of pool and build dependence relation
        member_ids = pool_info.get('members')

        if member_ids:
            lb_member_driver = \
                LoadbalanceMember(
                    self.context,
                    collected_resources=self._collected_resources,
                    collected_parameters=self._collected_parameters,
                    collected_dependencies=self._collected_dependencies)
            lb_member_driver.extract_loadbalanceMembers(member_ids, pool_name)

            # update collect resource
            self._collected_resources = \
                lb_member_driver.get_collected_resources()
            self._collected_dependencies = \
                lb_member_driver.get_collected_dependencies()

        # 4. query health monitor of pool and build dependence relation
        healthmonitor_ids = pool_info.get('health_monitors')

        if healthmonitor_ids:
            lb_healthmonitor_driver = \
                LoadbalanceHealthmonitor(
                    self.context,
                    collected_resources=self._collected_resources,
                    collected_parameters=self._collected_parameters,
                    collected_dependencies=self._collected_dependencies)
            ids = healthmonitor_ids
            h = lb_healthmonitor_driver.extract_loadbalanceHealthmonitors(ids)

            # update collect resource
            self._collected_resources = \
                lb_healthmonitor_driver.get_collected_resources()
            self._collected_dependencies = \
                lb_healthmonitor_driver.get_collected_dependencies()

            # add all healthmonitor to pool dependences
            monitors = []
            for res in h:
                dependencies.append(res.name)
                monitors.append({'get_resource': res.name})

            properties['monitors'] = monitors

        pool_type = "OS::Neutron::Pool"
        pool_name = 'loadbalancePool_%d' % self._get_resource_num(pool_type)

        pool_res = resource.Resource(pool_name, pool_type,
                                     pool_id, properties=properties)

        # remove duplicate dependencies
        dependencies = {}.fromkeys(dependencies).keys()
        pool_dep = resource.ResourceDependency(pool_id, pool_name, '',
                                               pool_type,
                                               dependencies=dependencies)

        self._collected_resources[pool_id] = pool_res
        self._collected_dependencies[pool_id] = pool_dep

class LoadbalanceListener(base.Resource):

    def __init__(self, context, collected_resources=None,
                 collected_parameters=None, collected_dependencies=None):
        self.context = context
        self.neutron_api = network.API()
        self._collected_resources = collected_resources or {}
        self._collected_parameters = collected_parameters or {}
        self._collected_dependencies = collected_dependencies or {}

    def extract_loadbalanceListeners(self, listener_ids, vip_id, vip_name):
        if not listener_ids:
            _msg = 'Create LB listener resource error: id is null.'
            LOG.error(_msg)
            raise exception.InvalidInput(reason=_msg)

        try:
            for listener_id in listener_ids:
                self.extract_loadbalanceListener(listener_id, vip_id,
                                                 vip_name)
        except exception.ResourceExtractFailed:
            raise
        except Exception as e:
            _msg = 'Create LB listener resource error: %s' % e
            LOG.error(_msg)
            raise exception.ResourceExtractFailed(_msg)

    def extract_loadbalanceListener(self, listener_id, vip_id, vip_name):

        listener_col = self._collected_resources.get(listener_id)

        if listener_col:
            return listener_col

        # 1 query listener info
        try:
            listener = self.neutron_api.show_listener(self.context,
                                                      listener_id,
                                                      vip_id)
        except Exception as e:
            _msg = 'Create LB listener resource error: %s' % e
            LOG.error(_msg)
            raise exception.ResourceExtractFailed(reason=_msg)

        # 2 build listener resource and dependence
        listener_info = listener.get('listener')

        properties = {}
        dependencies = []
        properties['protocol'] = listener_info.get('protocol')
        properties['protocol_port'] = listener_info.get('protocol_port')
        properties['vip_id'] = {'get_resource': vip_name}
        dependencies.append(vip_name)

        listener_type = "OS::Neutron::Listener"
        listener_name = 'loadbalanceListener_%d' % \
            self._get_resource_num(listener_type)

        listener_res = resource.Resource(listener_name, listener_type,
                                         listener_id,
                                         properties=properties)

        # remove duplicate dependencies
        dependencies = {}.fromkeys(dependencies).keys()
        listener_dep = resource.ResourceDependency(listener_id, listener_name,
                                                   '',
                                                   listener_type,
                                                   dependencies=dependencies)

        self._collected_resources[listener_id] = listener_res
        self._collected_dependencies[listener_id] = listener_dep

        return listener_res

class LoadbalanceMember(base.Resource):

    def __init__(self, context, collected_resources=None,
                 collected_parameters=None, collected_dependencies=None):
        self.context = context
        self.neutron_api = network.API()
        self.nova_api = compute.API()
        self._collected_resources = collected_resources or {}
        self._collected_parameters = collected_parameters or {}
        self._collected_dependencies = collected_dependencies or {}

    def extract_loadbalanceMember(self, member_id, pool_name):
        # check resource exist or not
        member_col = self._collected_resources.get(member_id)

        if member_col:
            return member_col

        properties = {}
        dependencies = []

        # query member info

        try:
            member = self.neutron_api.show_member(self.context, member_id)
        except Exception as e:
            _msg = 'Create LB member resource error: %s' % e
            LOG.error(_msg)
            raise exception.ResourceExtractFailed(reason=_msg)

        member_info = member.get('member')
        properties['address'] = member_info.get('address')
        properties['protocol_port'] = member_info.get('protocol_port')
        member_admin = member_info.get('admin_state_up')
        member_weight = member_info.get('weight')
        if member_admin:
            properties['admin_state_up'] = member_admin
        if member_weight:
            properties['weight'] = member_weight
        properties['pool_id'] = {'get_resource': pool_name}
        dependencies.append(pool_name)

        # if member relates to instances
        try:
            server_id = self._get_member_related_vm(properties['address'])

            if server_id:
                instance_driver = instances.InstanceResource(
                    self.context,
                    collected_resources=self._collected_resources,
                    collected_parameters=self._collected_parameters,
                    collected_dependencies=self._collected_dependencies)
                instance_ids = []
                instance_ids.append(server_id)
                reses = instance_driver.extract_instances(instance_ids)
                if reses:
                    instance_res_name = reses[0].name
                    dependencies.append(instance_res_name)

        except Exception as e:
            _msg = 'Create LB member resource error: %s' % e
            LOG.error(_msg)
            raise exception.ResourceExtractFailed(reason=_msg)

        # build member resource and build dependence relation
        member_type = "OS::Neutron::PoolMember"
        member_name = 'loadbalancePoolMember_%d' % \
            self._get_resource_num(member_type)

        member_res = resource.Resource(member_name, member_type,
                                       member_id, properties=properties)

        # remove duplicate dependencies
        dependencies = {}.fromkeys(dependencies).keys()
        member_dep = resource.ResourceDependency(member_id,
                                                 member_name,
                                                 '',
                                                 member_type,
                                                 dependencies=dependencies)

        self._collected_resources[member_id] = member_res
        self._collected_dependencies[member_id] = member_dep

        return member_res

    def extract_loadbalanceMembers(self, member_ids, pool_name):

        if not member_ids or not pool_name:
            LOG.error('Create LB member resource error: %s', pool_name)
            _msg = 'Create LB member resource error: member or pool is null.'
            raise exception.InvalidInput(reason=_msg)
        try:
            for member_id in member_ids:
                self.extract_loadbalanceMember(member_id, pool_name)
        except Exception as e:
            _msg = 'Create LB member resource error: %s' % e
            LOG.error(_msg)
            raise exception.ResourceExtractFailed(reason=_msg)

    def _get_member_related_vm(self, mem_address):

        # 1. query all vm
        servers = self.nova_api.get_all_servers(self.context)

        # 2. find the same address with mem_address of vm
        for server in servers:
            addresses = server.get('addresses', '')
            if not addresses:
                continue
            for k, addrs in addresses.items():
                if not addrs:
                    continue
                for addr in addrs:
                    ip_address = addr.get('addr', None)
                    if mem_address == ip_address:
                        return server.get('id')
        return None

class LoadbalanceHealthmonitor(base.Resource):

    def __init__(self, context, collected_resources=None,
                 collected_parameters=None, collected_dependencies=None):
        self.context = context
        self.neutron_api = network.API()
        self._collected_resources = collected_resources or {}
        self._collected_parameters = collected_parameters or {}
        self._collected_dependencies = collected_dependencies or {}

    def extract_loadbalanceHealthmonitors(self, healthmonitor_ids):

        if not healthmonitor_ids:
            _msg = 'Create LB health monitor resource error: id is null.'
            LOG.error(_msg)
            raise exception.InvalidInput(reason=_msg)

        res_list = []
        try:
            for healthmonitor_id in healthmonitor_ids:
                res = self.extract_loadbalanceHealthmonitor(healthmonitor_id)
                res_list.append(res)
        except exception.ResourceExtractFailed:
            raise
        except Exception as e:
            _msg = 'Create LB health monitor resource error: %s' % e
            LOG.error(_msg)
            raise exception.ResourceExtractFailed(reason=_msg)

        return res_list

    def extract_loadbalanceHealthmonitor(self, healthmonitor_id):

        # check health monitor resource exist or not
        healthmonitor_col = self._collected_resources.get(healthmonitor_id)
        if healthmonitor_col:
            return healthmonitor_col

        try:
            healthmonitor = \
                self.neutron_api.show_health_monitor(self.context,
                                                     healthmonitor_id)
        except Exception as e:
            _msg = 'Create LB health monitor resource error: %s' % e
            LOG.error(_msg)
            raise exception.ResourceExtractFailed(reason=_msg)

        properties = {}

        healthmonitor_info = healthmonitor.get('health_monitor')

        properties['delay'] = healthmonitor_info.get('delay')
        properties['type'] = healthmonitor_info.get('type')
        properties['max_retries'] = healthmonitor_info.get('max_retries')
        properties['timeout'] = healthmonitor_info.get('timeout')

        healthmonitor_code = healthmonitor_info.get('expected_codes')
        healthmonitor_http = healthmonitor_info.get('http_method')
        healthmonitor_admin = healthmonitor_info.get('admin_state_up')
        healthmonitor_url = healthmonitor_info.get('url_path')
        if healthmonitor_code:
            properties['expected_codes'] = healthmonitor_code

        if healthmonitor_http:
            properties['http_method'] = healthmonitor_http

        if healthmonitor_admin:
            properties['admin_state_up'] = healthmonitor_admin

        if healthmonitor_url:
            properties['url_path'] = healthmonitor_url

        healthmonitor_type = "OS::Neutron::HealthMonitor"
        healthmonitor_name = 'loadbalanceHealthmonitor_%d' \
                             % self._get_resource_num(healthmonitor_type)

        healthmonitor_res = resource.Resource(healthmonitor_name,
                                              healthmonitor_type,
                                              healthmonitor_id,
                                              properties=properties)

        healthmonitor_dep = resource.ResourceDependency(healthmonitor_id,
                                                        healthmonitor_name,
                                                        '',
                                                        healthmonitor_type)

        self._collected_resources[healthmonitor_id] = healthmonitor_res
        self._collected_dependencies[healthmonitor_id] = healthmonitor_dep

        return healthmonitor_res


from prefect.deployments import Deployment
from prefect.filesystems import GitHub

github_block = GitHub.load("github-block")

github_dep = Deployment.build_from_flow(
    flow=parent_flow,
    name="github-flow",
    storage = github_block
)

if __name__ == "__main__":
    github_dep.apply()


import sys
class Graph:
    def __init__(self, ver_no):
        self.v_count = ver_no
        self.graph = []

    def find_path(self, cur_vertex, parent):
        if parent[cur_vertex] == - 1:
            print (cur_vertex, end='')
            return
        self.find_path(parent[cur_vertex], parent)
        print (cur_vertex, end='')

    def print_vertex_and_distance(self, distance, parent):
        print("Vertex \t\tDistance from Source\tPath")
        for i in range(self.v_count):
            print("\n%d --> %d \t\t%d \t\t\t\t\t" %(0, i, distance[i]), end=''), self.find_path(i, parent)

    def select_min_u_vertex(self, selected_status, min_distance):

        current_min_value = sys.maxsize
        for vertex in range(self.v_count):
            if selected_status[vertex] == False and min_distance[vertex] < current_min_value:
                current_min_value = min_distance[vertex]
                vertex_with_min_dist_index = vertex

        return vertex_with_min_dist_index

    def dijkstra(self, src):
        selected_status =[False] * self.v_count
        min_distance = [sys.maxsize] * self.v_count
        min_distance[src] = 0
        parent = [-1] * self.v_count # simple and effective

        while False in selected_status:
            # initially only src will be selected
            u_vertex = self.select_min_u_vertex(selected_status, min_distance)

            selected_status[u_vertex] = True

            for v_vertex in range(self.v_count):
                if self.graph[u_vertex][v_vertex] > 0 and selected_status[v_vertex] == False and min_distance[u_vertex] + self.graph[u_vertex][v_vertex] < min_distance[v_vertex]:
                    min_distance[v_vertex] = min_distance[u_vertex] + self.graph[u_vertex][v_vertex]
                    parent[v_vertex] = u_vertex

        self.print_vertex_and_distance(min_distance, parent)
        # print(parent)

g = Graph(9)
g.graph = [[0, 4, 0, 0, 0, 0, 0, 8, 0],
		[4, 0, 8, 0, 0, 0, 0, 11, 0],
		[0, 8, 0, 7, 0, 4, 0, 0, 2],
		[0, 0, 7, 0, 9, 14, 0, 0, 0],
		[0, 0, 0, 9, 0, 10, 0, 0, 0],
		[0, 0, 4, 14, 10, 0, 2, 0, 0],
		[0, 0, 0, 0, 0, 2, 0, 1, 6],
		[8, 11, 0, 0, 0, 0, 1, 0, 7],
		[0, 0, 2, 0, 0, 0, 6, 7, 0]
		]

g.dijkstra(0)


import numpy as np
import math
import sys,os
sys.path.append('../../')
import pickle as pckl
from matplotlib import pyplot
#from tools.mkm_creator import *
#from tools_for_analysis.mkm_creator import *
from tools_for_analysis.FED_tools import plot_diagram,read_calculated_data
from intermediates_dict import ads_and_electron

#inputfile = 'C2model_h2o.pkl'
show = True
products=['C2H4','CH3CH2OH','CH4','H2']

pH = 13
SHE_potential=-0.77
show_reaction_steps=False
RHE_potential = SHE_potential+0.059*pH
c2h4color='b'
ethacolor='peru'
bothc='g'

ylabel='$\Delta$G$^\phi$ / eV'

xlabel='N$_\mathrm{H}$ + $\gamma$'
inputfile=None

figsize=(12,9)

alldata=read_calculated_data(inputfile,'100',1,pklfile='../../parsed_data_for_analysis/parsed_data.pkl',indict=ads_and_electron)
plot_diagram(ads_and_electron,'100',
                show_reaction_steps=show_reaction_steps,
                RHE_potential=RHE_potential,pH=pH,add_field=False,
                ylabel=ylabel,xlabel=xlabel,show=show,yrange=[-1.3,1.5],size=(16,6),
                clean_dict_from_dead_ends=True,add_CO_to_C1s=True,add_title=False,
                name='Figure1.pdf',
                highlights={'HCCHO':bothc, 'OHCCH2':bothc,'CO':bothc,'OCCO':bothc,'OCCOH':bothc,'CCO':bothc,'HCCO':bothc,
                    'OCHCH3':ethacolor,'CH3CH2OH_g':ethacolor,'H3CCH2O':ethacolor,
                    'H2CCH2O':{'self':c2h4color,'O':c2h4color},'O':c2h4color,'OH':c2h4color,'C2H4_g':c2h4color},
                add_potential_response_in_x=True,state_width=0.25)


from pprint import pformat
from io import BytesIO

import core
import settings

PLUGINVERSION = 2
# Always name this variable as `plugin`
# If you dont, module loader will fail to load the plugin!
plugin = core.Plugin()

@plugin.command(command="//msgdump",
                hidden=True)
def dump(bot, update, user, args):
    text = pformat(update.message.to_dict())
    text_obj = BytesIO()
    text_obj.name = "%s.txt" % update.update_id
    text_obj.write(bytes(text, "utf-8"))
    text_obj.seek(0)
    return core.message(file=text_obj)

@plugin.command(command="//delmsg",
                inline_supported=False,
                hidden=True)
def msgdel(bot, update, user, args):
    if user.id == settings.ADMIN:
        update.message.reply_to_message.delete()

@plugin.command(command="//exec",
                hidden=True)
def docode(bot, update, user, args):
    if user.id == settings.ADMIN:
        return core.message(eval(" ".join(args)))

@plugin.command(command="//plugins", hidden=True)
def coreplug_list(bot, *_):
    message = []
    for plugin in bot.modloader.plugins:
        txt = ''
        if plugin.state == core.constants.OK:
            txt += "✅"
        else:
            txt += "⛔"
        txt += plugin.name
        message.append(txt)
    message = sorted(message)
    message.reverse()
    return core.message("\n".join(message))

@plugin.command(command="//sideload", hidden=True)
def coreplug_load(self, bot, update, user, args):
    args = " ".join(args)
    if user.id == settings.ADMIN:
        self.logger.info("Reload requested.")
        update.message.reply_text("Loading " + args)
        self.load_plugin(args)
        return self.coreplug_list()
    else:
        return core.message("Access Denied.")


def powdgt(a,b):
    if type(a)==int and type (b)==int:
        r=a**b
        #to find the sum of all digits in the result:
        st=str(r)
        l=0
        for i in st:
            l+=int(i)
    return l
inp=input('Enter the base and the power separated by comma:')
x,y=inp.split(',')
x=int(x)
y=int(y)
print('The sum of digits of',x,'power',y,'is',powdgt(x,y))
#'The sum of digits of {0} power {1} is {2}'.format(x,y,powdgt(x,y))


n, m = map(int, input().split())
AB = [list(map(int, input().split())) for _ in range(m)]
edges = [[] for _ in range(n)]
for a, b in AB:
    edges[a - 1].append(b)
    edges[b - 1].append(a)
print(edges)

cnt = 0
ans = [-1] * n
for idx, edge in enumerate(edges):
    ans[idx] = 0
    for e in set(edge):
        if ans[e - 1] == -1:
            cnt += 1
            ans[e - 1] = 0
            print(ans, cnt)

print(cnt)
##############
#WA
n, m = map(int, input().split())
AB = [list(map(int, input().split())) for _ in range(m)]
edges = [[] for _ in range(n)]
for a, b in AB:
    edges[a - 1].append(b)
    edges[b - 1].append(a)
print(edges)

ans = 0
print(edges)
for edge in edges:
    ans = max(ans, len(set(edge)+1))

print(ans)

########################
n, m = map(int, input().split())
AB = [list(map(int, input().split())) for _ in range(m)]
edges = [[] for _ in range(n)]
for a, b in AB:
    edges[a - 1].append(b)
    edges[b - 1].append(a)

ans = 0
for edge in edges:
    print(len(set(edge)))
    ans = max(ans, len(set(edge))+1)

print(ans)
############################
# Union find
# かっつぱさんの解説と公式のスライドみた
# https://www.youtube.com/watch?time_continue=916&v=zxor0DdwoXA&feature=emb_logo
# https://www.slideshare.net/chokudai/union-find-49066733
# ユニオン
import sys

sys.setrecursionlimit(10 ** 9)

n, m = map(int, input().split())
AB = [list(map(int, input().split())) for _ in range(m)]

root = [-1] * n

def find(x):
    if root[x] < 0:
        return x
    else:
        root[x] = find(root[x])
        return root[x]

def unite(x, y):
    gx = find(x)
    gy = find(y)

    if gx == gy:
        return

    if root[gx] > root[gy]:
        gx, gy = gy, gx

    root[gx] += root[gy]
    root[gy] = gx

def size(x):
    x = find(x)
    return -root[x]

for a, b in AB:
    a -= 1
    b -= 1
    unite(a, b)

ans = 0
for i in range(n):
    ans = max(ans, size(i))
print(ans)


import frappe
from frappe import _
from frappe.utils import today

"""def get_context(context):
    context.name = "Administrator"
    data = frappe.form_dict

    if frappe.db.exists({'doctype': 'Lead', 'email_id': data["paymentEmail"]}):
        request_status = True
    else:
        request_status = False

    paid_request = frappe.get_doc({
        "doctype": "Request",
        "request_type": "Paid Request",
        #"event_name": data["eventTitle"],
        "event_name": "World Data Summit 2022",
        "already_exists": request_status,
        "first_name": data["paymentFirstName"],
        "last_name": data["paymentLastName"],
        "full_name": data["paymentFirstName"] + " " + data["paymentLastName"],
        "email_address": data["paymentEmail"],
        "number_of_tickets": data["numberOfTickets"],
        "payment_fee": data["paymentFee"],
        "currency": data["paymentCurrency"],
        "language":data["paymentLanguage"],
        "tax": data["paymentSalesTax"],
        "interest_type": "Attending",
        "type": "Attendee",
        "payment_type": data["paymentType"],
        "payment_status": "Paid",
        "source": "Xing Registration"
        })
    paid_request.insert(ignore_permissions=True)"""

def get_context(context):
    context.name = "Administrator"
    data = frappe.form_dict
    number_of_tickets = data["numberOfTickets"]

    i = 0
    while i < int(number_of_tickets):
        if frappe.db.exists({'doctype': 'Attendee', 'email_address': data["ticketEmail"+str(i)]}):
            attendee = frappe.get_doc('Attendee', data["ticketEmail"+str(i)])
            row = attendee.append("tickets", {
                "ticket_id": data["ticketId"+str(i)],
                "ticket_category": data["ticketCategory"+str(i)],
                "event_name": data["eventTitle"]
                })
            row.insert()

            event = frappe.db.get_list("Events", filters={'status':'Registration'}, fields=['event_name'])
            row = event.append("attendees", {
                "attendee_id": data["ticketEmail"+str(i)],
                "payment_status": "Paid",
                })
            row.insert()
            i+=1
        else:
            new_attendee = frappe.get_doc({
                "doctype": "Attendee",
                "first_name": data["ticketFirstName"+str(i)],
                "last_name": data["ticketLastName"+str(i)],
                "full_name": data["ticketFirstName"+str(i)] + " " + data["ticketLastName"+str(i)],
                "email_address": data["ticketEmail"+str(i)],
                "country": "Unknown",
                "tickets": [
                    {
                        "ticket_id": data["ticketId"+str(i)],
                        "ticket_category": data["ticketCategory"+str(i)],
                        "event_name": data["eventTitle"]
                    }
                    ]
                })
            new_attendee.insert(ignore_permissions=True)

            event = frappe.db.get_list("Events", filters={'status':'Registration'}, fields=['event_name'])
            row = event.append("attendees", {
                "attendee_id": data["ticketEmail"+str(i)],
                "payment_status": "Paid",
                })
            row.insert()
            i+=1


from dbkdev.core import DevelopmentClient
from dbkdev.core import DevelopmentEngine
from dbkcore.core import Log
from dbkcore.core import trace
from dbkcore.core import Singleton
from dbkdev.core import IdeEnvironment
from pyspark.sql import SparkSession
from typing import Dict

class Engine(metaclass=Singleton):

    """
    This is the core of the framework.
    It configures the environment to interact with the remote Databricks.
    """

    def __init__(self):
        """
        Instantiate the current object

        """
        self.__ide_environment = DevelopmentEngine().get_instance().ide_environment
        self.appi_ik = None

    def initialize_env(self):
        """
        Initializes the DevelopmentClient.
        That is, sets the dbutils and spark context accordingly if the code is runt on cluster or locally.
        """
        DevelopmentClient(
            dbutils=DevelopmentEngine().get_instance().dbutils,
            spark=DevelopmentEngine().get_instance().spark,
            ide_environment=self.__ide_environment
        )

    # def initialize_import(self):
    #     # Setting pipeline module path
    #     if self.__ide_environment == IdeEnvironment.DATABRICKS:
    #         import sys
    #         sys.path.append(str(self.pipelines_lib_path))

    def initialize_logger(
        self,
        pipeline_name: str,
        appi_ik_scope: str = 'config',
        appi_ik_secret: str = 'APPI_IK'
    ):
        """
        Initializes the logger

        Parameters
        ----------
        pipeline_name : str
            Name to use with the logger. It will be the base name used for all the upcoming logs and tracing
        appi_ik_scope : str, optional
            Databricks secret scope where the Application Insight key is stored, by default "dds"
        appi_ik_secret : str, optional
            Databricks secret name where the Application Insight key is stored, by default "appiik"

        Raises
        ------
        ValueError
            Unknown Ide Environment used
        """
        # Configuring application insight key
        if self.__ide_environment == IdeEnvironment.LOCAL:
            from dbkenv.core import Configuration
            configurations = Configuration()
            self.appi_ik = configurations.APPINSIGHT_CONNECTIONSTRING
        elif self.__ide_environment == IdeEnvironment.DATABRICKS:
            self.appi_ik = DevelopmentEngine().get_instance().dbutils.secrets.get(appi_ik_scope, appi_ik_secret)
        else:
            raise ValueError(f'ide_environment unknown: {self.__ide_environment}')
        # Instantiating logger
        Log(pipeline_name, self.appi_ik)

    def spark(self) -> SparkSession:
        """
        Current spark context

        Returns
        -------
        SparkSession
            Spark context
        """
        return DevelopmentClient().get_instance().spark

    def dbutils(self):
        """
        Current dbutils

        Returns
        -------
        DBUtils
            The DBUtils
        """
        return DevelopmentClient().get_instance().dbutils

    @classmethod
    def get_instance(cls):
        """
        Current singleton Engine

        Returns
        -------
        Engine
            The Engine
        """
        return Engine()

    @staticmethod
    def ide_environment() -> IdeEnvironment:
        """
        Current Ide Environment

        Returns
        -------
        IdeEnvironment
            The Ide Environment
        """
        return DevelopmentClient().get_instance().ide_environment

    @staticmethod
    def is_ide_dataricks() -> bool:
        """
        Checks if the current environment is Databricks

        Returns
        -------
        bool
            Check result
        """
        return DevelopmentClient().get_instance().ide_environment == IdeEnvironment.DATABRICKS

    @staticmethod
    def is_ide_local() -> bool:
        """
        Checks if the current environment is Local

        Returns
        -------
        bool
            Check result
        """
        return DevelopmentClient().get_instance().ide_environment == IdeEnvironment.LOCAL

    def run_notebook_with_retry(self, notebook: str, args: Dict, timeout=86400, max_retries=3):
        """
        Runs the specified notebook through dbutils

        Parameters
        ----------
        notebook : str
            Name or path of the notebook
        args : Dict
            [description]
        timeout : int, optional
            [description], by default 86400
        max_retries : int, optional
            [description], by default 3

        Returns
        -------
        [type]
            [description]

        Raises
        ------
        e
            [description]
        """
        num_retries = 0
        while True:
            try:
                return DevelopmentClient().get_instance().dbutils.notebook.run(notebook, timeout, args)
            except Exception as e:
                if num_retries > max_retries:
                    raise e
                else:
                    print("Retrying error"), e
                    num_retries += 1

    @trace
    # TODO: rename check
    def run_notebook(self, notebook: str, args: Dict, timeout=86400, error_raise=True):
        try:
            res = DevelopmentClient().get_instance().dbutils.notebook.run(notebook, timeout, args)
        except Exception as e:
            res = f"Notebook {notebook} failed"
            Log().get_instance().log_error(res)
            if error_raise:
                raise e
        return res


from sqlalchemy import String, Column, MetaData, Table, delete, select

from nova.openstack.common import log as logging

LOG = logging.getLogger(__name__)

def upgrade(migrate_engine):
    meta = MetaData()
    meta.bind = migrate_engine

    aggregates = Table('aggregates', meta, autoload=True)
    aggregate_metadata = Table('aggregate_metadata', meta, autoload=True)
    # migrate data
    record_list = list(aggregates.select().execute())
    for rec in record_list:
        row = aggregate_metadata.insert()
        row.execute({'created_at': rec['created_at'],
                    'updated_at': rec['updated_at'],
                    'deleted_at': rec['deleted_at'],
                    'deleted': rec['deleted'],
                    'key': 'availability_zone',
                    'value': rec['availability_zone'],
                    'aggregate_id': rec['id'],
                    })
    aggregates.drop_column('availability_zone')

def downgrade(migrate_engine):
    meta = MetaData()
    meta.bind = migrate_engine

    aggregates = Table('aggregates', meta, autoload=True)
    aggregate_metadata = Table('aggregate_metadata', meta, autoload=True)
    availability_zone = Column('availability_zone', String(255))
    aggregates.create_column(availability_zone)
    # migrate data
    aggregates.update().values(availability_zone=select(
        [aggregate_metadata.c.value]).where(aggregates.c.id ==
        aggregate_metadata.c.aggregate_id).where(aggregate_metadata.c.key ==
        'availability_zone')).execute()
    delete(aggregate_metadata,
           aggregate_metadata.c.key == 'availability_zone').execute()
    aggregates.c.availability_zone.alter(nullable=False)


from openerp import models, fields, api

class mrp_bom_update(models.TransientModel):
    _name = "mrp.bom.update"

    name = fields.Char('name')
    source_id = fields.Many2one('product.product', 'Source')
    target_id = fields.Many2one('product.product', 'Target')

    @api.one
    def button_update_bom(self):
        source = self.source_id
        target = self.target_id

        bom_lines = self.env['mrp.bom.line'].search([('product_id','=', source.id)])
        if target:
            bom_lines.write({'product_id': target.id})
        else:
            bom_lines.unlink()
        return {'type': 'ir.actions.act_window_close'}

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:


import crcmod
import sys
import struct
import zlib

HEADER_SIG = bytes([255, 0, 254, 1])
HEADER_LEN = 16
HEADER_VERSION = bytes([0, 0, 0, 3])

# '<<ClientClient>>?VoxelTerrain'
TERRAIN_MARK = [0x3C, 0x3C, 0x43, 0x6C, 0x69, 0x65, 0x6E, 0x74, 0x43, 0x6C, 0x69, 0x65, 0x6E,
    0x74, 0x3E, 0x3E, 0xF3, 0x56, 0x6F, 0x78,0x65, 0x6C, 0x54, 0x65, 0x72, 0x72, 0x61, 0x69, 0x6E, 0x00]

class PlanetCoasterSaveFile():
    def __init__(self, fn):
        self.raw_data = None
        self.body = None
        self.header = { }
        self.terrain_start = -1

        self.crc_func = crcmod.mkCrcFun(0x104C11DB7, 0, True, 0)
        self.read_compressed_file(fn)
        self.parse_raw_data()

    def read_compressed_file(self, fn):
        f = open(fn, 'rb')
        self.raw_data = zlib.decompress(f.read(), -15)
        self.body = [ b for b in self.raw_data[16:] ]
        f.close()

    def parse_raw_data(self):
        self.parse_header(self.raw_data[0:HEADER_LEN])
        body = self.raw_data[HEADER_LEN:]
        assert(self.header['body_len'] == len(body))
        crc = self.crc_func(body)
        assert(crc == self.header['crc'])

    def create_raw_data(self):
        body = bytes(self.get_body())
        crc = struct.pack('>I', self.crc_func(body))
        body_len = struct.pack('>I', self.header['body_len'])
        return HEADER_SIG + crc + HEADER_VERSION + body_len + body

    def write_raw_file(self, fn):
        out = self.create_raw_data()
        f_out = open(fn, 'wb')
        f_out.write(out)
        f_out.close()

    def parse_header(self, h):
        assert h[0:4] == HEADER_SIG
        self.header['crc'] = int.from_bytes(h[4:8], byteorder='big')
        assert h[8:12] == HEADER_VERSION
        self.header['body_len'] = int.from_bytes(h[12:16], byteorder='big')

    def get_body(self):
        return self.body

    def find_terrain(self):
        self.terrain_start = -1
        tm_len = len(TERRAIN_MARK)
        for i in range(0, len(self.body)):
            if self.body[i : i + tm_len] == TERRAIN_MARK:
                self.terrain_start = i + tm_len
                break
        return self.terrain_start

    def output_file(self, fn):
        out = self.create_raw_data()
        compressor = zlib.compressobj(level=4, method=zlib.DEFLATED, wbits=-15)
        f_out = open(fn, 'wb')
        f_out.write(compressor.compress(out))
        f_out.write(compressor.flush(zlib.Z_FINISH))
        f_out.close()

    def set_byte(self, value, abspos=None, bodypos=None):
        if abspos is not None:
            self.body[bodypos - HEADER_LEN] = b
        elif bodypos is not None:
            self.body[bodypos] = value

    def find_bytes(self, find, start=0):
        """Find first exact match of array of bytes in the body.
        Returns offset relative to the start of the body.

        """
        length = len(find)
        for i in range(start, self.header['body_len']):
            if self.body[i : i+length] == find:
                yield i

    def replace_bytes(self, find, replace=None, all=False):
        """Change the first (or all) exact matching sequences of bytes in the
        body with another sequence.  They must be the same length.  If replace is
        None, just return the match.

        """
        matched = 0
        length = len(find)
        if replace is not None:
            assert(length == len(replace))

        for m in self.find_bytes(find):
            matched += 1
            if replace is not None:
                for j in range(0, length):
                    self.set_byte(replace[j], bodypos=m+j)
            if all is False:
                break
        return matched


#!/usr/bin/env pnpython4

#
# Dump tables in ph5 file to kef format.
#
# Steve Azevedo, April 2007
#

import argparse
import string
import sys
import logging
import time
# This provides the base functionality
from ph5.core import experiment

# Timeseries are stored as numpy arrays

PROG_VERSION = '2019.036'
LOGGER = logging.getLogger(__name__)

#
# These are to hold different parts of the meta-data
#

def init_local():
    global EXPERIMENT_T, EVENT_T, OFFSET_T, SORT_T, RESPONSE_T, REPORT_T
    global ARRAY_T, DAS_T
    global RECEIVER_T, SOH_A, INDEX_T, M_INDEX_T, DASS, TIME_T, TABLE_KEY
    # /Experiment_g/Experiment_t
    EXPERIMENT_T = None
    # /Experiment_g/Sorts_g/Event_t
    EVENT_T = {}
    # /Experiment_g/Sorts_g/Offset_t
    OFFSET_T = {}
    # /Experiment_g/Sorts_g/Sort_t
    SORT_T = None
    # /Experiment_g/Responses_g/Response_t
    RESPONSE_T = None
    # /Experiment_g/Reports_g/Report_t
    REPORT_T = None
    # /Experiment_g/Sorts_g/Array_t_[nnn]
    ARRAY_T = {}
    # /Experiment_g/Receivers_g/Das_g_[sn]/Das_t (keyed on DAS)
    DAS_T = {}
    # /Experiment_g/Receivers_g/Receiver_t
    RECEIVER_T = None
    # /Experiment_g/Receivers_g/Das_g_[sn]/SOH_a_[n] (keyed on DAS then by
    # SOH_a_[n] name)
    SOH_A = {}
    # /Experiment_g/Receivers_g/Index_t
    INDEX_T = None
    # /Experiment_g/Maps_g/Index_t
    M_INDEX_T = None
    # A list of Das_Groups that refers to Das_g_[sn]'s
    DASS = []
    # /Experiment_g/Receivers_g/Time_t
    TIME_T = None
    #
    TABLE_KEY = None

#
# To hold table rows and keys
#

class Rows_Keys(object):
    __slots__ = ('rows', 'keys')

    def __init__(self, rows=None, keys=None):
        self.rows = rows
        self.keys = keys

    def set(self, rows=None, keys=None):
        if rows is not None:
            self.rows = rows
        if keys is not None:
            self.keys = keys

#
# To hold DAS sn and references to Das_g_[sn]
#

class Das_Groups(object):
    __slots__ = ('das', 'node')

    def __init__(self, das=None, node=None):
        self.das = das
        self.node = node

#
# Read Command line arguments
#

def get_args():
    global PH5, PATH, DEBUG, EXPERIMENT_TABLE, SORT_TABLE, OFFSET_TABLE, \
        EVENT_TABLE, ARRAY_TABLE, RESPONSE_TABLE, REPORT_TABLE, \
        RECEIVER_TABLE, DAS_TABLE, TIME_TABLE, TABLE_KEY, INDEX_TABLE, \
        M_INDEX_TABLE, ALL_ARRAYS, ALL_EVENTS, OFILE

    parser = argparse.ArgumentParser(
                                formatter_class=argparse.RawTextHelpFormatter)

    parser.usage = "Version: {0}\ntabletokef\
     --nickname ph5-file-prefix options".format(
        PROG_VERSION)

    parser.description = "Dump a table to a kef file."

    parser.add_argument("-n", "--nickname", dest="ph5_file_prefix",
                        help="The ph5 file prefix (experiment nickname).",
                        metavar="ph5_file_prefix", required=True)

    parser.add_argument("-p", "--path", dest="ph5_path",
                        help=("Path to ph5 files. Default to current "
                              "directory."), default=".",
                        metavar="ph5_path")

    parser.add_argument("-u", "--update_key", dest="update_key",
                        help="Set generated kef file to do an Update on key.",
                        metavar="update_key", type=str)

    parser.add_argument("-d", "--debug", dest="debug", action="store_true",
                        default=False)

    parser.add_argument("-E", "--Experiment_t", dest="experiment_t",
                        action="store_true",
                        default=False,
                        help="Dump /Experiment_g/Experiment_t to a kef file.")

    parser.add_argument("-S", "--Sort_t", dest="sort_t", action="store_true",
                        default=False,
                        help=("Dump /Experiment_g/Sorts_g/Sort_t to a kef "
                              "file."))

    parser.add_argument("-O", "--Offset_t", dest="offset_t_", metavar="a_e",
                        help=("Dump "
                              "/Experiment_g/Sort_g/Offset_t_[arrayID_eventID]"
                              " to a kef file."))

    parser.add_argument("-V", "--Event_t_", dest="event_t_", metavar="n",
                        type=int,
                        help=("Dump /Experiment_g/Sorts_g/Event_t_[n]"
                              "to a kef file."))

    parser.add_argument("--all_events", dest='all_events', action='store_true',
                        default=False,
                        help=("Dump all /Experiment_g/Sorts_g/Event_t_xxx "
                              "to a kef file."))

    parser.add_argument("-A", "--Array_t_", dest="array_t_", metavar="n",
                        type=int,
                        help=("Dump /Experiment_g/Sorts_g/Array_t_[n] "
                              "to a kef file."))

    parser.add_argument("--all_arrays", dest='all_arrays', action='store_true',
                        default=False,
                        help=("Dump all /Experiment_g/Sorts_g/Array_t_xxx "
                              "to a kef file."))

    parser.add_argument("-R", "--Response_t", dest="response_t",
                        action="store_true",
                        default=False,
                        help=("Dump /Experiment_g/Responses_g/Response_t "
                              "to a kef file."))

    parser.add_argument("-P", "--Report_t", dest="report_t",
                        action="store_true",
                        default=False,
                        help=("Dump /Experiment_g/Reports_g/Report_t "
                              "to a kef file."))

    parser.add_argument("-C", "--Receiver_t", dest="receiver_t",
                        action="store_true",
                        default=False,
                        help=("Dump /Experiment_g/Receivers_g/Receiver_t "
                              "to a kef file."))

    parser.add_argument("-I", "--Index_t", dest="index_t", action="store_true",
                        default=False,
                        help=("Dump /Experiment_g/Receivers_g/Index_t "
                              "to a kef file."))

    parser.add_argument("-M", "--M_Index_t", dest="m_index_t",
                        action="store_true",
                        default=False,
                        help=("Dump /Experiment_g/Maps_g/Index_t to a "
                              "kef file."))

    parser.add_argument("-D", "--Das_t", dest="das_t_", metavar="das",
                        help=("Dump /Experiment_g/Receivers_g/Das_g_[das]/"
                              "Das_t to a kef file."))

    parser.add_argument("-T", "--Time_t", dest="time_t", action="store_true",
                        default=False,
                        help=("Dump /Experiment_g/Receivers_g/Time_t "
                              "to a kef file."))

    parser.add_argument("-o", "--outfile", dest="output_file",
                        help="The kef output file to be saved at.",
                        metavar="output_file", default=None)

    args = parser.parse_args()

    PH5 = args.ph5_file_prefix
    PATH = args.ph5_path
    DEBUG = args.debug
    EXPERIMENT_TABLE = args.experiment_t
    SORT_TABLE = args.sort_t
    if args.offset_t_ is not None:
        try:
            OFFSET_TABLE = map(int, args.offset_t_.split("_"))
        except Exception as e:
            LOGGER.error(
                "Offset table should be entered as arrayID underscore "
                "shotLineID, eg. 1_2 or 0_0.")
            LOGGER.error(e.message)
            sys.exit()
    else:
        OFFSET_TABLE = None
    EVENT_TABLE = args.event_t_
    TIME_TABLE = args.time_t
    INDEX_TABLE = args.index_t
    M_INDEX_TABLE = args.m_index_t
    TABLE_KEY = args.update_key
    ARRAY_TABLE = args.array_t_
    ALL_ARRAYS = args.all_arrays
    ALL_EVENTS = args.all_events
    RESPONSE_TABLE = args.response_t
    REPORT_TABLE = args.report_t
    RECEIVER_TABLE = args.receiver_t
    DAS_TABLE = args.das_t_

    table_list = [EXPERIMENT_TABLE, SORT_TABLE, OFFSET_TABLE, EVENT_TABLE,
                  TIME_TABLE, INDEX_TABLE, M_INDEX_TABLE, ARRAY_TABLE,
                  ALL_ARRAYS, ALL_EVENTS, RESPONSE_TABLE, REPORT_TABLE,
                  RECEIVER_TABLE, DAS_TABLE]
    if all(not t for t in table_list):
        LOGGER.error("No table specified for output. See --help for more "
                     "details.")

    # define OFILE to write output
    o_filename = args.output_file
    if o_filename is None:
        OFILE = None
    else:
        OFILE = open(o_filename, 'w')

#
# Initialize ph5 file
#

def initialize_ph5(editmode=False):
    '''   Initialize the ph5 file   '''
    global EX, PATH, PH5

    EX = experiment.ExperimentGroup(PATH, PH5)
    EX.ph5open(editmode)
    EX.initgroup()

#
# Print out report
#
def print_report(text):
    global OFILE
    if OFILE is None:
        print(text)
    else:
        OFILE.write(text + '\n')

#
# Print Rows_Keys
#
def table_print(t, a, fh=None):
    global TABLE_KEY
    global PATH
    global EX
    i = 0
    s = ''
    s = s + \
        "#\n#\t%s\tph5 version: %s\n#\n" % (
            time.ctime(time.time()), EX.version())
    # Loop through table rows
    for r in a.rows:
        i += 1

        s = s + "#   Table row %d\n" % i
        # Print table name
        if TABLE_KEY in a.keys:
            s = s + "{0}:Update:{1} \n".format(t, TABLE_KEY)
        else:
            s = s + t + "\n"
        # Loop through each row column and print
        for k in a.keys:
            s = s + "\t" + str(k) + "=" + str(r[k]) + "\n"
        if fh is None:
            print_report(s)
            s = ''
        else:
            fh.write(s)
            s = ''

def read_time_table():
    global EX, TIME_T

    times, time_keys = EX.ph5_g_receivers.read_time()

    TIME_T = Rows_Keys(times, time_keys)

def read_report_table():
    global EX, REPORT_T

    reports, report_keys = EX.ph5_g_reports.read_reports()

    rowskeys = Rows_Keys(reports, report_keys)

    REPORT_T = rowskeys

def read_experiment_table():
    '''   Read /Experiment_g/Experiment_t   '''
    global EX, EXPERIMENT_T

    exp, exp_keys = EX.read_experiment()

    rowskeys = Rows_Keys(exp, exp_keys)

    EXPERIMENT_T = rowskeys

def read_event_table():
    '''   Read /Experiment_g/Sorts_g/Event_t   '''
    global EX, EVENT_T

    if EVENT_TABLE == 0:
        T = "Event_t"
    else:
        T = "Event_t_{0:03d}".format(EVENT_TABLE)

    try:
        events, event_keys = EX.ph5_g_sorts.read_events(T)
    except Exception:
        LOGGER.error("Can't read {0}.\nDoes it exist?\n".format(T))
        sys.exit()

    rowskeys = Rows_Keys(events, event_keys)

    EVENT_T[T] = rowskeys

def read_all_event_table():
    global EX, EVENT_T
    import re
    EVENT_T_NAME_RE = re.compile("Event_t.*")

    names = EX.ph5_g_sorts.namesRE(EVENT_T_NAME_RE)
    for name in names:
        try:
            events, event_keys = EX.ph5_g_sorts.read_events(name)
        except Exception:
            LOGGER.error("Can't read {0}. Does it exist?".format(name))
            continue

        rowskeys = Rows_Keys(events, event_keys)
        EVENT_T[name] = rowskeys

def read_offset_table():
    '''   Read /Experinent_t/Sorts_g/Offset_t   '''
    global EX, OFFSET_T

    if OFFSET_TABLE[0] == 0 or OFFSET_TABLE[1] == 0:
        name = "Offset_t"
    else:
        name = "Offset_t_{0:03d}_{1:03d}".format(
            OFFSET_TABLE[0], OFFSET_TABLE[1])

    try:
        rows, keys = EX.ph5_g_sorts.read_offset(name)
    except Exception:
        return

    OFFSET_T[name] = Rows_Keys(rows, keys)

def read_sort_table():
    '''   Read /Experiment_t/Sorts_g/Sort_g   '''
    global EX, SORT_T

    sorts, sorts_keys = EX.ph5_g_sorts.read_sorts()

    rowskeys = Rows_Keys(sorts, sorts_keys)

    SORT_T = rowskeys

def read_sort_arrays():
    '''   Read /Experiment_t/Sorts_g/Array_t_[n]   '''
    global EX, ARRAY_T

    # We get a list of Array_t_[n] names here...
    # (these are also in Sort_t)
    names = EX.ph5_g_sorts.names()
    for n in names:
        arrays, array_keys = EX.ph5_g_sorts.read_arrays(n)

        rowskeys = Rows_Keys(arrays, array_keys)
        # We key this on the name since there can be multiple arrays
        ARRAY_T[n] = rowskeys

def read_response_table():
    '''   Read /Experiment_g/Respones_g/Response_t   '''
    global EX, RESPONSE_T

    response, response_keys = EX.ph5_g_responses.read_responses()

    rowskeys = Rows_Keys(response, response_keys)

    RESPONSE_T = rowskeys

def read_receiver_table():
    global EX, RECEIVER_T

    # Read /Experiment_g/Receivers_g/Receiver_t
    receiver, receiver_keys = EX.ph5_g_receivers.read_receiver()
    rowskeys = Rows_Keys(receiver, receiver_keys)
    RECEIVER_T = rowskeys

def read_index_table():
    global EX, INDEX_T

    rows, keys = EX.ph5_g_receivers.read_index()
    INDEX_T = Rows_Keys(rows, keys)

def read_m_index_table():
    global EX, M_INDEX_T

    rows, keys = EX.ph5_g_maps.read_index()
    M_INDEX_T = Rows_Keys(rows, keys)

def read_receivers(das=None):
    '''   Read tables and arrays (except wiggles) in Das_g_[sn]   '''
    global EX, DAS_T, RECEIVER_T, DASS, SOH_A

    dasGroups = EX.ph5_g_receivers.alldas_g()
    if das is None:
        # Get references for all das groups keyed on das
        dass = sorted(dasGroups.keys())
        # Sort by das sn
    else:
        dass = [das]

    for d in dass:
        # Get node reference
        if "Das_g_" + d not in dasGroups:
            continue

        g = dasGroups["Das_g_" + d]
        dg = Das_Groups(d, g)
        # Save a master list for later
        DASS.append(dg)

        # Set the current das group
        EX.ph5_g_receivers.setcurrent(g)

        # Read /Experiment_g/Receivers_g/Das_g_[sn]/Das_t
        das, das_keys = EX.ph5_g_receivers.read_das()
        rowskeys = Rows_Keys(das, das_keys)
        DAS_T[d] = rowskeys

        # Read SOH file(s) for this das
        SOH_A[d] = EX.ph5_g_receivers.read_soh()

#####################################################
# def readPH5
# author: Lan Dam
# updated: 201802
# read data from exp(PH5) to use for KefUtility => KefEdit.py

def readPH5(exp, filename, path, tableType, arg=None):
    # print "readPH5"
    global EX, OFFSET_TABLE, EVENT_TABLE, ARRAY_TABLE, OFFSET_T, EVENT_T,\
        ARRAY_T
    global DAS_T, DASS, SOH_A
    init_local()  # innitiate values and clear cache

    EX = exp

    if tableType == "Experiment_t":
        read_experiment_table()
        return EXPERIMENT_T

    if tableType == "Sort_t":
        read_sort_table()
        return SORT_T

    if tableType == "Offset_t":
        if arg == "Offset_t":
            OFFSET_TABLE = [0]
        else:
            OFFSET_TABLE = map(int, arg.split("_"))
        # read_offset_table() will read from global var. OFFSET_TABLE to add
        # new item into dict. OFFSET_T
        read_offset_table()
        return OFFSET_T

    if tableType == "All_Offset_t":
        for o in EX.Offset_t_names:
            if o == "Offset_t":
                OFFSET_TABLE = [0]
                read_offset_table()
                break
            OFFSET_TABLE = map(int, o.replace("Offset_t_", "").split("_"))
            read_offset_table()
        return OFFSET_T

    if tableType == "Event_t":
        EVENT_TABLE = int(arg)

        # read_event_table() will read from global var. EVENT_TABLE to add new
        # item into dict. EVENT_T
        read_event_table()
        return EVENT_T

    if tableType == "All_Event_t":

        for n in EX.Event_t_names:
            if n == 'Event_t':
                EVENT_TABLE = 0
            else:
                EVENT_TABLE = int(n.replace('Event_t_', ''))
            read_event_table()

        return EVENT_T

    if tableType == "Index_t":
        read_index_table()
        return INDEX_T

    if tableType == "Map_Index_t":
        read_m_index_table()
        return M_INDEX_T

    if tableType == "Time_t":
        read_time_table()
        return TIME_T

    if tableType == "Array_t":
        ARRAY_TABLE = arg
        read_sort_table()
        read_sort_arrays()
        arrays = ARRAY_T.keys()
        for a in arrays:
            n = int(string.split(a, '_')[2])
            if n == int(ARRAY_TABLE):
                return ARRAY_T[a]

    if tableType == "All_Array_t":
        read_sort_table()
        read_sort_arrays()
        arrays = ARRAY_T.keys()
        return ARRAY_T

    if tableType == "Response_t":
        read_response_table()
        return RESPONSE_T

    if tableType == "Report_t":
        read_report_table()
        return REPORT_T

    if tableType == "Receiver_t":
        read_receiver_table()
        return RECEIVER_T

    if tableType == "Das_t":
        read_receivers(arg)
        return DAS_T

def main():
    global PH5, PATH, DEBUG, EXPERIMENT_TABLE, SORT_TABLE, OFFSET_TABLE,\
        EVENT_TABLE, \
        ARRAY_TABLE, RESPONSE_TABLE, REPORT_TABLE, RECEIVER_TABLE, DAS_TABLE,\
        TIME_TABLE, INDEX_TABLE

    init_local()

    get_args()

    initialize_ph5()

    if EXPERIMENT_TABLE:
        read_experiment_table()
        table_print("/Experiment_g/Experiment_t", EXPERIMENT_T)

    if SORT_TABLE:
        read_sort_table()
        table_print("/Experiment_g/Sorts_g/Sort_t", SORT_T)

    if OFFSET_TABLE:
        read_offset_table()
        keys = OFFSET_T.keys()
        for k in keys:
            table_print("/Experiment_g/Sorts_g/{0}".format(k), OFFSET_T[k])

    if EVENT_TABLE is not None:
        read_event_table()
        keys = EVENT_T.keys()
        for k in keys:
            table_print("/Experiment_g/Sorts_g/{0}".format(k), EVENT_T[k])
    elif ALL_EVENTS is not False:
        read_all_event_table()
        keys = EVENT_T.keys()
        for k in keys:
            table_print("/Experiment_g/Sorts_g/{0}".format(k), EVENT_T[k])

    if INDEX_TABLE:
        read_index_table()
        table_print("/Experiment_g/Receivers_g/Index_t", INDEX_T)

    if M_INDEX_TABLE:
        read_m_index_table()
        table_print("/Experiment_g/Maps_g/Index_t", M_INDEX_T)

    if TIME_TABLE:
        read_time_table()
        table_print("/Experiment_g/Receivers_g/Time_t", TIME_T)

    if ARRAY_TABLE:
        if not SORT_T:
            read_sort_table()

        read_sort_arrays()
        arrays = ARRAY_T.keys()
        for a in arrays:
            n = int(string.split(a, '_')[2])
            if n == int(ARRAY_TABLE):
                table_print("/Experiment_g/Sorts_g/" + a, ARRAY_T[a])
    elif ALL_ARRAYS:
        if not SORT_T:
            read_sort_table()

        read_sort_arrays()
        arrays = ARRAY_T.keys()
        for a in arrays:
            table_print("/Experiment_g/Sorts_g/" + a, ARRAY_T[a])

    if RESPONSE_TABLE:
        read_response_table()
        table_print("/Experiment_g/Responses_g/Response_t", RESPONSE_T)

    if REPORT_TABLE:
        read_report_table()
        table_print("/Experiment_g/Reports_g/Report_t", REPORT_T)

    if RECEIVER_TABLE:
        read_receiver_table()
        table_print("/Experiment_g/Receivers_g/Receiver_t", RECEIVER_T)

    if DAS_TABLE:
        read_receivers(DAS_TABLE)
        dass = DAS_T.keys()
        for d in dass:
            table_print("/Experiment_g/Receivers_g/Das_g_" +
                        d + "/Das_t", DAS_T[d])

    EX.ph5close()
    if OFILE is not None:
        OFILE.close()

if __name__ == '__main__':
    main()