# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import sevenbridges as sbg
from sevenbridges.errors import SbgError
from sevenbridges.http.error_handlers import rate_limit_sleeper, maintenance_sleeper
import sys
import re
import pdb
import concurrent.futures
from requests import request
import json
config = sbg.Config(profile='turbo')
api = sbg.Api(config=config, error_handlers=[rate_limit_sleeper, maintenance_sleeper])

# ## Create example bash calls from workflow
# Works best with non-restarted WF

project = "d3b-bixu/rs-vpf5jbc3-cov-irt-controlled-access-study"
task_id = "9aeaeaf3-0e59-4c92-a709-c6bd37431294"
out_file = open("/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/2020-06-22_VEP.tsv", "w")
# task_id = "3c20cc8e-18d7-43f2-bc2c-4a76d38a88f8"
task = api.tasks.get(task_id)
jobs = {}
temp = {}
for job in task.get_execution_details().jobs:
    if job.status == "COMPLETED":
        check = job.name.split('_')
        cmd = job.command_line
        if job.command_line == None:
            # pdb.set_trace()
            cmd = "embedded script or task retry"
            sys.stderr.write("WARN: Job " + job.name + " had null cmd\n")
        if check[-1] == "s":
            key = "_".join(check[:-2])
            if key not in temp:
                jobs[job.start_time] = {}
                jobs[job.start_time][key] = cmd
                temp[key] = 1
            else:
                temp[key] += 1
        else:
            jobs[job.start_time] = {}
            jobs[job.start_time][job.name] = cmd
out_file.write("Step\tType\tNum scatter\tCommand\n")
for rtime in sorted(jobs.keys()):
    for key in jobs[rtime]:
        rtype = "run step"
        sct = "NA"
        if key in temp and temp[key] > 1:
            rtype = "scatter"
            sct = str(temp[key])
        cmds = jobs[rtime][key].split('\n')
        for cmd in cmds:
            out_file.write(key + "\t" + rtype + "\t" + sct + "\t" + cmd + "\n")
out_file.close()


# ## Convert tsv to markdown table

# +
import sys

# max desired col width
max_w = 200
tsv_in = open("/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/2020-06-22_VEP.tsv")
out_md = open("/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/2020-06-22_VEP.md", "w")
data = []
max_len = []


for line in tsv_in:
    info = line.rstrip('\n').split('\t')
    data.append(info)
    if len(max_len) == 0:
        for item in info:
            max_len.append(len(item))
    else:
        for i in range(len(max_len)):
            if len(info[i]) > max_w:
                max_len[i] = max_w
            elif len(info[i]) > max_len[i]:
                max_len[i] = len(info[i])
# print header first
d_ct = []
for i in range(len(data[0])):
    d_ct.append(len(data[0][i]))
    out_md.write(" | " + data[0][i] + "".join([" "] * max_len[i]))
    d_ct[i] += max_len[i]
out_md.write(" |\n")
for i in range(len(data[0])):
    out_md.write(" | " + "".join(["-"] * d_ct[i]))
out_md.write(" |\n")
# pdb.set_trace()
for i in range(1, len(data), 1):
    for j in range(len(data[i])):
        d_ct = len(data[i][j]) + 2
        out_md.write(" | " + data[i][j] + "".join([" "] * max_len[j]))
        d_ct += max_len[j]
    out_md.write(" |\n")
out_md.close()


# -

# ## Get run times

# ### Get run times by step

def get_job_run_time(task, phrase):
    data = []
    if re.search(phrase, task.name):
        try:
            for job in task.get_execution_details().jobs:
                if job.status != "COMPLETED":
                    sys.stderr.write("Skipping job likely killed due to spot instance kill for " + job.name + " from task " + task.id + "\n")
                else:
                    data.append([job.name, str((job.end_time-job.start_time).seconds/3600)])
            # pdb.set_trace()
            hold=1
            return task.id, task.name, str(task.price.amount), str((task.end_time - task.start_time).seconds/3600), data
        except Exception as e:
            return [e, task.id]
    else:
        return []



# +
project = "d3b-bixu/rs-vpf5jbc3-cov-irt-controlled-access-study"
phrase = "VEP R100 ANNOTATE"

tasks = api.tasks.query(project=project, status="COMPLETED").all()
actual_out = open("/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/cost_est/VEP_cov-irt-actual_cost.txt", "w")
actual_out.write("Task name\tTask ID\tCost\tRun Time in hours\n")
step_run = open("/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/cost_est/VEP_cov-irt_step_run_times.txt", "w")
step_run.write("Run step\tRun time in hours\n")
# for task in tasks:
#     result = get_job_run_time(task, phrase)
#     if len(result) > 0:
#         pdb.set_trace()
#         actual_out.write("\t".join(result[0:4]) + "\n")
#         for step in result[4]:
#             step_run.write("\t".join(step) + "\n")
x = 0
m = 100
with concurrent.futures.ThreadPoolExecutor(16) as executor:
    results = {executor.submit(get_job_run_time, task, phrase): task for task in tasks}
    for result in concurrent.futures.as_completed(results):
        if len(result.result()) > 2:
            if x % m == 0:
                sys.stderr.write("Processed " + str(x) + " valid tasks\n")
            actual_out.write("\t".join(result.result()[0:4]) + "\n")
            for step in result.result()[4]:
                step_run.write("\t".join(step) + "\n")
            x += 1
        elif len(result.result()) == 2:
            sys.stderr.write(str(result.result()[0]) + "\tFailed processing task ID " + result.result()[1] + "\n")
            exit(1)
actual_out.close()
step_run.close()

# -

# ## Tag source files

def tag_file(info, header):
    try:
        meta = info.rstrip('\n').split('\t')
        f_obj = api.files.get(meta[0])
        metadata = {}
        for i in range(3, len(header), 1):
            metadata[header[i]] = meta[i]
        f_obj.metadata = metadata
        f_obj.save()
    except Exception as e:
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("Could not process " + info)
        exit(1)


project = "d3b-bixu/rs-vpf5jbc3-cov-irt-controlled-access-study"
manifest = open("/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/manifests/covwc_to_tag.txt")
head = next(manifest)
header = head.rstrip("\n").split("\t")
x = 1
m = 250
with concurrent.futures.ThreadPoolExecutor(16) as executor:
    results = {executor.submit(tag_file, line, header): line for line in manifest}
    for result in concurrent.futures.as_completed(results):
        if x % m == 0:
            sys.stderr.write('Processed ' + str(x) + ' files\n')
            sys.stderr.flush()


# ## Set up GATK4 RNAseq WF Tasks

def get_gatk_refs(api, project):
    try:
        ref_dict = {}
        ref_dict['reference_fasta'] = api.files.get('5eebc5d1e4b0a6d311357eb9')
        ref_dict['reference_dict'] = api.files.get('5eecb14ae4b0efd899f474da')
        known_sites = []
        known_sites.append(api.files.get('5eecd8c3e4b0efd899f4ae44'))
        known_sites.append(api.files.get('5eecd846e4b0efd899f4ae27'))
        known_sites.append(api.files.get('5eecd846e4b0efd899f4ae26'))
        known_sites.append(api.files.get('5eecd846e4b0efd899f4ae21'))
        ref_dict['knownsites'] = known_sites
        ref_dict['call_bed_file'] = api.files.get('5eebdeece4b0efd899f43eaa')
        ref_dict['dbsnp_vcf'] = api.files.get('5eecd846e4b0efd899f4ae24')
        ref_dict['tool_name'] = 'STAR_GATK4'
    except Exception as e:
        sys.stderr.write(str(e) + "\nFailed to get REFS\n")
        exit(1)
    return ref_dict


def draft_task(in_file):
    try:
        input_dict = {}
        for key in ref_obj:
            input_dict[key] = ref_obj[key]
        info = in_file.rstrip('\n').split('\t')
        input_dict['STAR_sorted_genomic_bam'] = api.files.get(info[0])
        task_name = "GATK RNAseq SNV: " + info[3]
        task = api.tasks.create(name=task_name, project=project, app=app_name, inputs=input_dict, run=False)
        task.inputs['output_basename'] = task.id
        task.save()
    except Exception as e:
        sys.stderr.write(str(e) + "\nfailed to set up task for " + in_file)
        exit(1)


project = 'd3b-bixu/rs-vpf5jbc3-cov-irt-controlled-access-study'
app_name = project + "/d3b-gatk-rnaseq-snv-wf"
manifest = open("/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/manifests/bams_for_gatk_to_run.tsv")
head = next(manifest)
ref_obj = get_gatk_refs(api, project)
x = 1
m = 250
with concurrent.futures.ThreadPoolExecutor(16) as executor:
    results = {executor.submit(draft_task, line ): line for line in manifest}
    for result in concurrent.futures.as_completed(results):
        if x % m == 0:
            sys.stderr.write('Processed ' + str(x) + ' tasks\n')
            sys.stderr.flush()


# ## Run VEP

def get_vep_refs(api, project):
    try:
        ref_dict = {}
        ref_dict['reference'] = api.files.get('5eebc5d1e4b0a6d311357eb9')
        ref_dict['cache'] = api.files.get('5eed0f54e4b0efd899f4afda')
        ref_dict['merged_cache'] = True
        ref_dict['tool_name'] = 'STAR_GATK4'
    except Exception as e:
        sys.stderr.write(str(e) + "\nFailed to get REFS\n")
        exit(1)
    return ref_dict


def draft_vep_task(in_file):
    try:
        input_dict = {}
        for key in ref_obj:
            input_dict[key] = ref_obj[key]
        info = in_file.rstrip('\n').split(',')
        input_dict['input_vcf'] = api.files.get(info[0])
        task_name = "VEP R100 ANNOTATE: " + info[3]
        task = api.tasks.create(name=task_name, project=project, app=app_name, inputs=input_dict, run=False)
        task.inputs['output_basename'] = task.id
        task.save()
    except Exception as e:
        sys.stderr.write(str(e) + "\nfailed to set up task for " + in_file)
        exit(1)


project = 'd3b-bixu/rs-vpf5jbc3-cov-irt-controlled-access-study'
app_name = project + "/vep-1oo-annotate"
manifest = open("/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/manifests/vcf_to_annotate-manifest.csv")
head = next(manifest)
ref_obj = get_vep_refs(api, project)
x = 1
m = 250
with concurrent.futures.ThreadPoolExecutor(16) as executor:
    results = {executor.submit(draft_vep_task, line ): line for line in manifest}
    for result in concurrent.futures.as_completed(results):
        if x % m == 0:
            sys.stderr.write('Processed ' + str(x) + ' tasks\n')
            sys.stderr.flush()


# ## Copy metadata to outputs

def add_metadata_to_outputs(task, phrase, in_key):
    if re.search(phrase, task.name):
        sys.stderr.write('Valid task found ' + task.name + '\n')
        metadata = task.inputs[in_key].metadata
        for out_key in task.outputs:
            # pdb.set_trace()
            try:
                if type(task.outputs[out_key]) is not list:
                    file_obj = api.files.get(task.outputs[out_key].id)
                    for key in metadata:
                        file_obj.metadata[key] = metadata[key]
                    file_obj.save()
                else:
                    for output in task.outputs[out_key]:
                        if type(output) is not list:
                            file_obj = api.files.get(output.id)
                            for key in metadata:
                                file_obj.metadata[key] = metadata[key]
                            file_obj.save()
                        else:
                            for item in output:
                                if item is not None:
                                    file_obj = api.files.get(item.id)
                                    for key in metadata:
                                        file_obj.metadata[key] = metadata[key]
                                    file_obj.save()

            except Exception as e:
                print(e)
                print("Skipping " + out_key + " for " + task.name + " due to error")


prefix = 'VEP R100 ANNOTATE'
project = "d3b-bixu/rs-vpf5jbc3-cov-irt-controlled-access-study"
key = '<KEY>'
print("You sure tag outputs with task prefix: " + prefix + "? Type \"YASS\" if so")
check = input()
if check == "YASS":
    tasks = api.tasks.query(project=project, status="COMPLETED").all()
    for task in tasks:
        add_metadata_to_outputs(task, prefix, key)
else:
    sys.stderr.write("User did not type YASS, skipping\n")


# ## Get task outputs

def write_to_manifest(out_fh, file_obj, out_key, task_name):
    out_fh.write(",".join([file_obj.id, file_obj.name, out_key, task_name]) + "\n")


# +
project = "d3b-bixu/rs-vpf5jbc3-cov-irt-controlled-access-study"
tasks = api.tasks.query(project=project, status="COMPLETED").all()
out = open('/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/manifests/test_vep_out.txt', 'w')
out.write("id,name,output_category,task_name\n")
phrase = "VEP R100 ANNOTATE: COVHA-20200311-P1-A01-CP"

for task in tasks:
    if re.search(phrase, task.name):
        sys.stderr.write('Processing task: ' + task.name + "\n")
        for out_key in task.outputs:
#             try:
            if type(task.outputs[out_key]) is not list:
                file_obj = task.outputs[out_key]
                write_to_manifest(out, file_obj, out_key, task.name)
                if task.outputs[out_key].secondary_files is not None:
                    write_to_manifest(out, task.outputs[out_key].secondary_files[0], out_key, task.name)
            else:
                for i in range(len(task.outputs[out_key])):
                    if type(task.outputs[out_key][i]) is not list:
                        write_to_manifest(out, task.outputs[out_key][i], out_key, task.name)
                        if task.outputs[out_key][i].secondary_files is not None:
                            write_to_manifest(out, task.outputs[out_key][i].secondary_files[0], out_key, task.name)
                    else:
                        for j in range(len(task.outputs[out_key][i])):
                            if task.outputs[out_key][i][j] is not None:
                                write_to_manifest(out, task.outputs[out_key][i][j], out_key, task.name)
                                if task.outputs[out_key][i][j].secondary_files is not None:
                                    write_to_manifest(out, task.outputs[out_key][i][j].secondary_files[0], out_key, task.name)
#             except Exception as e:
#                 print(e)
#                 print("Skipping " + out_key + " for " + task.name + " due to error")
out.close()

# -

# ## Abort unresponsive and restart tasks

# +
import datetime
import pytz

project = "d3b-bixu/rs-vpf5jbc3-cov-irt-controlled-access-study"
tasks = api.tasks.query(project=project, status="RUNNING")
current = datetime.datetime.now()
tz = pytz.timezone('America/New_York')
prefix = "GATK RNAseq SNV"
task_abort = open("/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/TASK_RUN/aborted_and_restarted2.log", 'w')
task_abort.write("Task ID\tTask name\tNew Task ID")
for task in tasks:
    if re.search(prefix, task.name):
        for job in task.get_execution_details().jobs:
            if job.name == "preprocess_rnaseq_bam_sambamba_md_sorted":
                if job.status == "RUNNING":
                    diff = (current-pytz.utc.localize(job.start_time, is_dst=None).astimezone(tz).replace(tzinfo=None)).seconds/3600
                    if diff > 2:
                        task_abort.write(task.id + "\t" + task.name)
                        in_dict = {}
                        sys.stderr.write("Aborting " + task.id + "\t" + task.name + "\n" )
                        task.abort()
                        new_task = task.clone(run=False)
                        new_task.inputs['output_basename'] = new_task.id
                        new_task.save()
                        task_abort.write("\t" + new_task.id + "\n")
                        task_abort.flush()
                    else:
                        break
                else:
                    break
task_abort.close()
# -

# ## Restart failed tasks

project = "d3b-bixu/rs-vpf5jbc3-cov-irt-controlled-access-study"
tasks = api.tasks.query(project=project, status="FAILED")
prefix = "GATK RNAseq SNV"
task_restart = open("/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/TASK_RUN/failed_and_restarted2.log", 'w')
task_restart.write("Task ID\tTask name\tNew Task ID\n")
run_list = ["GATK RNAseq SNV: COVHA-20200315-P9-F02-N",
             "GATK RNAseq SNV: COVHA-20200316-P12-F01-P","GATK RNAseq SNV: COVHA-20200403-P1-C06-P"]
for task in tasks:
    if re.search(prefix, task.name) and task.name in run_list:
        task_restart.write(task.id + "\t" + task.name)
        new_task = task.clone(run=False)
        new_task.inputs['output_basename'] = new_task.id
        new_task.save()
        task_restart.write("\t" + new_task.id + "\n")
        task_restart.flush()
task_restart.close()


# ### Get Failed list

project = "d3b-bixu/rs-vpf5jbc3-cov-irt-controlled-access-study"
tasks = api.tasks.query(project=project, status="FAILED")
prefix = "GATK RNAseq SNV"
task_failed = open("/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/TASK_RUN/failed.log", 'w')
task_failed.write("Task ID\tTask name\n")
for task in tasks:
    if re.search(prefix, task.name):
        task_failed.write(task.id + "\t" + task.name + "\n")
task_failed.close()

# ### Remove outputs from failed and aborted tasks

project = "d3b-bixu/rs-vpf5jbc3-cov-irt-controlled-access-study"
tasks = api.tasks.query(project=project, status="ABORTED")
prefix = "GATK RNAseq SNV"
print("You sure remove outputs from failed tasks with prefix: " + prefix + "? Type \"YASS\" if so")
check = input()
if check == "YASS":

    for task in tasks:
        if re.search(prefix, task.name):
            for key in task.outputs:
                if task.outputs[key] is not None:
                    sys.stderr.write("Found files to remove from failed task: " + task.id + " " + task.name + "\n")
                    try:
                        if task.outputs[key].secondary_files is not None:
                            sys.stderr.write("Removing secondary files\n")
                            for i in range(0, len(task.outputs[key].secondary_files), 1):
                                task.outputs[key].secondary_files[i].delete()
                    except Exception as e:
                        sys.stderr.write(str(e) + "\nFile with key " + key + " probably does not have secondaryFiles, skipping\n")
                    try:
                        task.outputs[key].delete()
                    except Exception as e:
                        sys.stderr.write(str(e) + "\nFile with key " + key + " was probably deleted before, skipping\n")            
                    sys.stderr.write("Finished processing " + task.id + "\n")


# ### Rename _\d_ files

manifest = open("/Users/brownm28/Documents/2020-Apr-8_RNAseq_snv_dev/manifests/to_rename-manifest.csv")
head = next(manifest)
print("You sure you want to rename the files in that manifest? Type \"YASS\" if so")
check = input()
if check == "YASS":
    for line in manifest:
        info = line.split(',')
        cur = api.files.get(info[0])
        new_name = cur.name[3:]
        sys.stderr.write("Renaming file with ID " + cur.id + " " + cur.name + " to " + new_name + "\n")
        cur.name = new_name
        cur.save()



StarcoderdataJupyterScriptsDedupFiltered