from datetime import datetime
from collections import OrderedDict
import json
import logging
import math
import os
from shutil import chown
import subprocess
import sys
import time
import uuid
import xml.etree.ElementTree as ET

import boto3
import requests
from sqlalchemy import create_engine, text

# Set the boto3 log levels to warning:
logging.getLogger("boto3").setLevel(logging.WARNING)
logging.getLogger("botocore").setLevel(logging.WARNING)
logging.getLogger("nose").setLevel(logging.WARNING)
logging.getLogger("s3transfer").setLevel(logging.WARNING)

logging.basicConfig(format='%(asctime)s %(levelname)s %(name)s: %(message)s', level=logging.INFO)
logger = logging.getLogger('snpfinder-pipeline_manager')

# Load the JobFlowId
with open('/mnt/var/lib/info/job-flow.json', 'r') as handle:
    JOBFLOW_ID = json.load(handle)['jobFlowId']
# Load the processing job parameters
with open('/mnt/parameters.json', 'r') as handle:
    ARGS = json.load(handle)
# Load the list of S3 paths
with open('/mnt/paths.json', 'r') as handle:
    S3_PATHS = json.load(handle)
# Create the SQL engine
ENGINE = create_engine(S3_PATHS["engine_path"])

def get_headers():
    """Retrieves an OIDC token and creates the headers needed to call the SNPFinder
    API.

    Args:
        None

    Returns:
        dict. A dictionary with the headers required to call the API.

    """
    # Get the Client ID and secret
    client_id = S3_PATHS['client_id']
    client_secret = S3_PATHS['client_secret']
    # Call the credentials endpoint
    response = requests.get(
        url=S3_PATHS['auth_token_url'],
        data={
            "grant_type": "client_credentials",
            "scope": f"api://{client_id}/.default"
        },
        auth=(client_id, client_secret)
    ).json()
    # Build the header
    access_token = response["access_token"]
    headers = {"Authorization": f"Bearer {access_token}"}
    return headers

def get_executor_vcpus(instance_type):
    """Return the number of vCPUs for the specified EC2 instance type.

    Args:
        instance_type (str): The EC2 instance type of the executor. See https://
        docs.aws.amazon.com/emr/latest/ManagementGuide/emr-supported-instance-types.html
        for a list of EMR-compatible instance types.

    Returns:
        int. The number of vCPUs on the specified EC2 instance.

    """
    vcpu_count = None
    pricing_client = boto3.client(service_name="pricing", region_name="us-east-1")
    # Get the description for the specified instanct type
    response = pricing_client.get_products(
        ServiceCode="AmazonEC2",
        Filters=[
            {
                "Type": "TERM_MATCH",
                "Field": "instanceType",
                "Value": instance_type,
            },
            {
                "Type": "TERM_MATCH",
                "Field": "location",
                "Value": "US East (N. Virginia)",
            },
            {
                "Type": "TERM_MATCH",
                "Field": "operatingSystem",
                "Value": "Linux",
            },
        ],
    )["PriceList"]
    # Parse the response for the vCPU count
    for item in response:
        vcpu_count = int(json.loads(item)["product"]["attributes"]["vcpu"])
        break
    return vcpu_count

def get_worker_memory():
    """Returns the memory available to the Spark worker nodes.

    Args:
        None.

    Returns:
        int. The memory available to Spark (in MB) on the worker instance.

    """
    # Load the YARN XML configuration file
    tree = ET.parse('/etc/hadoop/conf.empty/yarn-site.xml')
    root = tree.getroot()
    # Create a dictionary from the configuration file
    yarn_config = {x[0].text:x[1].text for x in root}
    # Get the executor memory
    worker_memory = yarn_config["yarn.scheduler.maximum-allocation-mb"]
    return float(worker_memory)

# Function to set the Spark-Submit parameters based on the EC2 instance type and reference assembly
def get_spark_parameters():
    # Get the total size of BT2 files associated with the reference assebly
    client = boto3.client('s3')
    r = client.list_objects_v2(
        Bucket=S3_PATHS['bucket'].replace('s3://', ''),
        Prefix='{}/{}/{}/'.format(S3_PATHS['bt2_directory'], ARGS['species'], ARGS['reference'])
    )

    # Sum up the sizes of the individual bt2 files
    size = 0
    for item in r['Contents']:
        size += item['Size']

    # Convert to GB and add a safety margin
    bt2_size = round(float(size) / (1024**3), 2) + 3

    # Get memory and vCPU count for the instance type
    available_memory = 0.9 * get_worker_memory() / 1024
    vcpus = get_executor_vcpus(ARGS['core_node_type'])

    # Calculate the number of vCPUs to use during the Samtools merge step (2, 3, or 4)
    num_samples_to_merge = len(ARGS['samples'].split(','))

    merge_cores = 0
    processing_time = 10**8
    for cores in range(1, 8):
        # Calculate the number of samtools instances that can fit on the cluster
        # Yarn seems to have problems handling more than 30 simultaneous tasks, so limit to 25
        instance_slots = math.floor(vcpus / cores)
        total_slots = instance_slots * int(ARGS["core_node_number"])
        if total_slots > 25:
            total_slots = 25
            instance_slots = math.floor(total_slots / int(ARGS["core_node_number"]))
        # Calculate the number of cycles required to merge all samples
        processing_cycles = math.ceil(num_samples_to_merge / total_slots)
        # Guesstimate the time to merge a sample using a given number of cores
        sample_metric = 0.0143 * (cores ** -0.678)
        # Guesstimate the total time to process all samples using a given number of cores
        total_time = processing_cycles * sample_metric
        if total_time < processing_time:
            merge_cores = cores
            processing_time = total_time
            merge_memory = math.floor(available_memory / instance_slots) - 1

    # Find the number of executors to minimize the number of wasted vCPUs
    config = {'wasted_cpus': 1000, 'vcpus_per_executor': 0, 'executors': 0}
    for executor_memory in range(int(math.ceil(bt2_size)), int(math.floor(available_memory))):
        # Calculate the number of executors
        executors = math.floor(available_memory / (executor_memory))
        # Calculate how many cores are available to each executor
        cores = math.floor(vcpus / executors)
        # Calculate how many Bowtie2 instances can run on the executor
        bowties = math.floor(executor_memory / bt2_size)
        if bowties > 5:
            bowties = 5
        # Calculate how many Bowtie2 instances can run on the instance
        total_bowties = bowties * executors
        # Calculate the number of threads for each Bowtie instance
        threads = math.floor(cores / bowties)
        # Get the number of wasted vCPUs
        waste = vcpus - (threads * bowties * executors)
        # Update config if this configuration minimizes wasted vCPU and maximizes cores/executor
        if (waste <= config['wasted_cpus']) and (executors > config['executors']):
            excess_memory = available_memory - (executors * executor_memory)
            #excess_memory -= additional_per_executor
            config = {
                'bowtie_instances': total_bowties,
                'bowties_per_executor': bowties,
                'executors': executors,
                'memory_per_executor': executor_memory,
                'vcpus_per_bowtie': threads,
                'vcpus_per_executor': cores,
                'unused_memory': round(excess_memory, 1),
                'wasted_cpus': waste,
            }
    # Return the selected configuration for the alignment jobs and the memory
    # settings required for one merge executor per node
    return config, merge_cores, merge_memory

# Function to check whether a BAM file exists in S3
def check_for_bam(sample):
    client = boto3.client('s3')
    bam_present = True

    # Check for the BAM, BAI, and JSON files
    bucket = S3_PATHS["bucket"].replace('s3://', '')
    for filetype in ['bam', 'bam.bai']:
        prefix = '{}/{}/{}/{}.{}'.format(
            S3_PATHS['bam_directory'],
            ARGS['species'],
            ARGS['reference'],
            sample,
            filetype
        )
        r = client.list_objects_v2(Bucket=bucket, Prefix=prefix)
        if r['KeyCount'] == 0:
            bam_present = False
        # .bam extension will appear for .bam.bai record, so check there are at least 2 matches
        elif (filetype == 'bam') and r['KeyCount'] < 2:
            bam_present = False

    return bam_present

# Function to generate alignment jobs for Spark
def submit_alignment_jobs():
    # Get the list of samples
    sample_list = [x.strip() for x in ARGS['samples'].split(',')]

    # Get optimized parameters for the EC2 instance type
    executor_config, _, _ = get_spark_parameters()
    if executor_config['memory_per_executor'] > 46:
        executor_config['executor_memory'] = 46
    log_message = f"Bowtie2 Configuration: {json.dumps(executor_config, indent=2)}"
    logger.info(log_message)

    # Create empty dictionary to keep track of alignment job status
    alignment_jobs = OrderedDict()

    # Generate the commands to create intermediate BAM files for each sample
    for sample in sample_list:
        # Check to see if we need to align this sample
        if ((ARGS['overwrite_BAMs'] == True) or
            (ARGS['job_type'].lower() == 'discovery') or
            (check_for_bam(sample.split('@')[0].strip()) == False)):
            # Create a tracking tag for the job
            tag = str(uuid.uuid1())
            # Build the spark-submit command
            cmd = "spark-submit --deploy-mode {} --master yarn ".format(ARGS['deploy_mode'])
            if ARGS["species"] == "canola":
                cmd += "--executor-memory 8G --executor-cores 3 "
            else:
                cmd += f"--executor-memory {executor_config['memory_per_executor']}G "
                cmd += f"--executor-cores {executor_config['bowties_per_executor']} "
            cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
            cmd += "--conf spark.yarn.tags={} ".format(tag)
            cmd += "--py-files /mnt/snpfinder-pipeline/Utilities/logger.py "
            cmd += "/mnt/snpfinder-pipeline/Alignment/spark_align.py "
            cmd += "--job_id {} --data_source {} ".format(ARGS["job_id"], ARGS["data_source"])
            cmd += "--bt2 {} --sample {} ".format( ARGS["reference"], sample)
            cmd += f"--bt2_threads {executor_config['vcpus_per_bowtie']} "
            if ("data_folder" in ARGS) and (ARGS["data_folder"]):
                cmd += "--data_folder {} ".format(ARGS["data_folder"])
            if ("max_insert" in ARGS) and (ARGS["max_insert"]):
                cmd += "--max_insert {} ".format(ARGS["max_insert"])
            if ("maponly" in ARGS) and (ARGS["maponly"] == True):
                cmd += "--maponly "
            if ("no_discordant" in ARGS) and (ARGS["no_discordant"] == True):
                cmd += "--no_discordant "
            # Add command to the task dictionary
            alignment_jobs[sample] = {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}

    # Remove any merged files from sample_list
    for index, sample in enumerate(sample_list):
        samples_to_merge = [x.strip() for x in sample.split('@')]
        if len(samples_to_merge) > 1:
            sample_list[index] = samples_to_merge[0]

    return sample_list, alignment_jobs

# Function to submit merge jobs to the cluster
def merge_intermediate_bams(sample_list):
    # Get optimized parameters for the EC2 instance type
    _, merge_cores, merge_memory = get_spark_parameters()

    # Check if any files need to be merged
    samples_to_merge, samples_to_download = (False, False)
    filtered_sample_list = []

    # Create empty dictionary to keep track of merge job status
    bam_merge_jobs = OrderedDict()

    # Create list of samples to merge and samples that need to be downloaded
    for sample in sample_list:
        if ((ARGS['overwrite_BAMs'] == True) or
            (ARGS['job_type'].lower() == 'discovery') or
            (check_for_bam(sample) == False)):
            samples_to_merge = True
            filtered_sample_list.append(sample)

    # Merge files for any samples that were aligned in this job
    if samples_to_merge == True:
        for sample in filtered_sample_list:
            # Create a tracking tag for the job
            tag = str(uuid.uuid1())

            # Build the spark-submit command
            cmd = "spark-submit --deploy-mode {} --master yarn ".format(ARGS['deploy_mode'])
            cmd += "--driver-memory 1G "
            cmd += "--executor-memory {}G --executor-cores 1 ".format(merge_memory)
            cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
            cmd += "--conf spark.yarn.tags={} ".format(tag)
            cmd += "--py-files /mnt/snpfinder-pipeline/Utilities/logger.py,"
            cmd += "/mnt/snpfinder-pipeline/Utilities/common_functions.py "
            cmd += "/mnt/snpfinder-pipeline/Alignment/spark_merge.py "
            cmd += "--job_id {} --species {} --merge_cores {} ".format(ARGS["job_id"], ARGS["species"], merge_cores)
            cmd += "--refasta {} --sample {} --data_source {} ".format(ARGS["reference"], sample, ARGS["data_source"])

            if ("data_folder" in ARGS) and (ARGS["data_folder"]):
                cmd += "--data_folder {}".format(ARGS["data_folder"])

            if ("dedup" in ARGS) and (ARGS["dedup"]):
                cmd += " --dedup"

            if ARGS["index"] == True:
                cmd += " --index"

            if ARGS['job_type'].lower() == 'discovery':
                cmd += " --bam_folder {}".format(ARGS['bam_folder'])

            # Add command to the task dictionary
            bam_merge_jobs[sample] = {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}

    return bam_merge_jobs

# Function to perform SNP-calling
def call_snps(sample_list):
    # Create empty dictionary to keep track of job statuses
    snp_calling_jobs = OrderedDict()
    # Build the list of required py-files
    pyfiles = ["common_functions", "logger", "yz_base", "yz_iupac"]
    pyfile_paths = [f"/mnt/snpfinder-pipeline/Utilities/{x}.py" for x in pyfiles]
    pyfile_string = ",".join(pyfile_paths)
    # Create the SNP-calling command for each sample
    for sample in sample_list:
        # Create a tracking tag for the job
        tag = str(uuid.uuid1())
        # Create the list of arguments
        cmd = f"spark-submit --deploy-mode {ARGS['deploy_mode']} --master yarn "
        cmd += "--executor-memory 1800M --executor-cores 1 "
        cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
        cmd += f"--conf spark.yarn.tags={tag} --py-files {pyfile_string} "
        cmd += "/mnt/snpfinder-pipeline/SNP-Calling/spark_call_snps.py "
        cmd += f"--job_id {ARGS['job_id']} --species {ARGS['species']} "
        cmd += f"--refasta {ARGS['reference']} --bam {sample} --qual {ARGS['qual']} "
        cmd += f"--min_cov {ARGS['min_cov']} --min_purity {ARGS['min_purity']} "
        cmd += f"--partition_size {ARGS['partition_size']}"
        # Add Indel commands if flag present
        if ARGS["indel"]:
            cmd += f" --indel --indel_cov {ARGS['indel_cov']}"
            cmd += f" --indel_purity {ARGS['indel_purity']}"
            cmd += f" --indel_perc {ARGS['indel_perc']}"
        # Add Hets commands if flag present
        if ARGS["hets"]:
            cmd += f" --hets --hets_cov {ARGS['hets_cov']}"
            cmd += f" --hets_purity {ARGS['hets_purity']}"
        # Add additional mpileup parameters, if present
        if ("param" in ARGS) and (ARGS["param"]):
            cmd += f" --param {ARGS['param']}"
        # option for discovery pipeline
        if ARGS["job_type"].lower() == "discovery":
            cmd += " --discovery"
            if "minBQ" in ARGS and ARGS["minBQ"]:
                cmd += f" --minBQ {ARGS['minBQ']}"
        else:
            cmd += " --minBQ 0"
        # Add command to the task dictionary
        snp_calling_jobs[sample] = {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}
    return snp_calling_jobs

# Function to merge intermediate SNP files into a single SNP file per sample
def combine_snps(sample_list):
    # Generate a tracking tag
    tag = str(uuid.uuid1())

    sample_string = '?'.join(sample_list)
    # Build the list of required py-files
    pyfiles = ["common_functions", "logger"]
    pyfile_paths = [f"/mnt/snpfinder-pipeline/Utilities/{x}.py" for x in pyfiles]
    pyfile_string = ",".join(pyfile_paths)
    # Create the list of arguments
    cmd = f"spark-submit --deploy-mode {ARGS['deploy_mode']} --master yarn "
    cmd += "--executor-memory 9G --executor-cores 2 "
    cmd += f"--conf spark.yarn.tags={tag} --py-files {pyfile_string} "
    cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
    cmd += "/mnt/snpfinder-pipeline/SNP-Calling/spark_combine_snps.py "
    cmd += f"--job_id {ARGS['job_id']} --partition_size {ARGS['partition_size']} "
    cmd += f"--bam {sample_string}"

    # Add command to the task dictionary
    snp_merge_jobs = {'task': {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}}

    return snp_merge_jobs

# Function to build the list of SNP positions for backfilling
def build_snp_list(sample_list):
    # Build the command
    tag = str(uuid.uuid1())
    sample_string = '?'.join(sample_list)

    cmd = "spark-submit --deploy-mode {} --master yarn ".format(ARGS['deploy_mode'])
    cmd += "--conf spark.yarn.tags={} ".format(tag)
    cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
    cmd += "--py-files /mnt/snpfinder-pipeline/Utilities/common_functions.py "
    cmd += "/mnt/snpfinder-pipeline/Utilities/build_snpset.py "
    cmd += "--job_id {} --sample {}".format(ARGS["job_id"], sample_string)

    # Add command to the task dictionary
    create_snpset_jobs = {'task': {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}}

    return create_snpset_jobs

# Function to sort a SNPset/SNARPset
def sort_snp_list():
    # Build the command
    tag = str(uuid.uuid1())
    cmd = "spark-submit --deploy-mode {} --master yarn ".format(ARGS['deploy_mode'])
    cmd += "--conf spark.yarn.tags={} ".format(tag)
    cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
    cmd += "--py-files /mnt/snpfinder-pipeline/Utilities/common_functions.py "
    cmd += "/mnt/snpfinder-pipeline/Utilities/sort_snpset.py "
    cmd += "--job_id {} --call_type {} ".format(ARGS["job_id"], ARGS['call_type'].lower())
    cmd += "--partition_size {} ".format(ARGS["partition_size"])

    if ARGS['call_type'].lower() == 'snarping':
        cmd += "--species {} --refasta {} --snarpset {} ".format(ARGS["species"], ARGS["reference"], ARGS["snarpset"])

    if ('snpset' in ARGS) and (ARGS['snpset'] is not None):
        cmd += " --snpset {} ".format(ARGS['snpset'])

    # Add command to the task dictionary
    sort_snpset_jobs = {'task': {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}}

    return sort_snpset_jobs

# Function to submit SNARPing/Filling jobs
def snarp_call(sample_list):
    # Create empty dictionary to keep track of job statuses
    snarping_jobs = OrderedDict()

    # Add the step to create intermediate SNP files for each sample, on a per-sample basis
    for sample in sample_list:
        # Create a tracking tag for the job
        tag = str(uuid.uuid1())

        # Create the list of required Python files
        pyfiles = [
            "common_functions",
            "logger",
            "yz_base",
            "yz_iupac",
            "yz_sample",
            "yz_snarpset",
            "yz_snarpset_locus"
        ]
        pyfile_paths = [f"/mnt/snpfinder-pipeline/Utilities/{x}.py" for x in pyfiles]
        pyfile_string = ",".join(pyfile_paths)

        # Build the spark-submit command
        call_type = ARGS["call_type"].lower()
        cmd = f"spark-submit --deploy-mode {ARGS['deploy_mode']} --master yarn "
        cmd += "--executor-memory 8G --executor-cores 1 "
        cmd += f"--py-files {pyfile_string} --conf spark.yarn.tags={tag} "
        cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
        cmd += "/mnt/snpfinder-pipeline/SNP-Calling/spark_snarp.py "
        cmd += f"--job_id {ARGS['job_id']} --species {ARGS['species']} "
        cmd += f"--refasta {ARGS['reference']} --bam {sample} "
        cmd += f"--partition_size {ARGS['partition_size']} --call_type {call_type}"
        cmd += f""

        # Add SNARPset if performing SNARPing
        if call_type == "snarping":
            cmd += f" --snarpset {ARGS['snarpset']}"

        # Look for an external SNPset
        if ("snpset" in ARGS) and (ARGS["snpset"]):
            cmd += f" --snpset {ARGS['snpset']}"

        # Options for the discovery pipeline
        if ARGS["job_type"].lower() == "discovery":
            cmd += " --discovery"
            if ("minBQ" in ARGS) and (ARGS["minBQ"]):
                cmd += f" --minBQ {ARGS['minBQ']}"
        else:
            cmd += " --minBQ 0"

        # Add command to the task dictionary
        snarping_jobs[sample] = {"tag": tag, "command": cmd, "status": "INCOMPLETE"}

    return snarping_jobs

# Function to merge SNARPing/Filling files
def merge_snarps(sample_list):
    snarp_merge_jobs = OrderedDict()
    for sample in sample_list:
        # Create a tracking tag for the job
        tag = str(uuid.uuid1())

        # Build the command
        cmd = "spark-submit --deploy-mode {} --master yarn ".format(ARGS['deploy_mode'])
        cmd += "--executor-memory 4G --executor-cores 1 "
        cmd += "--conf spark.yarn.tags={} ".format(tag)
        cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
        cmd += "--py-files /mnt/snpfinder-pipeline/Utilities/logger.py "
        cmd += "/mnt/snpfinder-pipeline/SNP-Calling/spark_merge_snarps.py "
        cmd += "--job_id {} --refasta {} ".format(ARGS["job_id"], ARGS["reference"])
        cmd += "--partition_size {} --call_type {} ".format(ARGS["partition_size"], ARGS['call_type'].lower())
        cmd += "--job_type {} --sample {}".format(ARGS['job_type'].lower(), sample)

        # Add command to the task dictionary
        snarp_merge_jobs[sample] = {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}

    return snarp_merge_jobs

# Function to create a project-level parquet file
def write_project_parquet(sample_list):
    # Create a tracking tag for the job
    tag = str(uuid.uuid1())

    sample_string = '?'.join(sample_list)

    # Build the command
    cmd = "spark-submit --deploy-mode {} --master yarn ".format(ARGS['deploy_mode'])
    cmd += "--conf spark.yarn.tags={} --conf spark.driver.maxResultSize=0 ".format(tag)
    if ("large_mem" in ARGS) and (ARGS["large_mem"]):
        cmd += "--driver-memory 122g --executor-memory 64g "
        cmd += "--conf spark.driver.memoryOverhead=16g "
        cmd += "--conf spark.executor.memoryOverhead=16g "
        cmd += "--conf spark.dynamicAllocation.maxExecutors=100 "
        cmd += "--conf spark.yarn.maxAppAttempts=1 "
        cmd += "--conf spark.shuffle.registration.maxAttempts=1 "
        cmd += "--conf spark.shuffle.registration.timeout=600000 "
        cmd += "--conf spark.network.timeout=10000000 "
        cmd += "--conf spark.executor.heartbeatInterval=1000000 "
        cmd += "--conf spark.sql.parquet.compression.codec=uncompressed "
        cmd += "--conf spark.local.dir=/mnt/tmp "
    else:
        cmd += "--conf spark.network.timeout=10000000 "
        cmd += "--conf spark.executor.heartbeatInterval=1000000 "
        cmd += "--conf spark.local.dir=/mnt/tmp "
        if len(sample_list) > 500:
            cmd += "--conf spark.sql.parquet.compression.codec=uncompressed "
        cmd += "--conf spark.serializer=org.apache.spark.serializer.KryoSerializer "
        cmd += "--conf spark.executor.extraJavaOptions='-XX:+UseG1GC -XX:InitiatingHeapOccupancyPercent=35' "
        cmd += "--conf spark.driver.extraJavaOptions='-XX:+UseG1GC -XX:InitiatingHeapOccupancyPercent=35' "
        cmd += "--executor-memory 14G --conf spark.executor.memoryOverhead=5g --executor-cores 1 "
    cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
    cmd += "--py-files /mnt/snpfinder-pipeline/Utilities/common_functions.py "
    cmd += "/mnt/snpfinder-pipeline/Utilities/write_project_parquet.py "
    cmd += "--job_id {} --species {} ".format(ARGS["job_id"], ARGS["species"])
    cmd += "--refasta {} --job_type {} ".format(ARGS["reference"], ARGS["job_type"])
    cmd += "--call_type {} --project_path {} ".format(ARGS["call_type"], ARGS["project_path"])
    cmd += "--bam {}".format(sample_string)

    # Add command to the task dictionary
    project_parquet_jobs = {'task': {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}}

    return project_parquet_jobs

def write_metaproject_parquet():
    """ Builds the spark-submit command and job status dictionary for the job to
    create the metaproject-level parquet file

    Args:
        None.

    Returns:
        dict. A job status dictionary with the format
        {'task': {'tag': <uuid>, 'command': <spark-submit_command>, 'status': 'INCOMPLETE'}}

    """

    # Check if this is a metaproject
    path_parts = ARGS["project_path"].split("/")
    if len(path_parts) > 1:
        # Create a tracking tag for the job
        tag = str(uuid.uuid1())
        # Build the command
        cmd = "spark-submit --deploy-mode {} --master yarn ".format(ARGS['deploy_mode'])
        cmd += "--conf spark.yarn.tags={} --conf spark.driver.maxResultSize=0 ".format(tag)
        cmd += "--conf spark.network.timeout=10000000 "
        cmd += "--conf spark.executor.heartbeatInterval=1000000 "
        cmd += "--conf spark.local.dir=/mnt/tmp --conf spark.yarn.maxAppAttempts=1 "
        cmd += "--conf spark.sql.parquet.compression.codec=uncompressed "
        cmd += "--conf spark.serializer=org.apache.spark.serializer.KryoSerializer "
        cmd += "--conf spark.executor.extraJavaOptions='-XX:+UseG1GC -XX:InitiatingHeapOccupancyPercent=35' "
        cmd += "--conf spark.driver.extraJavaOptions='-XX:+UseG1GC -XX:InitiatingHeapOccupancyPercent=35' "
        cmd += "--executor-memory 14G --conf spark.executor.memoryOverhead=5g  --executor-cores 1 "
        cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
        cmd += "--py-files /mnt/snpfinder-pipeline/Utilities/common_functions.py "
        cmd += "/mnt/snpfinder-pipeline/Utilities/write_metaproject_parquet.py "
        cmd += "--species {} --refasta {} --metaproject {}".format(
            ARGS["species"], ARGS["reference"], path_parts[0]
        )
        # Add command to the task dictionary
        metaproject_parquet_jobs = {'task': {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}}
    else:
        metaproject_parquet_jobs = {}
    return metaproject_parquet_jobs

# Function to update the locus data
def update_locus(sample_list):
    # Create a tracking tag for the job
    tag = str(uuid.uuid1())

    sample_string = '?'.join(sample_list)

    # Build the command
    cmd = "spark-submit --deploy-mode {} --master yarn ".format(ARGS['deploy_mode'])
    cmd += "--conf spark.yarn.tags={} ".format(tag)
    cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
    cmd += "--py-files /mnt/snpfinder-pipeline/Utilities/common_functions.py "
    cmd += "/mnt/snpfinder-pipeline/Utilities/project_update_locus.py "
    cmd += "--job_id {} --species {} ".format(ARGS["job_id"], ARGS["species"])
    cmd += "--refasta {} --project_path {} ".format(ARGS["reference"], ARGS["project_path"])
    cmd += " --sample {}".format(sample_string)

    # Add command to the task dictionary
    locus_update_jobs = {'task': {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}}

    return locus_update_jobs

# Function to update the different metadata
def update_metadata(sample_list):
    # Create a tracking tag for the job
    tag = str(uuid.uuid1())

    sample_string = '?'.join(sample_list)

    # Build the command
    cmd = "spark-submit --deploy-mode {} --master yarn ".format(ARGS['deploy_mode'])
    cmd += "--conf spark.yarn.tags={} ".format(tag)
    cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
    cmd += "--py-files /mnt/snpfinder-pipeline/Utilities/common_functions.py "
    cmd += "/mnt/snpfinder-pipeline/Utilities/update_metadata.py "
    cmd += "--job_id {} --project_path {} ".format(ARGS["job_id"], ARGS["project_path"])
    cmd += "--reference {} --species {} ".format(ARGS["reference"], ARGS["species"])
    cmd += "--samples {} ".format(sample_string)
    cmd += "--call_type {}".format(ARGS['call_type'].lower())

    # Add command to the task dictionary
    metadata_update_jobs = {'task': {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}}

    return metadata_update_jobs

# Function to send new loci to Lexicon
def submit_new_loci():
    # Create a tracking tag for the job
    tag = str(uuid.uuid1())

    # Build the command
    cmd = "spark-submit --deploy-mode {} --master yarn ".format(ARGS['deploy_mode'])
    cmd += "--conf spark.yarn.tags={} ".format(tag)
    cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
    cmd += "/mnt/snpfinder-pipeline/Utilities/submit_to_lexicon.py "
    cmd += "--refasta {} --species {}".format(ARGS["reference"], ARGS["species"])

    # Add command to the task dictionary
    lexicon_request_jobs = {'task': {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}}

    return lexicon_request_jobs

# Function to clean up intermediate files
def tidy_up(sample_list):
    # Create a tracking tag for the job
    tag = str(uuid.uuid1())
    sample_string = '?'.join(sample_list)

    # Build the command
    cmd = "spark-submit --deploy-mode {} --master yarn ".format(ARGS['deploy_mode'])
    cmd += "--conf spark.yarn.tags={} ".format(tag)
    cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
    cmd += "--py-files /mnt/snpfinder-pipeline/Utilities/common_functions.py "
    cmd += "/mnt/snpfinder-pipeline/Utilities/clean_EFS.py "
    cmd += "--job_id {} --samples {}".format(ARGS['job_id'], sample_string)

    # Add command to the task dictionary
    housekeeping_jobs = {'task': {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}}

    return housekeeping_jobs

# Function to write the Parquet file to CSV
def write_csv():
    # Create a tracking tag for the job
    tag = str(uuid.uuid1())

    # Build the command
    cmd = "spark-submit --deploy-mode {} --master yarn ".format(ARGS['deploy_mode'])
    cmd += "--conf spark.network.timeout=10000000 --conf spark.executor.heartbeatInterval=10000000 "
    cmd += "--conf spark.local.dir=/mnt/tmp --conf spark.driver.memory=4g --executor-memory 3g "
    cmd += "--conf spark.yarn.executor.memoryOverhead=5120 --executor-cores 5 "
    cmd += "--py-files /mnt/snpfinder-pipeline/Utilities/SNP.py "
    cmd += "--jars /usr/lib/spark/examples/jars/spark-examples.jar "
    cmd += "/mnt/snpfinder-pipeline/Utilities/parquet2csv.py "
    cmd += "--job_id {} --species {} --refasta {} ".format(ARGS['job_id'], ARGS["species"], ARGS["reference"])
    cmd += "--project_path {} --csv_folder {} ".format(ARGS["project_path"], ARGS["csv_folder"])
    if ("min_cov" in ARGS) and (ARGS["min_cov"]):
        cmd += "--min_cov {} ".format(ARGS["min_cov"])
    if ("min_purity" in ARGS) and (ARGS["min_purity"]):
        cmd += "--min_purity {} ".format(ARGS["min_purity"])
    if ("hets_cov" in ARGS) and (ARGS["hets_cov"]):
        cmd += "--hets_cov {} ".format(ARGS["hets_cov"])
    if ("hets_purity" in ARGS) and (ARGS["hets_purity"]):
        cmd += "--hets_purity {} ".format(ARGS["hets_purity"])
    if ("indel_cov" in ARGS) and (ARGS["indel_cov"]):
        cmd += "--indel_cov {} ".format(ARGS["indel_cov"])
    if ("indel_purity" in ARGS) and (ARGS["indel_purity"]):
        cmd += "--indel_purity {} ".format(ARGS["indel_purity"])
    # Update project_path parameter to point to temporary location of parquet file
    # in intermediate_files

    write_csv_jobs = {'task': {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}}

    return write_csv_jobs

# Function to build a dictionary of commands required to complete the job
def get_commands():
    # 1. Get the alignment commnads and list of files (excluding those that have been merged)
    # For discovery, don't download any files
    sample_list, alignment_jobs = submit_alignment_jobs()

    # 2. Get the BAM aggregation commnads and list of files to download
    # For discovery, save to a different bucket than SNPFinder
    bam_merge_jobs = merge_intermediate_bams(sample_list)

    # Only execute SNP-calling and SNP set creation for Filling
    if (ARGS['call_type'].lower() == 'filling') and ('snpset' not in ARGS):
        # 3. Get the SNP-calling job commands
        snp_calling_jobs = call_snps(sample_list)

        # 4. Get the SNP aggregation job commands
        snp_merge_jobs = combine_snps(sample_list)

        # 5. Build the sorted SNPset to use in filling
        create_snpset_jobs = build_snp_list(sample_list)
    else:
        snp_calling_jobs = {}
        snp_merge_jobs = {}
        create_snpset_jobs = {}

    # 6. Sort the SNARPset/SNPset
    sort_snpset_jobs = sort_snp_list()

    # 7. Get SNARPing/Filling commands
    snarping_jobs = snarp_call(sample_list)

    # 8. Get SNARPing/Filling aggregation commands
    snarp_merge_jobs = merge_snarps(sample_list)

    # 9a. Get the project parquet file command
    #     (Write this to a different location than SNPFinder (intermdediate_files?))
    project_parquet_jobs = write_project_parquet(sample_list)

    # 9b. If this is a metaproject, write the metaproject-level Parquet file
    metaproject_parquet_jobs = write_metaproject_parquet()

    if ARGS['job_type'].lower() == 'discovery':
        locus_update_jobs = {}
        metadata_update_jobs = {}
        lexicon_request_jobs = {}
        # 10. Write to user specified location
        csv_jobs = write_csv()
        # 11. Clean up Parquet file
        housekeeping_jobs = tidy_up(sample_list)
    else:
        # 10. Get Locus file update command
        locus_update_jobs = update_locus(sample_list)

        # 11. Get metadata update command
        metadata_update_jobs = update_metadata(sample_list)

        # 12. Get the command to submit new loci to lexicon
        lexicon_request_jobs = submit_new_loci()

        # 13. Get the housekeeping command
        housekeeping_jobs = tidy_up(sample_list)

        csv_jobs = {}

    # Assemble the list of commands
    command_list = {
        'alignment_jobs': alignment_jobs,
        'bam_merge_jobs': bam_merge_jobs,
        'snp_calling_jobs': snp_calling_jobs,
        'snp_merge_jobs': snp_merge_jobs,
        'create_snpset_jobs': create_snpset_jobs,
        'sort_snpset_jobs': sort_snpset_jobs,
        'snarping_jobs': snarping_jobs,
        'snarp_merge_jobs': snarp_merge_jobs,
        'project_parquet_jobs': project_parquet_jobs,
        'metaproject_parquet_jobs': metaproject_parquet_jobs,
        'locus_update_jobs': locus_update_jobs,
        'metadata_update_jobs': metadata_update_jobs,
        'lexicon_request_jobs': lexicon_request_jobs,
        'csv_jobs': csv_jobs,
        'housekeeping_jobs': housekeeping_jobs
    }
    # Add command for creating metaproject file

    return command_list, sample_list

# Function to get the SAMPLE_METRICS table IDs for each sample
# Function to get the SAMPLE_METRICS table IDs for each sample
def get_sample_metrics_ids(sample_list):
    api_url = S3_PATHS["api_base"]
    # Download any metadata from S3
    provided_metadata = None
    if ("metadata_path" in ARGS) and (ARGS["metadata_path"] is not None):
        s3 = boto3.resource('s3')
        metadata_path = ARGS["metadata_path"]
        path_parts = metadata_path.replace('s3://', '').split('/')
        bucket = path_parts[0]
        key = '/'.join(path_parts[1:])
        obj = s3.Object(bucket, key)
        raw_json = obj.get()
        provided_metadata = json.loads(raw_json['Body'].read().decode('utf-8'))
    """ Get the PROJECT table ID"""
    # Build the project service URL
    lims_project_id = ARGS["project_path"].split('/')[-1].strip()
    project_url = '{}/snpfinder/projects/metadata?lims_project_id={}'.format(api_url, lims_project_id)
    # Add project fields
    project_path = '{}/{}'.format(ARGS["reference"], ARGS["project_path"])
    data = {
        "assembly_name": ARGS["reference"],
        "path": project_path
    }
    logger.info('project_url = {}'.format(project_url))
    logger.info('project_path = {}'.format(project_path))
    # Update the PROJECT table and get a PROJECT_ID
    r = requests.post(project_url, data=json.dumps(data), headers=get_headers()).json()
    logger.info('{}'.format(json.dumps(r, indent=2, sort_keys=True)))
    project_id = r['data']['id']
    # Set the parquet path equal to the project path
    with ENGINE.connect() as conn:
        sql = text("update PROJECT set PARQUET_PATH=`PATH` where ID=:p1")
        _ = conn.execute(sql, {"p1": project_id})
    # Check to see if project metadata was provided
    if (provided_metadata is not None) and (lims_project_id in provided_metadata):
        # Execute the put request to update sample metadata
        logger.info('Updating project {} with user-provided metadata'.format(project_path))
        data = provided_metadata[lims_project_id]
        data['assembly_id'] = r['data']['assembly_id']
        data['lims_project_id'] = r['data']['lims_project_id']
        data['path'] = project_path
        project_update_url = '{}/snpfinder/projects/{}'.format(api_url, project_id)
        r = requests.put(project_update_url, data=json.dumps(data), headers=get_headers()).json()
        logger.info('{}'.format(json.dumps(r, indent=2, sort_keys=True)))
    """ Get the SAMPLE_METRICS IDs"""
    sample_dictionary = {}
    for sample in sample_list:
        data = {"lims_sample_ids": [sample],"path": project_path}
        samples_search_url = '{}/snpfinder/samples/search'.format(api_url)
        r = requests.post(samples_search_url, data=json.dumps(data), headers=get_headers()).json()
        if not r['data']:
            # Get the sample ID from the SAMPLE table
            samples_url = '{}/snpfinder/samples/metadata?lims_sample_id={}'.format(api_url, sample)
            r = requests.post(samples_url, headers=get_headers()).json()
            sample_id = r['data']['sample_id']
        else:
            sample_id = r['data'][0]['sample_id']
        logger.info('{}'.format(json.dumps(r, indent=2, sort_keys=True)))
        # Check to see if sample metadata was provided
        if (provided_metadata is not None) and (sample in provided_metadata):
            # Execute the put request to update sample metadata
            logger.info('Updating sample {} with user-provided metadata'.format(sample))
            data = provided_metadata[sample]
            sample_update_url = '{}/snpfinder/samples/{}'.format(api_url, sample_id)
            r = requests.put(sample_update_url, data=json.dumps(data), headers=get_headers()).json()
            logger.info('{}'.format(json.dumps(r, indent=2, sort_keys=True)))
        # Build the list of metrics
        data ={
            "sample_id": sample_id,
            "project_id": project_id,
            "total_coverage": None,
            "total_aligned_coverage": None,
            "total_reads": None,
            "snps_found": None,
            "sample_uuid": None
        }
        # Update the metrics table
        samples_metrics_url = '{}/snpfinder/samples/metrics'.format(api_url)
        r = requests.post(samples_metrics_url, data=json.dumps(data), headers=get_headers()).json()
        if ('Exception' in r['data']) or ('data' not in r):
            # Query the database directly if the API fails to respond correctly
            sql = text("SELECT ID FROM SAMPLE_METRICS WHERE PROJECT_ID=:p1 AND SAMPLE_ID=:p2;")
            values = {"p1": project_id, "p2": sample_id}
            with ENGINE.connect() as conn:
                sample_metrics_id = conn.execute(sql, values).fetchone()[0]
            logger.info('SAMPLE_METRICS.ID for {} already exists: {}'.format(sample, sample_metrics_id))
        else:
            logger.info('{}'.format(json.dumps(r, indent=2, sort_keys=True)))
            sample_metrics_id = r['data']['id']
        sample_dictionary[sample] = sample_metrics_id
    """ Write the dictionary to file """
    # Check to see if a project folder exists
    project_path = '/efs/{}/'.format(ARGS["job_id"])
    if not os.path.exists(project_path):
        os.makedirs(project_path, mode=0o777, exist_ok=True)
        chown(project_path, group='hadoop')
        os.chmod(project_path, 0o777)
    with open('{}/sample_dictionary.json'.format(project_path), 'w') as f:
        json.dump(sample_dictionary, f, indent=2, sort_keys=True)

# Function to return metrics on cluster activity
def get_metrics(yarn_metrics_api):
    success = False
    for attempt in range(0,10):
        if success == True:
            break
        else:
            try:
                r = requests.get(yarn_metrics_api).json()
                success = True
            except:
                # Wait and try again
                time.sleep(attempt * 2)

    metrics = {
        'active_containers': r['clusterMetrics']['containersAllocated'],
        'pending_containers': r['clusterMetrics']['containersPending'],
        'active_jobs': r['clusterMetrics']['appsRunning'],
        'available_memory': math.floor(float(r['clusterMetrics']['availableMB']) / 1024)
    }

    return metrics

# Function to load the initial list of commands to the database
def load_commands_to_database(command_list):
    # Build the list of steps in the job
    step_list = ['alignment_jobs', 'bam_merge_jobs']
    if len(command_list['snp_calling_jobs']) > 0:
        step_list += ['snp_calling_jobs', 'snp_merge_jobs', 'create_snpset_jobs']
    step_list += ['sort_snpset_jobs', 'snarping_jobs', 'snarp_merge_jobs', 'project_parquet_jobs']
    if len(command_list['metaproject_parquet_jobs']) > 0:
        step_list += ['metaproject_parquet_jobs']
    if len(command_list['locus_update_jobs']) > 0:
        step_list += ['locus_update_jobs', 'metadata_update_jobs', 'lexicon_request_jobs', 'housekeeping_jobs']
    else:
        step_list += ['csv_jobs', 'housekeeping_jobs']
    job_id = ARGS['job_id']

    with ENGINE.connect() as conn:
        # Check to see if tasks for the job already exist in the database (i.e. if a spot job was killed)
        sql = text(
            """
            SELECT TRACKING_ID, YARN_ID, STATUS, DESCRIPTION, ID
            FROM EMR_JOB_STEPS
            WHERE EMR_JOB_ID=:p1;
            """
        )
        values = {"p1": job_id}
        r = conn.execute(sql, values).fetchall()

        # Load commands to DB for new jobs
        if len(r) == 0:
            task_count = 1
            # Iterate over steps
            for step in step_list:
                # Iterate over samples
                for sample in command_list[step]:
                    # Generate a description of the task
                    task_description = step[0:-1].replace('_', ' ').upper()
                    if sample != 'task':
                        task_description += ' for {}'.format(sample)
                    # Get the status and command
                    task_status = command_list[step][sample]['status']
                    task_command = command_list[step][sample]['command']
                    task_tracking_id = command_list[step][sample]['tag']
                    # Build the insert statement
                    sql = text(
                        """
                        INSERT INTO EMR_JOB_STEPS (
                            EMR_JOB_ID, STEP_INDEX, TRACKING_ID, DESCRIPTION,
                            STATUS, YARN_COMMAND, CREATED_BY
                        ) VALUES (:p1, :p2, :p3, :p4, :p5, :p6, 'system');
                        """
                    )
                    values = {
                        "p1": job_id,
                        "p2": task_count,
                        "p3": task_tracking_id,
                        "p4": task_description,
                        "p5": task_status,
                        "p6": task_command
                    }
                    # Update the database
                    _ = conn.execute(sql, values)
                    # Increment the counter
                    task_count += 1
        else:
            # Update the command dictionary with the values from the previous run of the job
            for row in r:
                (task_tracking_id, application_id, task_status, task_description, task_id) = row
                step_sample = task_description.split(' for ')
                if len(step_sample) == 1:
                    step = step_sample[0]
                    sample = 'task'
                else:
                    (step, sample) = step_sample
                # Convert the step value to the key name
                step = step.strip().replace(' ', '_').lower() + 's'
                if sample in command_list[step]:
                    logger.info('Found {} for sample {}'.format(step, sample))
                    # Only change the status and application_id for tasks that succeeded
                    if (task_status == "SUCCEEDED"):
                        command_list[step][sample]['status'] = task_status
                        command_list[step][sample]['application_id'] = application_id
                    else:
                        # Reset the status values in the database
                        status = command_list[step][sample]['status']
                        task_tracking_id = command_list[step][sample]['tag']
                        sql = text(
                            """
                            UPDATE EMR_JOB_STEPS
                            SET STATUS=:p1, TRACKING_ID=:p2, YARN_ID=null
                            WHERE ID=:p3
                            """
                        )
                        values = {"p1": status, "p2": task_tracking_id, "p3": task_id}
                        _ = conn.execute(sql, values)
                else:
                    logger.info('Could not find {} for sample {}'.format(step, sample))

    return command_list

# Function to update a DB entry for a Yarn task with a Yarn ID
def update_task_id(application_id, tag):
    # Genreate update statement
    sql = text(
        """
        UPDATE EMR_JOB_STEPS
        SET YARN_ID=:p1, UPDATED_BY='system'
        WHERE TRACKING_ID=:p2;
        """
    )
    values = {"p1": application_id, "p2": tag}
    # Update the database
    with ENGINE.connect() as conn:
        _ = conn.execute(sql, values)

# Function to update the status for a Yarn task DB entry
def update_task_status(new_status, tag, application_id, start_time=None, end_time=None):
    # Genreate SQL update statement
    if (new_status in ('SUCCEEDED', 'FAILED')) and (application_id is not None):
        # Get the S3 paths
        log_dir = S3_PATHS["emrLogUri"]

        # Build the log path in S3 and generate the logs
        log_path = '{}{}/containers/{}/{}_logs.txt'.format(log_dir, JOBFLOW_ID, application_id, application_id)
        cmd = "yarn logs -applicationId {} | aws s3 cp - {}".format(application_id, log_path)
        with open(os.devnull, 'w') as FNULL:
            r = subprocess.Popen(cmd, shell=True, stdout=FNULL, stderr=subprocess.STDOUT)

        # Create the log URL
        log_url = 'https://s3.console.aws.amazon.com/s3/object/{}?region=us-east-1&tab=overview'.format(log_path.replace('s3://', ''))

        sql = text(
            """
            UPDATE EMR_JOB_STEPS
                SET STATUS=:p1,
                YARN_LOG_URL=:p2,
                START_TIME=:p3,
                END_TIME=:p4,
                UPDATED_BY='system'
            WHERE TRACKING_ID=:p5;
            """
        )
        values = {
            "p1": new_status,
            "p2": log_url,
            "p3": start_time,
            "p4": end_time,
            "p5": tag
        }
    else:
        sql = text(
            """
            UPDATE EMR_JOB_STEPS
                SET STATUS=:p1,
                START_TIME=:p2,
                END_TIME=:p3,
                UPDATED_BY='system'
            WHERE TRACKING_ID=:p4;
            """
        )
        values = {"p1": new_status, "p2": start_time, "p3": end_time, "p4": tag}

    # Update the database
    db_update_successful = 0
    sleep_time = 2
    while db_update_successful == 0:
        try:
            with ENGINE.connect() as conn:
                _ = conn.execute(sql, values)
                db_update_successful = 1
        except:
            logger.warning(f'Error connecting to the database. Waiting {sleep_time}s to retry')
            time.sleep(sleep_time)
            sleep_time = sleep_time * 2

# Function to get the status of a Yarn task from the API data
def get_status(api_data, command_list, job_key, sample_id, current_status, tag):
    # Command list has the following format:
    """
    command_list = {
        job_key: {
            sample_id: {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}
        }
    }
    """
    # Iterate over the API data to look for the tag
    if (api_data is not None) and ('apps' in api_data) and (api_data['apps'] is not None):
        application_id = None
        for application in api_data['apps']['app']:
            if application['applicationTags'] == tag:
                # Get the Yarn application ID
                application_id = application['id']
                start_time =  datetime.utcfromtimestamp(application['startedTime']/1000)
                end_time = None

                # Get the Yarn job status
                if application['finalStatus'] == 'UNDEFINED':
                    new_status = application['state']
                else:
                    new_status = application['finalStatus']
                    end_time = datetime.utcfromtimestamp(application['finishedTime']/1000)

                # If the job_is not present in the sample_id record, add it
                if 'application_id' not in command_list[job_key][sample_id]:
                    command_list[job_key][sample_id]['application_id'] = application_id

                    # Update the database
                    update_task_id(application_id, tag)

                # Update the dictionary if the status has changed
                if current_status != new_status:
                    command_list[job_key][sample_id]['status'] = new_status

                    job_string = job_key[:-1].replace('_', ' ').upper()
                    if sample_id == 'task':
                        log_message = (
                            f"The status of the {job_string} changed from "
                            f"{current_status} to {new_status}"
                        )
                    else:
                        log_message = (
                            f"The status of the {job_string} for sample {sample_id} "
                            f"changed from {current_status} to {new_status}"
                        )
                    logger.info(log_message)
                    # Update the database
                    if 'application_id' in command_list[job_key][sample_id]:
                        application_id = command_list[job_key][sample_id]['application_id']
                    else:
                        application_id = None

                    update_task_status(new_status, tag, application_id, start_time, end_time)

    return command_list

# Function to get the job ID and status from Yarn using an application tag
def update_status(yarn_apps_api, command_list):
    # Command list has the following format:
    """
    command_list = {
        job_key: {
            sample_id: {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}
        }
    }
    """
    # Get Yarn status from API
    api_data = requests.get(yarn_apps_api).json()

    # Check status of each job
    for job_key in command_list:
        for sample_id in command_list[job_key]:
            # Get the tag associated with the job
            tag = command_list[job_key][sample_id]['tag']
            current_status = command_list[job_key][sample_id]['status']

            # Update the status for each job
            command_list = get_status(api_data, command_list, job_key, sample_id, current_status, tag)

    return command_list

# Function to terminate the EMR cluster
def terminate_jobflow():
    # Establish a boto3 client
    client = boto3.client('emr')

    # Send the terminate request
    r = client.terminate_job_flows(JobFlowIds=[JOBFLOW_ID])

def get_cluster_metadata():
    """Returns the cluster's start time, ready time, and normalized instance hours.

    Args:
        None.

    Returns:
        tuple. A tuple of (<start_time>, <ready_time>, <normalized_instance_hours>)

    """
    response = boto3.client("emr").describe_cluster(ClusterId=JOBFLOW_ID)
    local_start_time = response["Cluster"]["Status"]["Timeline"]["CreationDateTime"]
    local_ready_time = response["Cluster"]["Status"]["Timeline"]["ReadyDateTime"]
    normalized_instance_hours = response["Cluster"]["NormalizedInstanceHours"]
    start_time = (
        local_start_time + local_start_time.tzinfo.utcoffset(local_start_time)
    ).replace(tzinfo=None)
    ready_time = (
        local_ready_time + local_ready_time.tzinfo.utcoffset(local_ready_time)
    ).replace(tzinfo=None)
    # Update the normalized instance hours value
    # Boto3 only updates the value after every hour of run time. If you are in
    # the next hour, the normalized instance hours will end up being 1 hour
    # greater than what boto3 reports.
    run_time = datetime.now() - start_time
    completed_run_hours = math.floor(run_time.total_seconds() / 3600)
    if completed_run_hours > 0:
        billed_run_hours = math.ceil(run_time.total_seconds() / 3600)
        run_rate = normalized_instance_hours / completed_run_hours
        # This value will be off if nodes are added to/removed from the cluster
        # because of auto-scaling.
        normalized_instance_hours = billed_run_hours * run_rate
    return start_time, ready_time, normalized_instance_hours

# Function to update the DB with the status of the processing job
def update_job_status(job_id, status):
    """Updates the "STATUS" field of the EMR_JOBS table.

    Args:
        job_id (int): The ID of the job to update.

        status (str): The value to set the status to.

    Returns:
        None.

    """
    # Get the cluster metadata
    start_time, _, norm_instance_hours = get_cluster_metadata()
    if status == "RUNNING":
        field_name = "START_TIME"
        param_2 = start_time
    else:
        field_name = "END_TIME"
        param_2 = datetime.now()
    sql = text(
        f"""
        UPDATE EMR_JOBS
            SET STATUS=:p1,
            {field_name}=:p2,
            NORM_INST_HOURS=:p3,
            UPDATED_BY='system'
        WHERE ID=:p4;
        """
    )
    values = {
        "p1": status,
        "p2": param_2,
        "p3": norm_instance_hours,
        "p4": job_id
    }
    with ENGINE.connect() as conn:
        conn.execute(sql, values)

# Function to notify the submitter of status changes via SNS
def send_sns_notification(failed_samples):
    with open('/mnt/var/lib/info/job-flow.json', 'r') as f:
        j = json.load(f)['jobFlowId']
    for item in j['instanceGroups']:
        node_role = item["instanceRole"].lower()
        node_type = item["instanceType"]
        node_count = item["requestedInstanceCount"]
        if node_role == "master":
            if node_count == 1:
                part_1 = '{}x {} {} node'.format(node_count, node_type, node_role)
            else:
                part_1 = '{}x {} {} nodes'.format(node_count, node_type, node_role)
        else:
            part_2 = '{}x {} {} nodes'.format(node_count, node_type, node_role)

    cluster_description = '{} and {}'.format(part_1, part_2)

    # Get cluster metadata from boto3
    start_time, ready_time, normalized_instance_hours = get_cluster_metadata()

    # Compute processing durations
    now = datetime.now(start_time.tzinfo)
    creation_time = ready_time - start_time
    running_time = now - ready_time

    # Build the SNS message body
    if len(failed_samples) == 0:
        msg = 'SNPFinder Processing Job {} completed successfully.\n'.format(ARGS['job_id'])
        msg += 'Total processing time was {} and the job consumed {:,} normalized instance hours.\n'.format(running_time, normalized_instance_hours)
        msg += 'The EMR cluster consisted of {} and its spin-up time was {}.'.format(cluster_description, creation_time)
    else:
        job_string = failed_samples[0][0][0:-1].replace('_',' ').upper()
        if failed_samples[0][1] != 'task':
            problem_sample = failed_samples[0][1]
            msg = 'SNPFinder Processing Job {} terminated at {} due to a problem with the {} for sample {}.\n'.format(ARGS['job_id'], now, job_string, problem_sample)
        else:
            msg = 'SNPFinder Processing Job {} terminated at {} due to a problem with the {}.\n'.format(ARGS['job_id'], now, job_string)

        msg += 'Total processing time until job termination was {} and the job consumed {:,} normalized instance hours.\n'.format(running_time, normalized_instance_hours)
        msg += 'The EMR cluster consisted of {} and its spin-up time was {}.'.format(cluster_description, creation_time)

        if len(failed_samples) == 1:
            msg += '\n\nFailed job:\n'
        else:
            msg += '\n\nFailed jobs:\n'

        for job in failed_samples:

            job_string = job[0][0:-1].replace('_',' ').upper()
            if job[1] != 'task':
                job_string += ' for {}'.format(job[1])

            # Get the log URL for the failed job
            sql = text(
                """
                SELECT YARN_LOG_URL
                FROM EMR_JOB_STEPS
                WHERE EMR_JOB_ID=:p1 AND DESCRIPTION=:p2;
                """
            )
            values = {"p1": ARGS['job_id'], "p2": job_string}
            with ENGINE.connect() as conn:
                try:
                    log_url = conn.execute(sql, values).fetchone()[0]
                    msg += '\t{}, Log URL: {}\n'.format(job_string, log_url)
                except:
                    logger.warning('Problem executing SQL: {}'.format(sql))

        logger.info(msg)

    # Send the SNS message
    if len(failed_samples) == 0:
        subject = 'SNPFinder Pipeline Job {} Has Completed'.format(ARGS['job_id'])
    else:
        subject = 'SNPFinder Pipeline Job {} Has Failed'.format(ARGS['job_id'])

    client = boto3.client('sns')
    r = client.publish(
        TopicArn=j['sns_arn'],
        Message=msg,
        Subject= subject
    )

# Function to terminate an EMR cluster and delete intermediate files
def terminate_cluster(yarn_metrics_api, yarn_apps_api, failed_samples, command_list):
    # Wait for running jobs to complete
    while get_metrics(yarn_metrics_api)['active_jobs'] > 0:
        time.sleep(30)

    command_list = update_status(yarn_apps_api, command_list)

    # 1. Remove temporary files from EFS/S3 if job completed successfully
    if len(failed_samples) == 0:
        command = command_list['housekeeping_jobs']['task']['command']
        subprocess.call(command, shell=True)

    # 2. Update the database, setting status to ABORTED for tasks that were not executed
    for job_tag in command_list:
        if job_tag == 'housekeeping_jobs':
            tag = command_list['housekeeping_jobs']['task']['tag']
            if len(failed_samples) == 0:
                command_list = update_status(yarn_apps_api, command_list)
            else:
                update_task_status('ABORTED', tag, None)
        else:
            for task in command_list[job_tag]:
                if command_list[job_tag][task]['status'] not in ('SUCCEEDED', 'FAILED'):
                    tag = command_list[job_tag][task]['tag']
                    update_task_status('ABORTED', tag, None)

    # Update the final job status
    if len(failed_samples) == 0:
        final_status = 'SUCCEEDED'
    else:
        final_status = 'FAILED'

    update_job_status(ARGS['job_id'], final_status)

    # 3. Send SNS notification
    try:
        send_sns_notification(failed_samples)
    except:
        pass

    # 4. Send the Terminate Cluster message
    terminate_jobflow()

    # 5. Kill this EMR step
    sys.exit(0)

# Function to submit jobs to Yarn based based on whether resources are available and previous steps have been completed
def submit_job_to_yarn(command_list, job_tag, dependencies, yarn_metrics_api, yarn_apps_api, wait_for_all=False):
    # Update the cluster status in the Database
    update_job_status(ARGS['job_id'], 'RUNNING')
    g = get_metrics(yarn_metrics_api)
    while True:
        # Get the command list for tasks in this task group
        job_list = command_list[job_tag]

        if  len(job_list) > 0:
            # Get the number of containers pending (proxy for whether the cluster has the resources to execute a job)
            metrics = get_metrics(yarn_metrics_api)
            if metrics != g:
                logger.info('{}'.format(metrics))

            # Check if resources are available
            if (metrics['pending_containers'] == 0) and (metrics['available_memory'] > 2) and (metrics['active_jobs'] <= 25):
                # Get a list of samples in this step that still need to be processed
                pending_samples = [x for x in job_list if job_list[x]['status'] == 'INCOMPLETE']

                available_samples, failed_samples = ([], [])
                if wait_for_all == True:
                    # Don't return the list of available samples/tasks until all samples/tasks have finished processing
                    completed_samples, all_samples = ([],[])
                    for dependency in dependencies:
                        completed_samples += [x for x in command_list[dependency] if command_list[dependency][x]['status'] == 'SUCCEEDED']
                        all_samples += [x for x in command_list[dependency]]
                        failed_samples += [(dependency, x) for x in command_list[dependency] if command_list[dependency][x]['status'] == 'FAILED']

                    if len(completed_samples) != len(all_samples):
                        available_samples = []
                    else:
                        available_samples = pending_samples

                elif len(dependencies) > 0:
                    # Check to see which samples from the previous step(s) are done and available for further processing
                    for dependency in dependencies:
                        if len(command_list[dependency]) > 0:
                            available_samples += [x for x in pending_samples if (x in command_list[dependency]) and (command_list[dependency][x]['status'] == 'SUCCEEDED')]
                            failed_samples += [(dependency, x) for x in pending_samples if (x in command_list[dependency]) and (command_list[dependency][x]['status'] == 'FAILED')]
                        else:
                            available_samples = pending_samples
                else:
                    available_samples = pending_samples

                # Check to see if there are any failures. If so, wait for any running jobs to complete andF terminate cluster
                if len(failed_samples) > 0:
                    terminate_cluster(yarn_metrics_api, yarn_apps_api, failed_samples, command_list)

                # Check to see if any tasks are available to execute
                if len(available_samples) > 0:
                    # Get the task dictionary: {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}
                    task = job_list[available_samples[0]]

                    # Submit the task to Yarn
                    job_string = job_tag[:-1].replace('_', ' ').title()
                    logger.info('Submitting {} for {}'.format(job_string.upper(), available_samples[0]))
                    with open(os.devnull, 'w') as FNULL:
                        r = subprocess.Popen(task['command'], shell=True, stdout=FNULL, stderr=subprocess.STDOUT, close_fds=True)

                    # Update the status
                    command_list[job_tag][available_samples[0]]['status'] = 'SUBMITTED'
                # If there are not tasks to execute, and no tasks waiting on other tasks to complete, exit the loop
                elif len(pending_samples) == 0:
                    # Update status of the alignment jobs
                    command_list = update_status(yarn_apps_api, command_list)

                    # Once all tasks in this task group have been submitted, return to the calling function
                    break

            # Update status of the alignment jobs
            command_list = update_status(yarn_apps_api, command_list)
            g = metrics

            # Wait to check again
            time.sleep(15)
        else:
            break

    return command_list

# The main function to submit EMR tasks to the Yarn resource manager
def main():
    # Get connection to the database
    logger.debug('engine_path={}'.format(S3_PATHS["engine_path"]))

    # Set the job status to RUNNING
    update_job_status(ARGS['job_id'], 'RUNNING')

    # Get IP address of master node and build Yarn API urls from it
    r = requests.get('http://169.254.169.254/latest/meta-data/local-ipv4')
    ip_address = r.content.decode('utf-8')
    yarn_apps_api = 'http://{}:8088/ws/v1/cluster/apps'.format(ip_address)
    yarn_metrics_api = 'http://{}:8088/ws/v1/cluster/metrics'.format(ip_address)
    logger.debug('yarn_apps_api={}'.format(yarn_apps_api))
    logger.debug('yarn_metrics_api={}'.format(yarn_metrics_api))

    # Generate the list of spark-submit commands for the project
    command_list, sample_list = get_commands()
    # Command list has the following format:
    """
    command_list = {
        job_key: {
            sample_id: {'tag': tag, 'command': cmd, 'status': 'INCOMPLETE'}
        }
    }
    """
    logger.info('sample_list={}'.format(sample_list))
    logger.debug('command_list={}'.format(command_list))

    # Add commands to table
    command_list = load_commands_to_database(command_list)

    # Get SAMPLE_METRICS IDs by adding data to the PROJECT, SAMPLE, and SAMPLE_METRICS tables
    if ARGS['job_type'].lower() != 'discovery':
        # Check to see if the dictionary already exists (this is a re-run of a previous job)
        sample_dictionary_path = '/efs/{}/sample_dictionary.json'.format(ARGS["job_id"])
        logger.debug('sample_dictionary_path={}'.format(sample_dictionary_path))
        if not os.path.isfile(sample_dictionary_path):
            get_sample_metrics_ids(sample_list)
    else:
        cmd = 'sudo chown hadoop:hadoop /efs'
        subprocess.call(cmd, shell=True)
        project_path = '/efs/{}/'.format(ARGS["job_id"])
        logger.debug('project_path={}'.format(project_path))
        if not os.path.exists(project_path):
            os.makedirs(project_path, mode=0o777, exist_ok=True)
            chown(project_path, group='hadoop')
            os.chmod(project_path, 0o777)

    """ SUBMIT JOBS TO YARN"""
    # 1. Submit the alignment jobs
    command_list = submit_job_to_yarn(command_list, 'alignment_jobs', [], yarn_metrics_api, yarn_apps_api, wait_for_all=False)

    # 2. Submit the BAM merge job commands
    command_list = submit_job_to_yarn(command_list, 'bam_merge_jobs', ['alignment_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=False)

    # terminate job if it is aln_only mode
    if ('aln_only' in ARGS) and (ARGS['aln_only'] == True):
        failed_samples = [('bam_merge_jobs', x) for x in command_list['bam_merge_jobs'] if command_list['bam_merge_jobs'][x]['status'] == 'FAILED']
        terminate_cluster(yarn_metrics_api, yarn_apps_api, failed_samples, command_list)
        return

     # Check to see if the SNP-calling commands need to be executed
    if len(command_list['snp_calling_jobs']) > 0:
        # 3. Call SNPs
        command_list = submit_job_to_yarn(command_list, 'snp_calling_jobs', ['bam_merge_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=False)

        # 4. Merge intermediate SNP files
        command_list = submit_job_to_yarn(command_list, 'snp_merge_jobs', ['snp_calling_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=True)

        # 5. Build SNPset for Filling
        command_list = submit_job_to_yarn(command_list, 'create_snpset_jobs', ['snp_merge_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=True)

        # 6. Sort the SNPset
        command_list = submit_job_to_yarn(command_list, 'sort_snpset_jobs', ['create_snpset_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=True)

        # 7. Perform Filling
        command_list = submit_job_to_yarn(command_list, 'snarping_jobs', ['sort_snpset_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=True)
    else:
        # 7a. Sort the SNPset
        command_list = submit_job_to_yarn(command_list, 'sort_snpset_jobs', ['bam_merge_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=True)

        # 7b. Perform SNARPing
        command_list = submit_job_to_yarn(command_list, 'snarping_jobs', ['sort_snpset_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=True)

    # 8. Merge the intermediate SNARPing/Filling files to produce a Sample-level Parquet file
    command_list = submit_job_to_yarn(command_list, 'snarp_merge_jobs', ['snarping_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=True)

    # 9a. Write the Project-level Parquet file
    command_list = submit_job_to_yarn(command_list, 'project_parquet_jobs', ['snarp_merge_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=True)

    # 9b. Write the Metaproject-level Parquet file
    if command_list['metaproject_parquet_jobs']:
        command_list = submit_job_to_yarn(
            command_list,
            'metaproject_parquet_jobs',
            ['project_parquet_jobs'],
            yarn_metrics_api,
            yarn_apps_api,
            wait_for_all=True
        )
        last_job = 'metaproject_parquet_jobs'
    else:
        last_job = 'project_parquet_jobs'

    if ARGS['job_type'].lower() != 'discovery':
        # 10. Update the Locus data for the Reference assembly
        command_list = submit_job_to_yarn(command_list, 'locus_update_jobs', [last_job], yarn_metrics_api, yarn_apps_api, wait_for_all=True)

        # 11. Update the databases with metadata for the project_parquet_jobs
        command_list = submit_job_to_yarn(command_list, 'metadata_update_jobs', ['locus_update_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=True)

        # Check to see if the last job failed
        failed_samples = [('metadata_update_jobs', x) for x in command_list['metadata_update_jobs'] if command_list['metadata_update_jobs'][x]['status'] == 'FAILED']

        # 12. Submit new Loci to Lexicon for assignment of Lexicon IDs
        #command_list = submit_job_to_yarn(command_list, 'lexicon_request_jobs', ['metadata_update_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=True)

        # Check to see if the last job failed
        #failed_samples = [('lexicon_request_jobs', x) for x in command_list['lexicon_request_jobs'] if command_list['lexicon_request_jobs'][x]['status'] == 'FAILED']
    else:
        # 10. Write the Project-level Parquet  to CSV
        command_list = submit_job_to_yarn(command_list, 'csv_jobs', ['project_parquet_jobs'], yarn_metrics_api, yarn_apps_api, wait_for_all=True)

        failed_samples = [('csv_jobs', x) for x in command_list['csv_jobs'] if command_list['csv_jobs'][x]['status'] == 'FAILED']

    # Terminate the cluster
    terminate_cluster(yarn_metrics_api, yarn_apps_api, failed_samples, command_list)

if __name__ == "__main__":
    main()


import pymongo
import logging
import argparse
import sys
import inspect
import networkx as nx

logger = logging.getLogger(__name__)
logger_debug_frequency = 10000

def parseArgs():
    parser = argparse.ArgumentParser(description='Build a dictionary, where keys are users and value the GPS coordinates of their tweet.')
    parser.add_argument('--mongoServerHost', default='localhost', help='Mongo server host')
    parser.add_argument('--mongoServerPort', default=27017, type=int, help='Mongo server port')
    parser.add_argument('--mongoServerUsername', type=str, help='Mongo server username')
    parser.add_argument('--mongoServerPassword', type=str, help='Mongo server password')
    parser.add_argument('--mongoDatabase', '-m', required=True, help='name of the mongo database. Host is assumed to be localhost, port 27017.')
    parser.add_argument('--output-file', '-o', required=True, help='name of the output file')
    parser.add_argument('--log-level', '-l', default='info', help='logger verbosity level', choices=['info', 'debug'])
    parser.add_argument('--log-destination', '-L', default='file', help='logger destination', choices=['file', 'stderr'])
    return parser.parse_args(), parser

def main():
    args, parser = parseArgs()

    if args.log_level == 'debug':
        logger.setLevel(logging.DEBUG)
    else:
        logger.setLevel(logging.INFO)
    if args.log_destination == 'file':
        handler = logging.FileHandler(inspect.getfile(inspect.currentframe()).split('.')[0]+'.log')
    else:
        handler = logging.StreamHandler(sys.stderr)
    handler.setFormatter(logging.Formatter("%(asctime)s; %(levelname)s; %(message)s"))
    logger.addHandler(handler)

    if args.mongoServerUsername and args.mongoServerPassword:
        uri = "mongodb://%s:%s@%s:%d/%s"%(args.mongoServerUsername, args.mongoServerPassword, args.mongoServerHost,args.mongoServerPort, args.mongoDatabase)
    else:
        uri = "mongodb://%s:%d/%s"%(args.mongoServerHost,args.mongoServerPort, args.mongoDatabase)
    logger.info("Connecting to %s"%uri)
    db = pymongo.MongoClient(uri)[args.mongoDatabase]
    logger.info("Connected")

    net = nx.DiGraph()
    for c in db.communities.find():
        net.add_node(c['id'], nSpeakers=len(c['users']))
    correspondances = {int(l.split(' ')[1]): int(l.split(' ')[0]) for l in open(db.network.find_one({'files':{'$exists':True}})['files']['correspondances'])}
    logger.debug("loading users")
    users = {u['id'] : u['community'] for u in db.users.find({'community':{'$exists':True}},{'id':True, 'community':True})}
    logger.debug("reading file")
    logger.debug("counting lines")
    num_lines = sum(1 for line in open(db.network.find_one({'files':{'$exists':True}})['files']['network']))
    logger.debug("%d lines"%num_lines)
    i=0
    for l in open(db.network.find_one({'files':{'$exists':True}})['files']['network']):
        u0,u1=l.split(" ")
        n0=users[correspondances[int(u0)]]
        n1=users[correspondances[int(u1)]]
        if net.has_edge(n0,n1):
            net[n0][n1]['weight']+=1
        else :
            net.add_edge(n0,n1,weight=1)
        if i%logger_debug_frequency==0:
            logger.debug("%d/%d lines read"%(i,num_lines))
        i+=1
    logger.info("writing file: %s"%args.output_file)
    nx.write_gexf(net,args.output_file)
    logger.info("done.")

if __name__ == '__main__':
    main()


import itertools
import pdb
import re
import tensorflow as tf
from .base_dataset import VideoDataset
from collections import OrderedDict
from .softmotion_dataset import SoftmotionVideoDataset

class CartgripperVideoDataset(SoftmotionVideoDataset):
    def __init__(self, *args, **kwargs):
        VideoDataset.__init__(self, *args, **kwargs)
        if self.hparams.image_view != -1:
            self.state_like_names_and_shapes['images'] = \
                '%%d/image_view%d/encoded' % self.hparams.image_view, (48, 64, 3)
        else:
            # infer name of image feature(s)
            from google.protobuf.json_format import MessageToDict
            example = next(tf.python_io.tf_record_iterator(self.filenames[0]))
            dict_message = MessageToDict(tf.train.Example.FromString(example))
            feature = dict_message['features']['feature']
            image_names = set()
            for name in feature.keys():
                m = re.search('\d+/(\w+)/encoded', name)
                if m:
                    image_names.add(m.group(1))
            for image_name in image_names:
                m = re.search('image_view(\d+)', image_name)
                if m:
                    i = int(m.group(1))
                    suffix = '%d' % i if i > 0 else ''
                    self.state_like_names_and_shapes['images' + suffix] = \
                        '%%d/%s/encoded' % image_name, (48, 64, 3)

        if self.hparams.use_state:
            self.state_like_names_and_shapes['states'] = '%d/endeffector_pos', (self.hparams.sdim,)
        self.action_like_names_and_shapes['actions'] = '%d/action', (self.hparams.adim,)
        self._check_or_infer_shapes()

    def get_default_hparams_dict(self):
        default_hparams = super(CartgripperVideoDataset, self).get_default_hparams_dict()
        hparams = dict(
            context_frames=2,
            sequence_length=15,
            time_shift=3,
            use_state=True,
            sdim=6,
            adim=3,
            image_view=-1,
            autograsp=-1, # only take first n dimensions of action vector
            ignore_touch=False,
            saturate_touch=True,
            touch_no_state=False,
        )
        return dict(itertools.chain(default_hparams.items(), hparams.items()))

    def parser(self, serialized_example):
        state_like_seqs, action_like_seqs = super(CartgripperVideoDataset, self).parser(serialized_example)
        if 'actions' in action_like_seqs:
            if self.hparams.autograsp != -1:
                assert action_like_seqs['actions'].get_shape().as_list()[1] == 5
                action_like_seqs['actions'] = action_like_seqs['actions'][:,:self.hparams.autograsp]

        if 'states' in state_like_seqs:
            if self.hparams.ignore_touch:
                state_like_seqs['states'] = state_like_seqs['states'][:,:-2]

            if self.hparams.saturate_touch and not self.hparams.ignore_touch:
                assert state_like_seqs['states'].get_shape().as_list()[1] == 7
                state = state_like_seqs['states'][:,:5]
                touch = state_like_seqs['states'][:,5:]
                touch = tf.nn.sigmoid(touch)

                if self.hparams.touch_no_state:
                    state_like_seqs['states'] = touch
                else:
                    state_like_seqs['states'] = tf.concat([state, touch], axis=1)

        return state_like_seqs, action_like_seqs


import tkinter as tk

import matplotlib.pyplot as plt

from Architecture_roads.Architecture_matrix import Architecture_matrix
from Service.Service import Service

class SimpleTableInput(tk.Frame):
    def __init__(self, parent, rows, columns):
        tk.Frame.__init__(self, parent)

        self._entry = {}
        self.rows = rows
        self.columns = columns

        # register a command to use for validation
        vcmd = (self.register(self._validate), '%P')

        # create the table of widgets
        for row in range(self.rows):
            for column in range(self.columns):
                index = (row, column)
                e = tk.Entry(self, validate="key", validatecommand=vcmd)
                e.grid(row=row, column=column, stick="nsew")
                self._entry[index] = e
        # adjust column weights so they all expand equally
        for column in range(self.columns):
            self.grid_columnconfigure(column, weight=1)
        # designate a final, empty row to fill up any extra space

    def get_matrix(self):
        """
        Return a list of lists, containing the data in the table
        """
        result = []
        for row in range(self.rows):
            current_row = []
            for column in range(self.columns):
                index = (row, column)
                value = self._entry[index].get()
                if value == "":
                    current_row.append(0)
                elif value == 'Z':
                    current_row.append(-1)
                else:
                    current_row.append(int(value))
            result.append(current_row)
        return result

    def _validate(self, value_string):
        """
        Perform input validation.

        Allow only an empty value, pin Z or a value that can be converted to an integer
        """

        if value_string.strip() == "" or value_string == 'Z':
            return True
        try:
            i = int(value_string)
        except ValueError:
            self.bell()
            return False
        return True

class StepFour(tk.Frame):
    def __init__(self, parent, rows, columns):
        tk.Frame.__init__(self, parent)
        self.frows = rows
        self.fcolumns = columns
        self.table = SimpleTableInput(self, rows, columns)
        self.submit = tk.Button(self, text="Submit", command=self.on_submit)
        self.table.pack(side="top", fill="both", expand=True)
        self.submit.pack(side="bottom")

    def on_submit(self):
        matrix = self.table.get_matrix()
        architecture = Architecture_matrix(matrix, self.frows, self.fcolumns)
        service = Service(architecture)
        service.show_road_point()
        service.represent_roads()
        self.draw(matrix)

    def draw(self, matrix):
        cmap = 'rainbow'
        plt.matshow(matrix, cmap=cmap)
        plt.title(" Resulted Matrix ")
        plt.show()


###############################################################
#
# Author:  Mamata Anil Parab
# Project: Finding the frequency of given element from list
# Input:   10,70,30,10,80,90,50,10  Value=10
# Output:  3
#
###############################################################

def Frequency(arr,Value):

    iCount=0;
    for i in arr:
        if(i==Value):
            iCount=iCount+1;
    return iCount;

arr=list();

print("Enter count of elements in list");
Value=int(input());   # Accepting count of list

print("Enter actual elements of list");
for i in range(0,Value):
    No=input("");    # Accepting actual element of list
    arr.append(int(No));  # Inserting elements in list

print("Enter the element whose frequency need to be checked");
element=int(input())

Ret=Frequency(arr,element);      #Method call
print("Frequency of given element is {} ".format(Ret));


import pygame
import sys
from Panel import *

class Instruction():
    def __init__(self):
        self.header = Panel('INSTRUCTION', 500, 128, BLUE, None, 60)
        self.body = []
        self.body.append(Panel(
            '1.Press Escape to return to main menu', 0, 300, BLUE, None, 40))
        self.body.append(Panel(
            '2.Press s to hold your character,it cannot move', 0, 350, BLUE, None, 40))
        self.body.append(Panel(
            '3.Press r to reload your gun', 0, 400, BLUE, None, 40))

        self.instruction_group = pygame.sprite.Group()
        self.instruction_group.add(
            self.header, self.body)

    def check_key(self):
        for event in pygame.event.get():
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE or event.key == pygame.K_BACKSPACE:
                    return 'opening_game'
            if event.type == pygame.QUIT:
                sys.exit()

    def draw(self, screen):
        self.instruction_group.draw(screen)


# -*- coding: utf-8 -*-

import flask
from flask_login import login_user, logout_user, current_user, login_required, fresh_login_required, login_fresh
from forms.blacklist import EditForm, NewForm
from database import db, Role, Blacklist
from tools.Acl import Acl
from task import create_celery

__author__ = "Adam Schubert"
__date__ = "$26.7.2017 19:33:05$"

PER_PAGE = 20

blacklist = flask.Blueprint('blacklist', __name__)

@blacklist.route('/', methods=['GET'], defaults={'page': 1})
@blacklist.route('/page/<int:page>', methods=['GET'])
@login_required
@Acl.validate_path([Role.ADMIN], current_user)
def get_blacklist(page):
    pagination = Blacklist.query.filter().order_by(Blacklist.created.desc()).paginate(page, PER_PAGE)
    return flask.render_template('blacklist.html', pagination=pagination)

@blacklist.route('/new', methods=['GET', 'POST'])
@login_required
@Acl.validate_path([Role.ADMIN], current_user)
def new_blacklist():
    form = NewForm(flask.request.form)
    if flask.request.method == 'POST' and form.validate():
        blacklist_new = Blacklist()
        blacklist_new.dns = form.dns.data
        blacklist_new.date = form.date.data
        blacklist_new.bank_account = form.bank_account.data
        blacklist_new.note = form.note.data
        blacklist_new.ssl = form.ssl.data
        blacklist_new.signed = form.signed.data
        db.session.add(blacklist_new)
        db.session.commit()

        celery = create_celery(flask.current_app)
        res = celery.send_task('tasks.generate_thumbnail', args=(blacklist_new.id,))

        flask.flash('New blacklist item was added successfully.', 'success')
        return flask.redirect(flask.url_for('blacklist.get_blacklist'))

    return flask.render_template('blacklist_new.html', form=form)

@blacklist.route('/edit/<int:blacklist_id>', methods=['GET', 'POST'])
@Acl.validate_path([Role.ADMIN], current_user)
@login_required
def edit_blacklist(blacklist_id):
    blacklist_detail = Blacklist.query.filter_by(id=blacklist_id).first_or_404()

    form = EditForm(
        flask.request.form,
        dns=blacklist_detail.dns,
        id=blacklist_detail.id,
        date=blacklist_detail.date,
        bank_account=blacklist_detail.bank_account,
        note=blacklist_detail.note,
        ssl=blacklist_detail.ssl,
        signed=blacklist_detail.signed
    )
    if flask.request.method == 'POST' and form.validate():
        blacklist_detail.date = form.date.data
        blacklist_detail.bank_account = form.bank_account.data
        blacklist_detail.note = form.note.data
        blacklist_detail.dns = form.dns.data
        blacklist_detail.ssl = form.ssl.data
        blacklist_detail.signed = form.signed.data
        db.session.add(blacklist_detail)
        db.session.commit()
        celery = create_celery(flask.current_app)
        res = celery.send_task('tasks.generate_thumbnail', args=(blacklist_detail.id,))
        flask.flash('Domain was saved successfully.', 'success')
        return flask.redirect(flask.url_for('blacklist.get_blacklist'))

    return flask.render_template('blacklist_edit.html', form=form, blacklist_detail=blacklist_detail)

@blacklist.route('/delete/<int:blacklist_id>', methods=['GET'])
@login_required
@Acl.validate_path([Role.ADMIN], current_user)
def delete_blacklist(blacklist_id):
    blacklist_detail = Blacklist.query.filter_by(id=blacklist_id).first_or_404()
    db.session.delete(blacklist_detail)
    db.session.commit()
    flask.flash('Domain was deleted successfully.', 'success')

    return flask.redirect(flask.url_for('blacklist.get_domains'))


import pandas as pd
import spotipy
import numpy as np

sp = spotipy.Spotify() # Load the Spotipy Client

# Get specific artist object
artistId = '3TVXtAsR1Inumwj472S9r4' # Specify artist id: Drake
artist = sp.artist(artistId)
print(artist)

# Get specific album object
albumId = '40GMAhriYJRO1rsY4YdrZb' # Specific album id: Views
album = sp.album(albumId)
print(album)

# Get top 10 tracks for a given artist
top10 = sp.artist_top_tracks(artistId)
print(type(top10['tracks'])) # a list
top_track = top10['tracks'][0]
print(top_track) # Fake Love

# Get related artist based on "analysis of the Spotify community's listening history"
related_artists = sp.artist_related_artists(artistId)
print(type(related_artists['artists']))
print(related_artists['artists'][0]) # Most related artist

# Get specific track's features
trackId = top_track['id']
track = sp.track(trackId)
print(track)


# Network 7

from keras import Sequential
from keras.layers import Conv2D, Activation, MaxPooling2D, Flatten, Dense, Dropout
from keras.preprocessing.image import ImageDataGenerator

dimensions = [64, 64, 3]
batch_size = 32

model = Sequential()
model.add(Conv2D(32, 3, input_shape=(dimensions[0], dimensions[1], dimensions[2]),
                 data_format="channels_last"))
model.add(Activation("relu"))
model.add(MaxPooling2D(pool_size=2))

model.add(Conv2D(32, 3))
model.add(Activation("relu"))
model.add(MaxPooling2D(pool_size=2))

model.add(Conv2D(64, 3))
model.add(Activation("relu"))
model.add(MaxPooling2D(pool_size=2))

model.add(Conv2D(128, 3))
model.add(Activation("relu"))
model.add(MaxPooling2D(pool_size=2))

model.add(Flatten())
model.add(Dense(128))
model.add(Activation("relu"))
model.add(Dropout(0.5))

model.add(Dense(4))
model.add(Activation("softmax"))

model.compile(loss="categorical_crossentropy",
              optimizer="rmsprop",
              metrics=["accuracy"])

train_image_data_generator = ImageDataGenerator(
    rescale=1. / 255,
    shear_range=0.2,
    zoom_range=0.2
)

train_image_data_flow = train_image_data_generator.flow_from_directory(
    "./Images/Training",
    target_size=(dimensions[0], dimensions[1]),
    batch_size=batch_size,
    class_mode="categorical"
)

validation_image_data_generator = ImageDataGenerator(
    rescale=1. / 255
)

validation_image_data_flow = validation_image_data_generator.flow_from_directory(
    "./Images/Validation",
    target_size=(dimensions[0], dimensions[1]),
    batch_size=batch_size,
    class_mode="categorical"
)

sample_amount = 3399
validation_sample_amount = 1021
epochs = 50
epoch_steps = sample_amount // batch_size
validation_steps = validation_sample_amount // batch_size

history = model.fit_generator(
    train_image_data_flow,
    steps_per_epoch=epoch_steps,
    epochs=epochs,
    validation_data=validation_image_data_flow,
    validation_steps=validation_steps
)

print(history.history)

model.save("cnn_model_v12_50epoch.h5")


import logging

import bs4
from scrapers.utils import time_to_datetime, get_sha_hash, get_article, get_rss

logger = logging.getLogger("scraper.finance")

class FinanceScraper(object):
    FINANCE_RSS_URL = "http://www.finance.si/xml-servis/finance.rss"

    def parse_source(self, existing_ids=None):
        news = []
        feed_content = get_rss(self.FINANCE_RSS_URL)

        for feed_entry in feed_content.entries:
            link = feed_entry["link"]

            if existing_ids and get_sha_hash(link) in existing_ids:
                logger.debug("Skipping %s", link)
                continue

            published_date = time_to_datetime(feed_entry["published_parsed"])
            news.append((link, {"published": published_date}))

        return news

    def parse_article(self, article_url):
        link, data = article_url
        article = self.get_article_text(link)

        if article is None: return
        published_date = data["published"]
        article["published"] = published_date
        article["source"] = "Finance"
        article["source_url"] = link
        article["language"] = "si"
        article["id"] = get_sha_hash(link)
        return article

    def get_article_text(self, link):
        logger.debug("Grabbing article %s", link)
        article_html = get_article(link)
        result = {}
        result["raw_html"] = article_html
        article = bs4.BeautifulSoup(article_html)
        if article.body is None:
            return None

        title = article.body.find(class_="article-title")
        result["title"] = title.text.strip()

        author = article.body.find(class_="author")
        if author is not None:
            result["author"] = author.text.strip()
        else:
            result["author"] = None

        subtitle = article.body.find(class_="article-flash")
        if subtitle is not None:
            result["subtitles"] = [subtitle.string.strip()]

        content = article.body.find(class_="art-content")
        if content is None:
            return None
        else:
            result["text"] = u" ".join(content.stripped_strings)
            return result


import Tile
import Point
import Constants
import random
import math
from Terrains import terrains

class world(object):
    def __init__(self, start_size):
        self.tiles = dict()
        for x in range(start_size):
            for y in range(start_size):
                terrain = terrains(random.randint(1, 4))
                self.tiles[(x, y)] = Tile.tile(self, (x, y), terrain)

    def paint_world(self, center, terrain=terrains.grass, radius=1):
        pointLists = dict()

        cX, cY = center

        for x in range(-radius, radius + 1):
            pX = cX + x
            for y in range(-radius, radius + 1):
                pY = cY + y
                dist = math.sqrt(x ** 2 + y ** 2)
                if self.check_inside(pX, pY) and dist <= radius:
                    tileX, tileY = pX // Constants.TILE_SIZE, pY // Constants.TILE_SIZE
                    try:
                        pointLists[(tileX, tileY)].append(Point.point((pX, pY)))
                    except:
                        pointLists[(tileX, tileY)] = list()
                        pointLists[(tileX, tileY)].append(Point.point((pX, pY)))
        try:
            for tile, points in list(pointLists.items()):
                self.tiles[tile].paint_points(points, terrain)
        except:
            pass  # tile does not exist

    def check_inside(self, px, py):
        if px < 0 or py < 0 or px > Constants.WORLD_SIZE or py > Constants.WORLD_SIZE:
            return False
        return True


from setuptools import setup
import os
import sys
import zipfile
import urllib
from setuptools.command.develop import develop
from setuptools.command.install import install

if sys.version_info[0] >= 3:
    from urllib.request import urlretrieve
else:
    from urllib import urlretrieve

def fetch_viphreeqc():
    print("Fetching latest VIPhreeqc build from ci...")
    from_zip = False
    if sys.platform == "darwin":
        dll_name = "viphreeqc.dylib"
    elif "linux" in sys.platform:
        dll_name = "viphreeqc.so"
    else:
        dll_name = "viphreeqc.zip"

    dll_path = os.path.dirname(os.path.abspath(__file__)) + os.sep + "phreeqpython" + os.sep + "lib" + os.sep +dll_name
    print("Downloading to:", dll_path)

    urlretrieve("http://ci.abelheinsbroek.nl/"+dll_name, dll_path)

    if dll_name == "viphreeqc.zip":
        with zipfile.ZipFile(dll_path,"r") as archive:
            extract_path = os.path.dirname(os.path.abspath(__file__)) + "\phreeqpython\lib"
            print("Unpacking library to:", extract_path)
            archive.extractall(extract_path)
        os.remove(dll_path)

    print("Done!")

class PostDevelopCommand(develop):
    def run(self):
        fetch_viphreeqc()
        develop.run(self)

class PostInstallCommand(install):
    def run(self):
        fetch_viphreeqc()
        install.run(self)

setup(name='phreeqpython',
      version='0.2',
      description='Vitens viphreeqc wrapper and utilities',
      url='https://github.com/VitensTC/phreeqpython',
      author='Abel Heinsbroek',
      author_email='abel.heinsbroek@vitens.nl',
      license='Apache Licence 2.0',
      packages=['phreeqpython'],
      zip_safe=False,
      cmdclass={
          'develop': PostDevelopCommand,
          'install': PostInstallCommand
          },
      install_requires=['periodictable']
      )


# -*- coding: utf-8 -*-

import json
import os

from pyltp import SentenceSplitter
import os
from pyltp import Segmentor
from pyltp import Postagger

LTP_DATA_DIR = './ltp_data'  # ltp模型目录的路径
cws_model_path = os.path.join(LTP_DATA_DIR, 'cws.model')  # 分词模型路径，模型名称为`cws.model`
pos_model_path = os.path.join(LTP_DATA_DIR, 'pos.model')

segmentor = Segmentor()  #初始化实例
segmentor.load(cws_model_path) #加载实例
postagger = Postagger()
postagger.load(pos_model_path)

def load_and_split(comment_path):

    # -- 读入评论 --
    with open(comment_path,'r',encoding='utf-8') as f:
        data = f.read().strip().split("\n")

    # data的格式 [{},{},...{}]

    # -- 删除错误行 --
    wrong_index = []
    for i in range(len(data)):
        data[i] = json.loads(data[i])
        if data[i].__contains__("sku_id"):

            wrong_index.append(i)
    for i in range(len(wrong_index)):
        ind = wrong_index[i]-i
        del data[ind]

    if len(data)==0:
        return "empty"

    # -- NLP分词 写入文件--
    for item in data:
        comment = item['content']
        if comment=="":
            continue
        words = segmentor.segment(comment)
        words = list(words)
        words = [word.replace('.','').replace(',','').replace('?','').replace('!','').replace('@','').lower() for word in words]
        while '' in words:
            words.remove('')

        postags = postagger.postag(words)
        postags = list(postags)

        corpus_list = [words[i] for i in range(len(words)) if postags[i]!='wp']
        output.write(" ".join(corpus_list)+"\n")

folder_path = "./raw_comments/"
files= os.listdir(folder_path)
lens = len(files)
output = open('corpus.txt', 'w+',encoding = 'utf-8')

for i in range(lens):
    comment_path = folder_path+files[i]
    load_and_split(comment_path)

output.close()

segmentor.release()
postagger.release()


from django.views.generic import TemplateView, ListView, DetailView

from django.shortcuts import render_to_response, HttpResponse

from goods.models import Category, Goods

from django.core.mail import send_mail

from django.contrib import messages

from cart import Cart

class MainView(ListView):
    model = Category
    template_name = 'main.html'

    def post(self, request, *args, **kwargs):
        if request.POST['name'] and request.POST['phone'] and request.POST['message']:
            send_mail(
                'From DiyLab by {}'.format(request.POST['name']),
                '{}\nMy phone is: {}'.format(request.POST['message'], request.POST['phone']),
                request.POST['from_email'],
                ['qwaswexexortinvoke@gmail.com'],
                fail_silently=False,
            )
            messages.success(self.request, 'Message sent!')
        else:
            messages.warning(self.request, 'Form is incorrect!')
        return self.get(request, *args, **kwargs)

class GoodsView(ListView):
    model = Category
    template_name = 'goods.html'

class CategoryView(DetailView):
    model = Category
    template_name = 'category.html'

class GoodsDetailView(DetailView):
    model = Goods
    template_name = 'goods_detail.html'

class NewsView(TemplateView):
    template_name = 'news.html'

def add_to_cart(request, pk, quantity=1):
    product = Goods.objects.get(pk=pk)
    cart = Cart(request)
    cart.add(product, product.cost, quantity)
    return HttpResponse('OK')

def remove_from_cart(request, pk):
    product = Goods.objects.get(pk=pk)
    cart = Cart(request)
    cart.remove(product)
    return HttpResponse('OK')

def count(request):
    count = Cart(request).count()
    return HttpResponse(count)

def get_cart(request):
    return render_to_response('cart.html', dict(cart=Cart(request)))