# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/lsf_interactive.ipynb.

# %% auto 0
__all__ = ['lsf_job', 'submit', 'update_status', 'get_result', 'LSF']

# %% ../nbs/lsf_interactive.ipynb 3
import glob
import gzip
import json
import logging
import os
import pickle
import re
import shutil
import subprocess
import time
import multiprocessing
from collections import Counter

from fastcore.script import call_parse

# %% ../nbs/lsf_interactive.ipynb 7
class lsf_job:
    def __init__(
        self,
        bsub_args,
        job_name,
        output_file,
        args,
        job_id=None,
        status="unsend",
        iteration=1,
        iteration_lim=3,
        exit_code=None,
        read_func=None
    ):
        self.bsub_args = bsub_args
        self.job_name = job_name
        self.output_file = output_file
        self.args = args
        self.job_id = job_id
        self.stat = status
        self.iteration = iteration
        self.iter_lim = iteration_lim
        self.result = None
        self.stderr = None
        self.stdout = None
        self.exit_code = exit_code
        self.read_func = read_func

# %% ../nbs/lsf_interactive.ipynb 8
def submit(tup):
    bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func = tup
    cmd = "bsub -q {} -M {} -n {} -J {} -o {} -e {} {} {} {}".format(
        bsub_args["queue"],
        int(bsub_args["mem"] + (bsub_args["mem"] * (iteration - 1))),
        bsub_args["n_cpus"],
        job_name,
        f'{bsub_args["output_dir"]}/{job_name}.stdout',
        f'{bsub_args["output_dir"]}/{job_name}.stderr',
        bsub_args["interpreter"],
        bsub_args["script"],
        " ".join(map(str, args)),
    )
    process = subprocess.Popen(
        cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    output, err = process.communicate()
    
    if err != b"":
        # raise ValueError("error in sending job")
        stat = "unsend"
    else:
        job_id = re.split("[<>]", output.decode())[1]
        exit_code = None
        stat = f"send it:{iteration}"
    return bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func

# %% ../nbs/lsf_interactive.ipynb 9
def update_status(tup):
    bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func = tup
    if stat == "unsend":
        bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func = submit((bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func))
        return bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func
    if stat == "done":
        return None
    cmd = f"bjobs {job_id}"
    process = subprocess.Popen(
        cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    output, err = process.communicate()
    stat = (
        dict(list(zip(*[x.split() for x in output.decode().split("\n")][:2])))
        .get("STAT","EXIT")[:4]
        .lower()
    )
    if stat in [None, "exit"]:
        stdout_file = f'{bsub_args["output_dir"]}/{job_name}.stdout'
        if os.path.isfile(stdout_file) == False:
            stat = False
            return
        with open(stdout_file) as h:
            rev = list(reversed(list(h)))
            for ix, x in enumerate(rev):
                if "Sender" in x:
                    stat = rev[ix - 1].split()[-1][:4].lower()
                    break
    if stat == "done":
        bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func = get_result((bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func))
    if stat == "exit":
        for ix, x in enumerate(rev):
            if "Exited with exit code" in x:
                exit_code = int(float(x.split()[-1]))
                if (
                    "TERM_MEMLIMIT" in rev[ix + 1]
                ):  # exit_code == 1 and iter < iter_lim:
                    iteration += 1
                    bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func = submit((bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func))
                    # stat = None
                break
    if exit_code not in [None]:
        stderr_file = f'{bsub_args["output_dir"]}/{job_name}.stderr'
        if os.path.isfile(stderr_file) == False:
            stderr = False
            return
        with open(stderr_file) as h:
            stderr = list(h)
    return bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func

# %% ../nbs/lsf_interactive.ipynb 10
def get_result(tup):
    bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func = tup
    result = None
    if read_func==None:
        pickle_file = output_file
        if os.path.isfile(pickle_file) == False:
            stat = False
            return output_file,stat,result
        with open(pickle_file, "rb") as h:
            result = pickle.load(h)
    else:
        result = read_func(output_file)
    return bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func

# %% ../nbs/lsf_interactive.ipynb 11
class LSF:
    " send a series of lsf jobs, collect the objects, and clean. Output objects must be a pickle "

    def __init__(
        self,
        script,
        jobs,
        interpreter="python3",
        monitor_file="/homes/fragoso/generic_monitor_file.txt",
        output_dir="/hps/nobackup/rdf/metagenomics/research-team/santiago/.std.generic",
        mem=1000,
        n_cpus=1,
        queue="research",
        n_jobs =1,
    ):
        self.monitor_file = monitor_file
        self.bsub_args = {
            "script": script,
            "interpreter": interpreter,
            "output_dir": output_dir,
            "mem": mem,
            "n_cpus": n_cpus,
            "queue": queue,
        }
        self.n_jobs = n_jobs
        self.jobs_ = jobs
        self.jobs = {job_name:lsf_job(self.bsub_args, job_name, output_file, args)
            for job_name, output_file, args in self.jobs_
                    }
        try:
            os.mkdir(output_dir)
        except FileExistsError as err:
            print(err)
            print("Force gather mode")
            for _,j in self.jobs.items():
                j.stat=None

    def update(self,read_func=None,loop=False,verbose=1):
        summary_ = {'start'}
        it = 0
        while set(summary_)!={'done'}:
            chunk = [(j.bsub_args, j.job_name, j.output_file, j.args, j.job_id, j.stat, j.iteration, j.iter_lim, j.result, j.stderr, j.stdout, j.exit_code, j.read_func if read_func==None else read_func )
                        for k,j in self.jobs.items() if j.stat!="done"]
            if verbose>1:
                print(f"processing {len(chunk)} jobs")
            step = 0
            size = 500
            while step*size<len(chunk):
                with multiprocessing.Pool(self.n_jobs) as pool:
                    res = pool.map(update_status,chunk[step*size:(step+1)*size])
                for bsub_args, job_name, output_file, args, job_id, stat, iteration, iter_lim, result, stderr, stdout, exit_code, read_func in res:
                    self.jobs[job_name].bsub_args = bsub_args
                    self.jobs[job_name].job_name = job_name
                    self.jobs[job_name].output_file = output_file
                    self.jobs[job_name].args = args
                    self.jobs[job_name].job_id = job_id
                    self.jobs[job_name].stat = stat
                    self.jobs[job_name].iteration = iteration
                    self.jobs[job_name].iter_lim = iter_lim
                    self.jobs[job_name].result = result
                    self.jobs[job_name].stderr = stderr
                    self.jobs[job_name].stdout = stdout
                    self.jobs[job_name].exit_code = exit_code
                    self.jobs[job_name].read_func = read_func
                if verbose>1:    
                    print(f"step {step} done")
                step += 1
                del res

            summary = dict(Counter([job.stat for _,job in self.jobs.items()]))
            summary_ = dict(summary)
            if verbose>1:
                print(summary)
            if loop==False:
                sumary_= {'done'}
            else:
                if verbose>0:
                    # print(f"",end='\r',flush=True)
                    print(f"iter: {it}\t{str(summary)}\t\t\t",end='\r',flush=True)
                time.sleep(20)
                it += 1
        return summary

    def clean(self):
        shutil.rmtree(self.bsub_args["output_dir"])

    def retake(self, code):
        " keep gathering results after the original kernel is dead "
        pass

