#!/usr/bin/env python
from __future__ import print_function # for python2

import sys
import os, os.path
import shutil
import time
import string
import subprocess
import re
import copy
from . import options

job_mgr_opts = options.Options("job_manager")
job_mgr_opts.add("template", "job",
                 "Filename for job template to look for in search path")
job_mgr_opts.add("templatepath", None, "Full path to a job template", str)
job_mgr_opts.add("resumetemplate", "resumejob",
                 "Filename for resume job template to look for in search path")
job_mgr_opts.add("resumetemplatepath", None, "Full path to a resume job template", str)
job_mgr_opts.add("multitemplate", "multijob",
                 "Base filename for multijob templates to look for in search path")
job_mgr_opts.add("multitemplatepath", None, "Full path to the base name for multijob templates", str)
job_mgr_opts.add("resumemultitemplate", "resumemultijob",
                 "Base filename for resume multijob templates to look for in search path")
job_mgr_opts.add("resumemultitemplatepath", None, "Full path to the base name for resume multijob templates", str)
job_mgr_opts.add("createjob", False, "Whether to create new job directory")
#job_mgr_opts.add("resumejob", 0, "Whether to resume a previously checkpointed job", int)
job_mgr_opts.add("jobdir", "run", "Job directory")
#job_mgr_opts.add("resumedir", "run", "Directory containing checkpointed files", str)
job_mgr_opts.add("numproc", 1, "Number of processes")
job_mgr_opts.add("procspernode", 1, "Number of processes per node")
job_mgr_opts.add("submit", False, "Whether to immediately submit job")
job_mgr_opts.add("queue", None, "Batch system queue", str)
job_mgr_opts.add("account", None, "Batch system account", str)
job_mgr_opts.add("run", False, "Whether to immediately run job")
job_mgr_opts.add("setupsh", '/usr/local/share/synergia/setup.sh',
                 "Path to Synergia2 setup.sh file")
job_mgr_opts.add("overwrite", False, "Whether to overwrite existing job directory")
job_mgr_opts.add("walltime", None, "Limit job to given wall time", str)
job_mgr_opts.add("synergia_executable", "synergia", "Name or path of synergia executable")
job_mgr_opts.add("synergia_resume_executable", "synergia-pyresume", "Name or path of synergia resume executable")
job_mgr_opts.add("jmverbosity", 1, "Job manager verbosity")

#job_mgr_opts.add("checkpoint", 0, "enable generation of checkpoint", int)
#job_mgr_opts.add("checkpoint_freq", 100, "frequency of checkpoint generation", int)

def get_synergia_directory(die_on_failure=True):
    if "SYNERGIA2DIR" in os.environ:
        synergia_dir = os.environ["SYNERGIA2DIR"]
    else:
        if os.path.isfile("./local_paths.py"):
            synergia_dir = "."
        else:
            if os.path.isfile("../local_paths.py"):
                synergia_dir = ".."
            else:
                print("Unable to determine Synergia directory. Defaults")
                print("are \".\" and \"..\". Set the environment variable")
                print("SYNERGIA2DIR to point to the directory containing")
                print("local_paths.py.")
                if die_on_failure:
                    sys.exit(1)
                else:
                    synergia_dir = None
    if synergia_dir:
        synergia_dir = os.path.abspath(synergia_dir)
    return synergia_dir

def get_script_templates_dir():
    if 'SYNERGIA2TEMPLATES' in os.environ:
        return os.environ['SYNERGIA2TEMPLATES']
    else:
        return os.path.join(get_synergia_directory(),
                            'synergia-script-templates')

def search_job_manager_paths(filename):
    found_path = None
    here = os.path.join(os.getcwd(), filename)
    if os.path.exists(here):
        found_path = here
    if not found_path:
        if 'SYNERGIA2TEMPLATES' in os.environ:
            env_path = os.path.join(os.environ['SYNERGIA2TEMPLATES'], filename)
            if os.path.exists(env_path):
                found_path = env_path
    if not found_path:
        if 'HOME' in os.environ:
            home_path = os.path.join(os.environ['HOME'],
                                     ".local",
                                     "share",
                                     "synergia", filename)
            if os.path.exists(home_path):
                found_path = home_path
    if not found_path:
        synergia_path = os.path.join(get_synergia_directory(),
                                     'synergia-script-templates',
                                     filename)
        if os.path.exists(synergia_path):
            found_path = synergia_path
    if found_path:
        found_path = os.path.abspath(found_path)
    return found_path

def get_default_script_templates_dir():
    return os.path.join(get_synergia_directory(), 'synergia-script-templates')

def add_local_opts(verbose):
    local_options_path = search_job_manager_paths('local_opts.py')
    if local_options_path:
        sys.path.insert(0, os.path.dirname(local_options_path))
        import local_opts
        sys.path.pop(0)
        found_options = False
        for name in dir(local_opts):
            is_options = False
            object = getattr(local_opts, name)
            if hasattr(object, "is_options"):
                is_options = object.is_options
            if is_options:
                job_mgr_opts.add_suboptions(object)
                found_options = True
            if hasattr(object, "is_override"):
                if object.is_override:
                    job_mgr_opts.override(object)
        if verbose and (job_mgr_opts.jmverbosity > 0):
            print('read local options from:', local_options_path)
        if not found_options:
            if verbose:
                print('warning: no options object(s) found in', local_opts.__file__)

def expand_multiple(arg):
    s = arg.split('=')
    lhs = s[0]
    rhs = s[1][1:-1]
    if rhs.find(',,') > 0:
        rhss = rhs.split(',,')
    else:
        rhss = rhs.split(',')
    return [lhs + '=' + x for x in rhss]

def is_multiple(arg):
    retval = False
    s = arg.split('=')
    if len(s) > 1:
        # if arg is "option=" then s[1] is None
        if s[1] and s[1][0] == '_':
            if s[1][-1] == '_':
                retval = True
            else:
                sys.stderr.write('argument "%s" missing closing _\n' % arg)
                sys.exit(1)
    return retval

def has_multiple(argv):
    retval = False
    for arg in argv:
        if is_multiple(arg):
            retval = True
    return retval

# take a single argv (list) and returns a lists of argvs
def expand_multiples(argv):
    retval = [[]]
    for i in range(0,len(argv)):
        if is_multiple(argv[i]):
            vals = expand_multiple(argv[i])
            retval = [prev + [val] for prev in retval
                        for val in vals]
        else:
            retval = [prev + [argv[i]] for prev in retval]
    return retval

# extra_opt_files and extra_opt_dirs will be copied if present, but ignored if missing
class Job_manager:
    '''The main Synergia workflow class.

    :param script: the name of the Synergia Python script or executable.
    :param opts: an Options object containing the job options. The internally-defined job manager options
        will be added as a suboption to opts.
    :param extra_files: a list of additional files required to run the job.
    :param extra_dirs: a list of additional directories whose contents are required to run the job.
    :param argv: the command-line arguments. Use sys.argv if :code:`None`.
    :param extra_opt_files: a list of additional files to be included iff found.
    :param extra_opt_dirs:  list of additional directories to be included iff found.
    :param standalone: for internal use.
    :param subjob: for internal use.
    :param subjob_index: for internal use.

    To list extra files or directories that are part of an Options object opt, use, e.g.,
        extra_files=[...,[opt,"name"],...]
    to get the option named "name".
    '''
    def __init__(self, script, opts, extra_files=None, extra_dirs=None,
                 argv=None,
                 extra_opt_files=None, extra_opt_dirs=["lattice_cache"],
                 standalone=False, subjob=False, subjob_index=0):
        if argv == None:
            argv = sys.argv
        self.subjob = subjob
        self.subjob_index = subjob_index
        if has_multiple(argv):
            self.submanagers = []
            index = 0
            for sub_argv in expand_multiples(argv):
                self.submanagers.append(Job_manager(script, copy.deepcopy(opts),
                                                    extra_files, extra_dirs, sub_argv,
                                                    extra_opt_files, extra_opt_dirs, standalone,
                                                    True, index))
                index += 1
        else:
            self.submanagers = []
        self.opts = opts
        creating_job = False
        if len(sys.argv) > 1:
            if sys.argv[1] == "--strip-options-file":
                creating_job = True
        if not self.subjob:
            add_local_opts(creating_job)
        self.opts.add_suboptions(job_mgr_opts)
        self.directory = None
        self.subdirectory = "sub%02d" % self.subjob_index
        self.real_script = os.path.abspath(script)
        options_file = os.path.splitext(script)[0] + '_options.py'
        if os.path.exists(options_file):
            self.real_options_file = os.path.abspath(options_file)
        else:
            self.real_options_file = None
        self.argv = argv
        self.standalone = standalone
        if self.standalone:
            job_mgr_opts.add("executable", script, "Name or path of standalone executable", str)
        if len(self.argv) > 2:
            if (self.argv[2] == "--create-cxx-options-source"):
                self.create_cxx_options_source()
                sys.exit(0)
        if len(self.argv) > 1:
            if (self.argv[1] == "--create-cxx-options-source"):
                self.create_cxx_options_source()
                sys.exit(0)
            if self.argv[1] == "--strip-options-file":
                self.argv = self.argv[2:]
        self.synergia_dir = get_synergia_directory()

        self.opts.parse_argv(self.argv)

        # if we are resuming a checkpointed job, get the absolute
        # path of the resumedir to pass to the resumed job.  I don't check
        # whether the resumedir exists because I might be submitting a
        # series of dependent jobs to run one after another.
        if self.opts.get("resumejob"):
            real_resumedir = os.path.abspath(self.opts.get("resumedir"))
            # edit the absolute resumedir path into the resumedir= argument
            foundresumedir = False
            for argidx in range(len(self.argv)):
                splitarg = self.argv[argidx].split("=")
                if len(splitarg) > 1:
                    if splitarg[0] == "resumedir":
                        foundresumedir = True
                        self.argv[argidx] = "resumedir=" + real_resumedir
            # if I've made it all the way through the list without finding
            # a resumedir argument, add it now
            if not foundresumedir:
                self.argv.append("resumedir=" + real_resumedir)

        if self.opts.get("createjob") and not self.subjob:
            job_dir = self.create_job(self.opts.get("jobdir"), extra_files,
                                      extra_dirs, extra_opt_files, extra_opt_dirs)
            if not opts.job_manager.run:
                print("created", end=' ')
                if opts.get("submit"):
                    print("and submitted", end=' ')
                print("job in directory", job_dir)
            sys.exit(0)

    def _args_to_string(self, args, strip=[None]):
        retval = ""
        for arg in args:
            argout = arg
            splitarg = arg.split("=")
            if len(splitarg) > 1:
                if splitarg[0] in strip:
                    argout = None
                elif(splitarg[1].count(" ")) > 0:
                    argout = splitarg[0] + '="'
                    splitarg.pop(0)
                    argout = argout + "=".join(splitarg) + '"'
            if argout:
                if retval != "":
                    retval += " "
                retval += argout
        return retval

    def create_script(self, template_path, name, directory, subs):
        output_path = os.path.join(directory, name)
        output = open(output_path, 'w')
        process_template(template_path, output, subs)
        output.close()
        os.chmod(output_path, 0o755)

    def create_composite_script(self,
                                header_template_path,
                                subjob_template_path,
                                footer_template_path,
                                name, directory, subs):
        output_path = os.path.join(directory, name)
        output = open(output_path, 'w')
        process_template(header_template_path, output, subs)
        subsubs = subs.copy()
        for submanager in self.submanagers:
            subsubs["subjobindex"] = submanager.subjob_index
            subsubs["subjobdir"] = os.path.join(subsubs["jobdir"],
                                                 submanager.subdirectory)
            subsubs["args"] = submanager._args_to_string(submanager.argv[1:],
                                                          job_mgr_opts.options())
            process_template(subjob_template_path, output, subsubs)
        process_template(footer_template_path, output, subs)
        output.close()
        os.chmod(output_path, 0o755)

    def create_job_script(self, job_name, directory, subs):
        if self.opts.job_manager.templatepath:
            job_template_path = self.opts.job_manager.templatepath
            if self.opts.job_manager.jmverbosity > 0:
                print("using job template:", job_template_path)
        else:
            job_template_path = search_job_manager_paths(self.opts.job_manager.template)
            if job_template_path:
                if self.opts.job_manager.jmverbosity > 0:
                    print("using job template:", job_template_path)
            else:
                job_template_path = os.path.join(get_script_templates_dir(),
                                                 "job_example")
                if self.opts.job_manager.jmverbosity > 0:
                    print("warning: using example job template:")
                    print("  ", job_template_path)
                    print("because no job template was found in search path")

        self.create_script(job_template_path, job_name, directory, subs)

    def create_multijob_script(self, job_name, directory, subs):
        if self.opts.job_manager.multitemplatepath:
            multijob_header_template_path = self.opts.job_manager.multitemplatepath + '_header'
            multijob_subjob_template_path = self.opts.job_manager.multitemplatepath + '_subjob'
            multijob_footer_template_path = self.opts.job_manager.multitemplatepath + '_footer'
            if self.opts.job_manager.jmverbosity > 0:
                print("using multijob templates:")
                print("    ", multijob_header_template_path)
                print("    ", multijob_subjob_template_path)
                print("    ", multijob_footer_template_path)
        else:
            multijob_header_template_path = search_job_manager_paths(self.opts.job_manager.multitemplate
                                                                     + '_header')
            if multijob_header_template_path:
                basedir = os.path.dirname(multijob_header_template_path)
                multijob_subjob_template_path = os.path.join(basedir,
                                                             self.opts.job_manager.multitemplate + \
                                                             '_subjob')
                multijob_footer_template_path = os.path.join(basedir,
                                                             self.opts.job_manager.multitemplate + \
                                                             '_footer')
                if self.opts.job_manager.jmverbosity > 0:
                    print("using multijob templates:")
                    print("    ", multijob_header_template_path)
                    print("    ", multijob_subjob_template_path)
                    print("    ", multijob_footer_template_path)
            else:
                multijob_header_template_path = os.path.join(get_script_templates_dir(),
                                                 "multijob_example_header")
                multijob_subjob_template_path = os.path.join(get_script_templates_dir(),
                                                 "multijob_example_subjob")
                multijob_footer_template_path = os.path.join(get_script_templates_dir(),
                                                 "multijob_example_footer")
                if self.opts.job_manager.jmverbosity > 0:
                    print("warning: using multijob templates:")
                    print("  ", multijob_header_template_path)
                    print("  ", multijob_subjob_template_path)
                    print("  ", multijob_footer_template_path)
                    print("because no multijob templates were found in search path")

        self.create_composite_script(multijob_header_template_path,
                                     multijob_subjob_template_path,
                                     multijob_footer_template_path,
                                     job_name, directory, subs)

    def create_resumejob_script(self, directory, subs):
        resumejob_name = os.path.basename(directory) + "_resumejob"
        if self.opts.job_manager.resumetemplatepath:
            resumejob_template_path = self.opts.job_manager.resumetemplatepath
            if self.opts.job_manager.jmverbosity > 0:
                print("using resume job template:", resumejob_template_path)
        else:
            resumejob_template_path = search_job_manager_paths(self.opts.job_manager.resumetemplate)
            if resumejob_template_path:
                if self.opts.job_manager.jmverbosity > 0:
                    print("using resume job template:", resumejob_template_path)
            else:
                resumejob_template_path = os.path.join(get_script_templates_dir(),
                                                 "resumejob_example")
                if self.opts.job_manager.jmverbosity > 0:
                    print("warning: using example resume job template:")
                    print("  ", resumejob_template_path)
                    print("because no resume job template was found in search path")

        self.create_script(resumejob_template_path, resumejob_name, directory, subs)

    def create_resumemultijob_script(self, directory, subs):
        resumejob_name = os.path.basename(directory) + "_resumejob"
        if self.opts.job_manager.resumemultitemplatepath:
            resumemultijob_header_template_path = self.opts.job_manager.resumemultitemplatepath + '_header'
            resumemultijob_subjob_template_path = self.opts.job_manager.resumemultitemplatepath + '_subjob'
            resumemultijob_footer_template_path = self.opts.job_manager.resumemultitemplatepath + '_footer'
            if self.opts.job_manager.jmverbosity > 0:
                print("using resumemultijob templates:")
                print("    ", resumemultijob_header_template_path)
                print("    ", resumemultijob_subjob_template_path)
                print("    ", resumemultijob_footer_template_path)
        else:
            resumemultijob_header_template_path = search_job_manager_paths(self.opts.job_manager.resumemultitemplate
                                                                     + '_header')
            if resumemultijob_header_template_path:
                basedir = os.path.dirname(resumemultijob_header_template_path)
                resumemultijob_subjob_template_path = os.path.join(basedir,
                                                             self.opts.job_manager.resumemultitemplate + \
                                                             '_subjob')
                resumemultijob_footer_template_path = os.path.join(basedir,
                                                             self.opts.job_manager.resumemultitemplate + \
                                                             '_footer')
                if self.opts.job_manager.jmverbosity > 0:
                    print("using resumemultijob templates:")
                    print("    ", resumemultijob_header_template_path)
                    print("    ", resumemultijob_subjob_template_path)
                    print("    ", resumemultijob_footer_template_path)
            else:
                resumemultijob_header_template_path = os.path.join(get_script_templates_dir(),
                                                 "resumemultijob_example_header")
                resumemultijob_subjob_template_path = os.path.join(get_script_templates_dir(),
                                                 "resumemultijob_example_subjob")
                resumemultijob_footer_template_path = os.path.join(get_script_templates_dir(),
                                                 "resumemultijob_example_footer")
                if self.opts.job_manager.jmverbosity > 0:
                    print("warning: using resumemultijob templates:")
                    print("  ", resumemultijob_header_template_path)
                    print("  ", resumemultijob_subjob_template_path)
                    print("  ", resumemultijob_footer_template_path)
                    print("because no resumemultijob templates were found in search path")

        self.create_composite_script(resumemultijob_header_template_path,
                                     resumemultijob_subjob_template_path,
                                     resumemultijob_footer_template_path,
                                     resumejob_name, directory, subs)

    def submit_job(self, job_name, subs):
        submitter = "qsub"
        for line in open(job_name, 'r'):
            m = re.match(".*synergia_workflow_submitter:(.*)", line)
            if m:
                submitter = m.group(1)
        submitter = process_line(submitter, subs)
        os.system(submitter + ' ./' + job_name)

    def create_job(self, directory, extra_files, extra_dirs,
                   extra_opt_files, extra_opt_dirs):
###        real_script = os.path.abspath(self.argv[0])
        old_cwd = os.getcwd()
        overwrite = self.opts.get("overwrite")
        if self.subjob:
            directory = os.path.join(directory, self.subdirectory)
            os.mkdir(directory)
        else:
            directory = create_new_directory(directory, 0, overwrite)
        self.directory = directory
        os.chdir(directory)
        shutil.copy(self.real_script, ".")
        if self.real_options_file:
            shutil.copy(self.real_options_file, ".")
        commandfile = open("command", "w")
        commandfile.write("%s\n" % self._args_to_string(self.argv))
        commandfile.close()
        alloptionsfile = open("alloptions", "w")
        keys = self.opts.options()
        keys.sort()
        for key in keys:
            alloptionsfile.write(key + "=" + str(self.opts.get(key)) + "\n")
        alloptionsfile.close()
        os.chdir(old_cwd)
        subs = {}
        for sub in job_mgr_opts.options():
            subs[sub] = job_mgr_opts.get(sub)
        subs["numproc"] = self.opts.get("numproc")
        numnode = int((self.opts.get("numproc") + self.opts.get("procspernode") - 1) / \
                  self.opts.get("procspernode"))
        subs["procspernode"] = self.opts.get("procspernode")
        subs["numnode"] = numnode
        subs["synergia2dir"] = self.synergia_dir
        subs["args"] = self._args_to_string(self.argv[1:], job_mgr_opts.options())
        subs["jobdir"] = os.path.abspath(self.directory)
        subs["subjobindex"] = self.subjob_index
        subs["subjobdir"] = os.path.join(subs["jobdir"], self.subdirectory)
        numsubjobs = len(self.submanagers)
        subs["numsubjobs"] = numsubjobs
        if numsubjobs > 0:
            val = self.opts.get("numproc") / numsubjobs
            if val > 0:
                subs["subnumproc"] = val
            else:
                subs["subnumproc"] = 1
            val = (subs["subnumproc"] + self.opts.get("procspernode") - 1) / \
                  self.opts.get("procspernode")
            if val > 0:
                subs["subnumnode"] = val
            else:
                subs["subnumnode"] = 1
        else:
            subs["subnumproc"] = self.opts.get("numproc")
            subs["subnumnode"] = subs["numnode"]
        if self.standalone:
            subs["synergia_executable"] = self.real_script
            subs["script"] = ""
        else:
            subs["script"] = os.path.basename(self.real_script)
        job_name = os.path.basename(directory) + "_job"
        if not self.subjob:
            if self.submanagers == []:
                self.create_job_script(job_name, directory, subs)
                self.create_resumejob_script(directory, subs)
            else:
                self.create_multijob_script(job_name, directory, subs)
                self.create_resumemultijob_script(directory, subs)

        self.create_script(search_job_manager_paths("cleanup"), "cleanup",
                            directory, subs)
        if extra_files:
            self.copy_extra_files(extra_files)
        if extra_dirs:
            self.copy_extra_dirs(extra_dirs)
        if extra_opt_files:
            self.copy_extra_files(extra_opt_files, optional=True)
        if extra_opt_dirs:
            self.copy_extra_dirs(extra_opt_dirs, optional=True)
        for submanager in self.submanagers:
            submanager.create_job(directory,extra_files, extra_dirs,
                   extra_opt_files, extra_opt_dirs)
        if self.opts.get("submit") and not self.subjob:
            os.chdir(directory)
            self.submit_job(job_name, subs)
            os.chdir(old_cwd)
        if self.opts.get("run"):
            os.chdir(directory)
            os.system("./" + job_name)
            os.chdir(old_cwd)
        return directory

    def copy_extra_files(self, files, optional=False):
        for item in files:
            if type(item) == tuple or type(item) == list:
                file = getattr(item[0], item[1])
            else:
                file = item
            if os.path.exists(file):
                shutil.copy(file, self.directory)
            else:
                if not optional:
                    raise RuntimeError("Job_manager: required file " + file + " not found")

    def copy_extra_dirs(self, dirs, optional=False):
        for item in dirs:
            if type(item) == tuple or type(item) == list:
                dir = getattr(item[0], item[1])
            else:
                dir = item
            if os.path.exists(dir):
                shutil.copytree(dir, self.directory + os.sep + dir)
            else:
                if not optional:
                    raise RuntimeError("Job_manager: required directory " + dir + " not found")

    def create_cxx_options_source(self):
        filename_base = self.real_script + "_options"
        classname = os.path.basename(self.real_script) + "_options"
        classname = classname[0].upper() + classname[1:]
        includeguard = classname.upper() + "_H_"
        disclaimer = '// this file was automatically generated by the command\n' + \
            '//     synergia ' + sys.argv[0] + ' ' + sys.argv[-1] + '\n' + \
            '// DO NOT EDIT\n'

        header_filename = filename_base + ".h"
        header = open(header_filename, "w")
        header.write('#ifndef ' + includeguard + '\n')
        header.write('#define ' + includeguard + '\n')
        header.write('#include <string>\n')
        header.write('\n')
        header.write(disclaimer)
        header.write('\n')
        header.write('struct ' + classname + '\n')
        header.write('{\n')
        header.write('    ' + classname + '(int argc, char **argv);\n')
        for optname in self.opts.dict:
            opt = self.opts.dict[optname]
            header.write('    ' + cxx_typename(opt.val_type) + ' ' + optname + ';\n')
        header.write('};\n')
        header.write('#endif /* ' + includeguard + ' */\n')
        header.close()

        source = open(filename_base + ".cc", "w")
        source.write('#include "' + os.path.basename(header_filename) + '"\n')
        source.write('#include "synergia/utils/command_line_arg.h"\n')
        source.write('\n')
        source.write(disclaimer)
        source.write('\n')
        source.write(classname + '::' + classname + '(int argc, char **argv) :\n')
        count = 0
        for optname in self.opts.dict:
            if (count > 0):
                source.write(',\n')
            source.write('    ' + optname + '(' + \
                         cxx_source_value(self.opts.get(optname)) + ')')
            count += 1
        source.write('\n')
        source.write('{\n')
        source.write('    for (int i = 1; i < argc; ++i) {\n')
        source.write('        Command_line_arg arg(argv[i]);\n')
        source.write('        if (arg.is_equal_pair()) {\n')
        count = 0
        for optname in self.opts.dict:
            opt = self.opts.dict[optname]
            if (count > 0):
                source.write(' else ')
            else:
                source.write('            ')
            source.write('if (arg.get_lhs() == "' + optname + '") {\n')
            source.write('                ' + optname + ' = arg.extract_value<' + \
                         cxx_typename(opt.val_type) + ' >();\n')
            source.write('            }')
            count += 1
# jfa: next section is redundant now that job_mgr opts are being stripped.
#      fixme.
        source.write(''' else if (arg.get_lhs() == "synergia_executable") {
                // ignore
            } else if (arg.get_lhs() == "run") {
                // ignore
            } else if (arg.get_lhs() == "submit") {
                // ignore
            } else if (arg.get_lhs() == "overwrite") {
                // ignore
            } else if (arg.get_lhs() == "numproc") {
                // ignore
            } else if (arg.get_lhs() == "procspernode") {
                // ignore
            } else if (arg.get_lhs() == "jobdir") {
                // ignore
            } else {
                throw std::runtime_error("Unknown argument " + arg.get_lhs());
            }
        } else {
            throw std::runtime_error("Bad argument " + std::string(argv[i]));
        }
    }
}
''')
        source.close()

def cxx_typename(type):
    if type == int:
        retval = "int"
    elif type == float:
        retval = "double"
    elif type == bool:
        retval = "bool"
    elif type == str:
        retval = "std::string"
    else:
        raise RuntimeError("Unable to map type " + str(type) + " to C++ type")
    return retval

def cxx_source_value(val):
    retval = str(val)
    if type(val) == bool:
        if val:
            retval = "true"
        else:
            retval = "false"
    if type(val) == str:
        retval = '"' + val + '"'
    return retval

def process_line(line, subs, unknown_vars = []):
    match = re.search("@@[A-z0-9]+@@", line)
    while match:
        var = match.group().replace("@@", "")
        if var in subs:
            replacement = str(subs[var])
        else:
            replacement = ""
            unknown_vars.append(var)
        original = match.group()
        line = line.replace(original, replacement)
        match = re.search("@@[A-z0-9]+@@", line)
    match = re.search("__([A-z0-9]+){{(.*)}}{{(.*)}}__", line)
    while match:
        var = match.group(1)
        have_var = False
        if var in subs:
            if (subs[var] != None) and (subs[var] != False):
                have_var = True
        if have_var:
            replacement = match.group(2)
        else:
            replacement = match.group(3)
            if unknown_vars.count(var) > 0:
                unknown_vars.remove(var)
        original = "__%s{{%s}}{{%s}}__" % (match.group(1), match.group(2),
                                           match.group(3))
        line = line.replace(original, replacement)
        match = re.search("@@[A-z0-9]+@@", line)
    return line

def process_template(template_name, output, subs):
    template = open(template_name, "r")
    unknown_vars = []
    for line in template.readlines():
        output.write(process_line(line, subs, unknown_vars))
    for var in unknown_vars:
        print("process_template warning: variable \"%s\" unknown." % var)
    template.close()

def create_new_directory(directory, version, overwrite):
    if version > 9999:
        print("Sanity check failure: attempt to create directory version %d."\
              % version)
        print("Maximum is 9999.")
        sys.exit(1)
    if overwrite:
        created_directory = directory
    else:
        created_directory = "%s.%02d" % (directory, version)
    if os.path.isdir(created_directory):
        if overwrite:
            shutil.rmtree(created_directory)
            os.mkdir(created_directory)
        else:
            created_directory = create_new_directory(directory,
                                                     version + 1,
                                                     overwrite)
    else:
        os.mkdir(created_directory)
    return created_directory


class Solution:
    def numDistinctIslands(self, grid: List[List[int]]) -> int:
        # transform the shape into a str
        # str of directions
        ans = 0
        r_offset = [-1,0,1,0]
        c_offset = [0,1,0,-1]
        s_offset = ['U','R','D','L']
        R = len(grid)
        C = len(grid[0])
        def DFS(grid, r, c, strval):
            grid[r][c] = 0
            for i in range(4):
                new_r = r+r_offset[i]
                new_c = c+c_offset[i]
                if new_r>=0 and new_r<R and new_c>=0 and new_c<C and grid[new_r][new_c]==1:
                    strval += s_offset[i]
                    # print(strval)
                    strval = DFS(grid, new_r, new_c, strval)
            strval += 'B'   # need this BACKTRACKING symbol
            return strval

        str_l = []
        for i in range(R):
            for j in range(C):
                if grid[i][j]==1:
                    str_island = DFS(grid, i, j, '')
                    str_l.append(str_island)
        ans = set(str_l)
        return len(ans)

Number of Distinct Islands/Number of Distinct Islands.py


import argparse
from collections import OrderedDict

import torch

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('input', default='~/userfolder/mmsegmentation-lesion/work_dirs/xxx/iter_40000.pth')
    parser.add_argument('output', default='~/userfolder/HRNet-Semantic-Segmentation/temp.pth')
    args = parser.parse_args()

    key_dict = {
        'decode_head.convs.0.conv.': 'model.last_layer.0.',
        'decode_head.convs.0.bn.': 'model.last_layer.1.',
        'decode_head.conv_seg.': 'model.last_layer.3.',
        'backbone.': 'model.'
    }

    new_pth = OrderedDict()
    pth = torch.load(args.input, map_location=torch.device('cpu'))['state_dict']
    for key, weight in pth.items():
        new_key = key
        for k, v in key_dict.items():
            new_key = new_key.replace(k, v)
        new_pth[new_key] = weight

        print(key, '->', new_key)

    torch.save(new_pth, args.output)


from pathlib import Path

import jsonpickle

class Cache:

    def __init__(self):
        self.ids = {}

    @staticmethod
    def from_file(cache_path: Path):
        try:
            with open(str(cache_path), 'r') as cache_file:
                cache = jsonpickle.decode(cache_file.read())
                cache.ids = {int(key): cache.ids[key] for key in cache.ids}
                return cache
        except:
            return Cache()

    def add(self, torrent_id: str, reason: str, cache_path: Path):
        self.ids[torrent_id] = reason
        self.write(cache_path)

    def write(self, cache_path: Path):
        with open(str(cache_path), 'w') as cache_file:
            encoded = jsonpickle.encode(self)
            cache_file.write(encoded)


from django_dominate.django_tags import *

from dominate.tags import *
from dominate.util import text

from goal.templates.dominate_tags import *

from suggestion.templates.dominate_tags import *

suggestion_url = \
    "{% url 'suggestion' request.goal.slug suggestion.slug %}"

def result():
    with django_block("content") as content:
        goal_header()

        with django_with("suggestion as the_suggestion"):
            with django_with("revision as the_revision"):
                with div(_class="row small-gap-below"):
                    column(4)
                    with column(4):
                        suggestion_image()

                with div(_class="row small-gap-below"):
                    column(2)
                    with column(8):
                        with p():
                            text("This is a previous version of a ")
                            a("suggestion", href=suggestion_url)
                            text(" by {{ the_suggestion.owner.name }}")
                        text("{{ the_revision.description|markdown }}")

    return (
        "{% extends 'base.html' %}\n",
        "{% load humanize %}",
        "{% load markdown_deux_tags %}",
        "{% load notification_tags %}",
        content,
    )

# done123


import decimal
import json
import os
import urllib.request

import boto3

DATA_URL = 'https://supercharge.info/service/supercharge/allSites'

def charger_eq(existing, update):
    return all((
        key in existing and existing[key] == update[key]
        for key in update.keys()))

def main(event, context):
    try:
        local_filename, headers = urllib.request.urlretrieve(DATA_URL)
    except Exception as e:
        # some other error happened
        error = str(e)
        return "We had an error: %s" % (error)
    else:
        dynamodb = boto3.resource('dynamodb')
        table = dynamodb.Table(os.environ['CHARGERS_TABLE'])

        num_chargers = num_chargers_new = num_chargers_updated = 0
        with open(local_filename) as json_file:
            chargers = json.load(json_file, parse_float=decimal.Decimal)
            for charger in chargers:
                num_chargers += 1
                charger['latitude'] = charger['gps']['latitude']
                charger['longitude'] = charger['gps']['longitude']
                del charger['gps']

                try:
                    existing_charger = table.get_item(
                        Key={'id': int(charger['id'])})['Item']
                except Exception as e:
                    num_chargers_new += 1
                    table.put_item(Item=charger)
                else:
                    if charger_eq(existing_charger, charger):
                        # no need to update
                        continue

                    expressionAttributeNames = dict([
                        ('#%s' % key, key)
                        for key in charger.keys() if key != 'id'])

                    updateExpression = 'set %s' % (', '.join([
                        '#%s = :%s' % (key, key)
                        for key in charger.keys() if key != 'id']))

                    expressionAttributeValues = dict([(
                        ':%s' % key, val)
                        for key, val in charger.items()
                        if key != 'id'])

                    table.update_item(
                        Key={'id': int(charger['id'])},
                        UpdateExpression=updateExpression,
                        ExpressionAttributeNames=expressionAttributeNames,
                        ExpressionAttributeValues=expressionAttributeValues,
                        ReturnValues='NONE',
                    )
                    num_chargers_updated += 1

    return ("We updated %s chargers and created %s chargers for a total "
            "of %s!") % (num_chargers_updated, num_chargers_new, num_chargers)


# -*- coding: utf-8 -*-
"""
Spyder Editor
Problem:Longest Common Subset
Time Complexity:n**2

This is a temporary script file.
"""
from prettytable import PrettyTable

def LCSubset(x,y,m,n):
    for i in range(m+1):
        for j in range(n+1):
            if i==0 or j==0:
                l[i][j]=0
            elif x[i-1]==y[j-1]:
                l[i][j]=1+l[i-1][j-1]
            else:
                l[i][j]=0
    max_len=0

    for i in range(m+1):
        for j in range(n+1):
            if max_len<l[i][j]:
                max_len=l[i][j]

    return max_len

def PrintSubset(ans):
    for i in range(m+1):
        for j in range(n+1):
            if l[i][j]==ans:
                req_i=i
                req_j=j

    subset=[]
    while ans!=0:
        if x[req_i-1]==y[req_j-1]:
            subset.insert(0,x[req_i-1])
            req_i-=1
            req_j-=1
            ans-=1

    return subset

x=input("Enter First String:")
y=input("Enter Second Sting:")
m=len(x)
n=len(y)
l=[[-1]*(n+1) for i in range(m+1)]
ans=LCSubset(x,y,m,n)
print(ans)

Table=PrettyTable()
Table.field_names=[i for i in range(n+1)]
for i in l:
    Table.add_row(i)

print(Table)

subset=PrintSubset(ans)
print(subset)


#----------------------------------------------------
#   CH_12_8_EX_1_Pyplot_Semilogx_Semilogy_Loglog.py
#----------------------------------------------------
import numpy as np
import matplotlib.pyplot as plt
#
x = np.linspace(0., 10., 11)
y1 = np.exp(-0.5*x)
#----- 繪製實數座標曲線 (X-Y)
plt.subplot(2,2,1)
plt.plot(x, y1, 'r-', lw = 3, label = 'x-y')
plt.ylabel('y')
plt.grid()
plt.legend(loc='best')
#----- 繪製半對數座標曲線 (logX-Y)
plt.subplot(2,2,2)
plt.semilogx(x, y1, 'g-', lw = 3, label = 'logx-y')
plt.xlabel('x')
plt.grid()
plt.legend(loc='best')
#----- 繪製半對數座標曲線 (X-logY)
plt.subplot(2,2,3)
plt.semilogy(x, y1, 'b-', lw = 3, label = 'x-logy')
plt.xlabel('x')
plt.ylabel('logy')
plt.grid()
plt.legend(loc='best')
#----- 繪製雙對數座標曲線 (logX-logY)
plt.subplot(2,2,4)
plt.loglog(x, y1, 'k-', lw = 3, label = 'logx-logy')
plt.grid()
plt.xlabel('logx')
plt.legend(loc='best')
#
plt.savefig('CH_12_8_EX_1_Pyplot_Semilogx_Semilogy_Loglog.png')
plt.show()


"""
Simple function for building ensembles of iterators that
represent disjoint partitions of an overall Cartesian product.
"""
from __future__ import annotations
from typing import Tuple, Optional, Collection, Sequence, Iterable
import doctest
import collections.abc as collections_abc # Avoid collision with function argument.
import itertools
from parts import parts

def products(
        *collections: Tuple[Collection, ...],
        number: Optional[int] = None
    ) -> Sequence[Iterable]:
    """
    Accept zero or more :obj:`~collections.abc.Collection` instances as arguments
    and return a :obj:`~collections.abc.Sequence` of the specified number of disjoint
    subsets of the `Cartesian product <https://en.wikipedia.org/wiki/Cartesian_product>`__
    of the supplied :obj:`~collections.abc.Collection` instances. Each subset is
    represented as an :obj:`~collections.abc.Iterable` and the union of the disjoint
    subsets is equal to the overall Cartesian product.

    :param collections: Zero or more arguments that represent the
        factor sets of the Cartesian product.
    :param number: Number of disjoint subsets to return.

    >>> ss = products(range(1, 3), {'a', 'b'}, (False, True), number=3)
    >>> for s in sorted([sorted(list(s)) for s in ss]):
    ...     for t in s:
    ...         print(t)
    (1, 'a', False)
    (1, 'a', True)
    (1, 'b', False)
    (1, 'b', True)
    (2, 'a', False)
    (2, 'a', True)
    (2, 'b', False)
    (2, 'b', True)

    Two additional basic examples are presented below.

    >>> (x, y, z) = ([1, 2], ['a', 'b'], [True, False])
    >>> [list(s) for s in products(x, y, number=2)]
    [[(1, 'a'), (1, 'b')], [(2, 'a'), (2, 'b')]]
    >>> for s in [list(s) for s in products(x, y, z, number=2)]:
    ...     print(s)
    [(1, 'a', True), (1, 'a', False), (1, 'b', True), (1, 'b', False)]
    [(2, 'a', True), (2, 'a', False), (2, 'b', True), (2, 'b', False)]

    By default (if the ``number`` argument is not assigned a value), the number
    of disjoint subsets is one. Note that the union of the disjoint subsets is
    equivalent to the output of the :obj:`itertools.product` function.

    >>> p = itertools.product([1, 2], {'a', 'b'}, (True, False))
    >>> ss = products([1, 2], {'a', 'b'}, (True, False))
    >>> list(p) == list(list(ss)[0])
    True

    If no sets are specified, the Cartesian product consists of a single empty
    tuple. If there is one set, the Cartesian product consists of a set of
    one-element tuples. In both cases, a list of disjoint subsets is returned
    as in all other cases (even though the number of disjoint subsets may be
    one).

    >>> list(list(products())[0])
    [()]
    >>> list(list(products([1, 2]))[0])
    [(1,), (2,)]

    It is possible to confirm that the returned subsets are disjoint, and that
    the union of the disjoint subsets is the Cartesian product.

    >>> (x, y, z) = ([1, 2], ['a', 'b'], [True, False])
    >>> ss = [set(s) for s in products(x, y, z, x, y, z, number=5)]
    >>> set([len(ss[i] & ss[j]) for i in range(5) for j in range(5) if i != j])
    {0}
    >>> s = ss[0] | ss[1] | ss[2] | ss[3] | ss[4]
    >>> s == set(itertools.product(x, y, z, x, y, z))
    True
    >>> len(products(*[[1, 2, 3]]*1000, number=5))
    5
    >>> ls = [len(products(*[[1, 2]]*1000, number=n)) for n in range(1, 100)]
    >>> ls == list(range(1, 100))
    True

    If the requested number of disjoint subsets exceeds the number of elements in
    Cartesian product, the number of disjoint subsets will be equivalent to the
    number of elements in the Cartesian product.

    >>> len(products([1, 2], ['a', 'b'], number=10))
    4

    Any attempt to apply this function to arguments that have unsupported types
    raises an exception.

    >>> products([1, 2], number='abc')
    Traceback (most recent call last):
      ...
    TypeError: number of disjoint subsets must be an integer
    >>> products((i for i in range(3)), number=2)
    Traceback (most recent call last):
      ...
    TypeError: arguments must be collections
    >>> products([1, 2], number=0)
    Traceback (most recent call last):
      ...
    ValueError: number of disjoint subsets must be a positive integer
    >>> products([1, 2], number=0)
    Traceback (most recent call last):
      ...
    ValueError: number of disjoint subsets must be a positive integer
    """
    factors = list(collections) # Ensure factor sets are reusable and their quantity is known.

    if not all(isinstance(factor, collections_abc.Collection) for factor in factors):
        raise TypeError('arguments must be collections')

    if number is not None:
        if not isinstance(number, int):
            raise TypeError('number of disjoint subsets must be an integer')

        if number < 1:
            raise ValueError('number of disjoint subsets must be a positive integer')

    # If no target number of disjoint subsets has been supplied, simply
    # return a single product. Note that this function is not equivalent to
    # :obj:`itertools.product` because this function always returns a sequence
    # of iterables (the subsets). In this case, there is only one element
    # in that sequence.
    if number is None or number == 1:
        return [itertools.product(*factors)]

    # Determine the product of which prefix of arguments to break up
    # based on the target number of disjoint subsets.
    number_ = 1
    index = len(factors)
    for (i, a) in enumerate(factors):
        number_ = number_ * len(a)
        if number_ >= number:
            index = min(len(factors), i + 1)
            break

    # Create an iterable for each prefix.
    prefixes = list(parts(list(itertools.product(*factors[0: index])), number))

    # For each prefix, create an iterable for that prefix by concatenating
    # elements of the prefix to elements of suffix product.
    def generator(prefix):
        for p in prefix:
            for s in itertools.product(*factors[index:]):
                yield p + s # Concatenation of tuples.

    return [generator(prefix) for prefix in prefixes]

if __name__ == '__main__':
    doctest.testmod() # pragma: no cover


from torch.utils.data import Dataset
from torchvision import transforms
from scipy import misc
import numpy as np
import imageio
import torch
import os

to_tensor = transforms.Compose([
    transforms.ToTensor()
])

def extract_bayer_channels(raw):

    # Reshape the input bayer image

    ch_B  = raw[1::2, 1::2]
    ch_Gb = raw[0::2, 1::2]
    ch_R  = raw[0::2, 0::2]
    ch_Gr = raw[1::2, 0::2]

    RAW_combined = np.dstack((ch_B, ch_Gb, ch_R, ch_Gr))
    RAW_norm = RAW_combined.astype(np.float32) / (4 * 255)

    return RAW_norm

class LoadData(Dataset):

    def __init__(self, dataset_dir, dataset_size, dslr_scale, test=False):

        if test:
            self.raw_dir = os.path.join(dataset_dir, 'test', 'huawei_raw')
            self.dslr_dir = os.path.join(dataset_dir, 'test', 'canon')
            self.dataset_size = dataset_size
        else:
            self.raw_dir = os.path.join(dataset_dir, 'train', 'huawei_raw')
            self.dslr_dir = os.path.join(dataset_dir, 'train', 'canon')

        self.dataset_size = dataset_size
        self.scale = dslr_scale
        self.test = test

    def __len__(self):
        return self.dataset_size

    def __getitem__(self, idx):

        raw_image = np.asarray(imageio.imread(os.path.join(self.raw_dir, str(idx) + '.png')))
        raw_image = extract_bayer_channels(raw_image)
        raw_image = torch.from_numpy(raw_image.transpose((2, 0, 1)))

        dslr_image = misc.imread(os.path.join(self.dslr_dir, str(idx) + ".jpg"))
        dslr_image = np.asarray(dslr_image)
        dslr_image = np.float32(misc.imresize(dslr_image, self.scale / 2.0)) / 255.0
        dslr_image = torch.from_numpy(dslr_image.transpose((2, 0, 1)))

        return raw_image, dslr_image

class LoadVisualData(Dataset):

    def __init__(self, data_dir, size, scale, level, full_resolution=False):

        self.raw_dir = os.path.join(data_dir, 'test', 'huawei_full_resolution')

        self.dataset_size = size
        self.scale = scale
        self.level = level
        self.full_resolution = full_resolution
        self.test_images = os.listdir(self.raw_dir)

        if level > 1 or full_resolution:
            self.image_height = 1440
            self.image_width = 1984
        elif level > 0:
            self.image_height = 1280
            self.image_width = 1280
        else:
            self.image_height = 960
            self.image_width = 960

    def __len__(self):
        return self.dataset_size

    def __getitem__(self, idx):

        raw_image = np.asarray(imageio.imread(os.path.join(self.raw_dir, self.test_images[idx])))
        raw_image = extract_bayer_channels(raw_image)

        if self.level > 1 or self.full_resolution:
            raw_image = raw_image[0:self.image_height, 0:self.image_width, :]
        elif self.level > 0:
            raw_image = raw_image[80:self.image_height + 80, 352:self.image_width + 352, :]
        else:
            raw_image = raw_image[240:self.image_height + 240, 512:self.image_width + 512, :]

        raw_image = torch.from_numpy(raw_image.transpose((2, 0, 1)))

        return raw_image


import warnings
import torch
import torch.nn as nn
from mmcv.cnn.bricks.transformer import (BaseTransformerLayer,
                                         TransformerLayerSequence,
                                         build_transformer_layer_sequence,
                                         build_attention,
                                         build_feedforward_network)
from mmcv.cnn.bricks.drop import build_dropout
from mmdet.models.utils.builder import TRANSFORMER
from mmcv.cnn import build_norm_layer, xavier_init
from mmcv.runner.base_module import BaseModule
from mmcv.cnn.bricks.registry import (ATTENTION,TRANSFORMER_LAYER,
                                      TRANSFORMER_LAYER_SEQUENCE)
from mmcv.utils import deprecated_api_warning, ConfigDict
import copy
from torch.nn import ModuleList
from .attention import FlashMHA
import torch.utils.checkpoint as cp

from mmcv.runner import auto_fp16

@ATTENTION.register_module()
class PETRMultiheadFlashAttention(BaseModule):
    """A wrapper for ``torch.nn.MultiheadAttention``.
    This module implements MultiheadAttention with identity connection,
    and positional encoding  is also passed as input.
    Args:
        embed_dims (int): The embedding dimension.
        num_heads (int): Parallel attention heads.
        attn_drop (float): A Dropout layer on attn_output_weights.
            Default: 0.0.
        proj_drop (float): A Dropout layer after `nn.MultiheadAttention`.
            Default: 0.0.
        dropout_layer (obj:`ConfigDict`): The dropout_layer used
            when adding the shortcut.
        init_cfg (obj:`mmcv.ConfigDict`): The Config for initialization.
            Default: None.
        batch_first (bool): When it is True,  Key, Query and Value are shape of
            (batch, n, embed_dim), otherwise (n, batch, embed_dim).
             Default to False.
    """

    def __init__(self,
                 embed_dims,
                 num_heads,
                 attn_drop=0.,
                 proj_drop=0.,
                 dropout_layer=dict(type='Dropout', drop_prob=0.),
                 init_cfg=None,
                 batch_first=True,
                 **kwargs):
        super(PETRMultiheadFlashAttention, self).__init__(init_cfg)
        if 'dropout' in kwargs:
            warnings.warn(
                'The arguments `dropout` in MultiheadAttention '
                'has been deprecated, now you can separately '
                'set `attn_drop`(float), proj_drop(float), '
                'and `dropout_layer`(dict) ', DeprecationWarning)
            attn_drop = kwargs['dropout']
            dropout_layer['drop_prob'] = kwargs.pop('dropout')

        self.embed_dims = embed_dims
        self.num_heads = num_heads
        self.batch_first = True

        self.attn = FlashMHA(embed_dims, num_heads, attn_drop, dtype=torch.float16, device='cuda',
                                          **kwargs)

        self.proj_drop = nn.Dropout(proj_drop)
        self.dropout_layer = build_dropout(
            dropout_layer) if dropout_layer else nn.Identity()

    @deprecated_api_warning({'residual': 'identity'},
                            cls_name='MultiheadAttention')
    def forward(self,
                query,
                key=None,
                value=None,
                identity=None,
                query_pos=None,
                key_pos=None,
                attn_mask=None,
                key_padding_mask=None,
                **kwargs):
        """Forward function for `MultiheadAttention`.
        **kwargs allow passing a more general data flow when combining
        with other operations in `transformerlayer`.
        Args:
            query (Tensor): The input query with shape [num_queries, bs,
                embed_dims] if self.batch_first is False, else
                [bs, num_queries embed_dims].
            key (Tensor): The key tensor with shape [num_keys, bs,
                embed_dims] if self.batch_first is False, else
                [bs, num_keys, embed_dims] .
                If None, the ``query`` will be used. Defaults to None.
            value (Tensor): The value tensor with same shape as `key`.
                Same in `nn.MultiheadAttention.forward`. Defaults to None.
                If None, the `key` will be used.
            identity (Tensor): This tensor, with the same shape as x,
                will be used for the identity link.
                If None, `x` will be used. Defaults to None.
            query_pos (Tensor): The positional encoding for query, with
                the same shape as `x`. If not None, it will
                be added to `x` before forward function. Defaults to None.
            key_pos (Tensor): The positional encoding for `key`, with the
                same shape as `key`. Defaults to None. If not None, it will
                be added to `key` before forward function. If None, and
                `query_pos` has the same shape as `key`, then `query_pos`
                will be used for `key_pos`. Defaults to None.
            attn_mask (Tensor): ByteTensor mask with shape [num_queries,
                num_keys]. Same in `nn.MultiheadAttention.forward`.
                Defaults to None.
            key_padding_mask (Tensor): ByteTensor with shape [bs, num_keys].
                Defaults to None.
        Returns:
            Tensor: forwarded results with shape
            [num_queries, bs, embed_dims]
            if self.batch_first is False, else
            [bs, num_queries embed_dims].
        """

        if key is None:
            key = query
        if value is None:
            value = key
        if identity is None:
            identity = query
        if key_pos is None:
            if query_pos is not None:
                # use query_pos if key_pos is not available
                if query_pos.shape == key.shape:
                    key_pos = query_pos
                else:
                    warnings.warn(f'position encoding of key is'
                                  f'missing in {self.__class__.__name__}.')
        if query_pos is not None:
            query = query + query_pos
        if key_pos is not None:
            key = key + key_pos

        # Because the dataflow('key', 'query', 'value') of
        # ``torch.nn.MultiheadAttention`` is (num_query, batch,
        # embed_dims), We should adjust the shape of dataflow from
        # batch_first (batch, num_query, embed_dims) to num_query_first
        # (num_query ,batch, embed_dims), and recover ``attn_output``
        # from num_query_first to batch_first.
        if self.batch_first:
            query = query.transpose(0, 1)
            key = key.transpose(0, 1)
            value = value.transpose(0, 1)

        out = self.attn(
            q=query,
            k=key,
            v=value,
            key_padding_mask=None)[0]

        if self.batch_first:
            out = out.transpose(0, 1)

        return identity + self.dropout_layer(self.proj_drop(out))

class MultiheadAttentionWrapper(nn.MultiheadAttention):
    def __init__(self, *args, **kwargs):
        super(MultiheadAttentionWrapper, self).__init__(*args, **kwargs)
        self.fp16_enabled = True

    @auto_fp16(out_fp32=True)
    def forward_fp16(self, *args, **kwargs):
        return super(MultiheadAttentionWrapper, self).forward(*args, **kwargs)

    def forward_fp32(self, *args, **kwargs):
        return super(MultiheadAttentionWrapper, self).forward(*args, **kwargs)

    def forward(self, *args, **kwargs):
        if self.training:
            return self.forward_fp16(*args, **kwargs)
        else:
            return self.forward_fp32( *args, **kwargs)

@ATTENTION.register_module()
class PETRMultiheadAttention(BaseModule):
    """A wrapper for ``torch.nn.MultiheadAttention``.
    This module implements MultiheadAttention with identity connection,
    and positional encoding  is also passed as input.
    Args:
        embed_dims (int): The embedding dimension.
        num_heads (int): Parallel attention heads.
        attn_drop (float): A Dropout layer on attn_output_weights.
            Default: 0.0.
        proj_drop (float): A Dropout layer after `nn.MultiheadAttention`.
            Default: 0.0.
        dropout_layer (obj:`ConfigDict`): The dropout_layer used
            when adding the shortcut.
        init_cfg (obj:`mmcv.ConfigDict`): The Config for initialization.
            Default: None.
        batch_first (bool): When it is True,  Key, Query and Value are shape of
            (batch, n, embed_dim), otherwise (n, batch, embed_dim).
             Default to False.
    """

    def __init__(self,
                 embed_dims,
                 num_heads,
                 attn_drop=0.,
                 proj_drop=0.,
                 dropout_layer=dict(type='Dropout', drop_prob=0.),
                 init_cfg=None,
                 batch_first=False,
                 fp16 = False,
                 **kwargs):
        super(PETRMultiheadAttention, self).__init__(init_cfg)
        if 'dropout' in kwargs:
            warnings.warn(
                'The arguments `dropout` in MultiheadAttention '
                'has been deprecated, now you can separately '
                'set `attn_drop`(float), proj_drop(float), '
                'and `dropout_layer`(dict) ', DeprecationWarning)
            attn_drop = kwargs['dropout']
            dropout_layer['drop_prob'] = kwargs.pop('dropout')

        self.embed_dims = embed_dims
        self.num_heads = num_heads
        self.batch_first = batch_first
        self.fp16_enabled = True
        if fp16:
            self.attn = MultiheadAttentionWrapper(embed_dims, num_heads, attn_drop,  **kwargs)
        else:
            self.attn = nn.MultiheadAttention(embed_dims, num_heads, attn_drop,  **kwargs)

        self.proj_drop = nn.Dropout(proj_drop)
        self.dropout_layer = build_dropout(
            dropout_layer) if dropout_layer else nn.Identity()

    @deprecated_api_warning({'residual': 'identity'},
                            cls_name='MultiheadAttention')
    def forward(self,
                query,
                key=None,
                value=None,
                identity=None,
                query_pos=None,
                key_pos=None,
                attn_mask=None,
                key_padding_mask=None,
                **kwargs):
        """Forward function for `MultiheadAttention`.
        **kwargs allow passing a more general data flow when combining
        with other operations in `transformerlayer`.
        Args:
            query (Tensor): The input query with shape [num_queries, bs,
                embed_dims] if self.batch_first is False, else
                [bs, num_queries embed_dims].
            key (Tensor): The key tensor with shape [num_keys, bs,
                embed_dims] if self.batch_first is False, else
                [bs, num_keys, embed_dims] .
                If None, the ``query`` will be used. Defaults to None.
            value (Tensor): The value tensor with same shape as `key`.
                Same in `nn.MultiheadAttention.forward`. Defaults to None.
                If None, the `key` will be used.
            identity (Tensor): This tensor, with the same shape as x,
                will be used for the identity link.
                If None, `x` will be used. Defaults to None.
            query_pos (Tensor): The positional encoding for query, with
                the same shape as `x`. If not None, it will
                be added to `x` before forward function. Defaults to None.
            key_pos (Tensor): The positional encoding for `key`, with the
                same shape as `key`. Defaults to None. If not None, it will
                be added to `key` before forward function. If None, and
                `query_pos` has the same shape as `key`, then `query_pos`
                will be used for `key_pos`. Defaults to None.
            attn_mask (Tensor): ByteTensor mask with shape [num_queries,
                num_keys]. Same in `nn.MultiheadAttention.forward`.
                Defaults to None.
            key_padding_mask (Tensor): ByteTensor with shape [bs, num_keys].
                Defaults to None.
        Returns:
            Tensor: forwarded results with shape
            [num_queries, bs, embed_dims]
            if self.batch_first is False, else
            [bs, num_queries embed_dims].
        """

        if key is None:
            key = query
        if value is None:
            value = key
        if identity is None:
            identity = query
        if key_pos is None:
            if query_pos is not None:
                # use query_pos if key_pos is not available
                if query_pos.shape == key.shape:
                    key_pos = query_pos
                else:
                    warnings.warn(f'position encoding of key is'
                                  f'missing in {self.__class__.__name__}.')
        if query_pos is not None:
            query = query + query_pos
        if key_pos is not None:
            key = key + key_pos

        # Because the dataflow('key', 'query', 'value') of
        # ``torch.nn.MultiheadAttention`` is (num_query, batch,
        # embed_dims), We should adjust the shape of dataflow from
        # batch_first (batch, num_query, embed_dims) to num_query_first
        # (num_query ,batch, embed_dims), and recover ``attn_output``
        # from num_query_first to batch_first.
        if self.batch_first:
            query = query.transpose(0, 1).contiguous()
            key = key.transpose(0, 1).contiguous()
            value = value.transpose(0, 1).contiguous()

        out = self.attn(
            query=query,
            key=key,
            value=value,
            attn_mask=attn_mask,
            key_padding_mask=key_padding_mask)[0]

        if self.batch_first:
            out = out.transpose(0, 1).contiguous()

        return identity + self.dropout_layer(self.proj_drop(out))

@TRANSFORMER_LAYER_SEQUENCE.register_module()
class PETRTransformerEncoder(TransformerLayerSequence):
    """TransformerEncoder of DETR.
    Args:
        post_norm_cfg (dict): Config of last normalization layer. Default：
            `LN`. Only used when `self.pre_norm` is `True`
    """

    def __init__(self, *args, post_norm_cfg=dict(type='LN'), **kwargs):
        super(PETRTransformerEncoder, self).__init__(*args, **kwargs)
        if post_norm_cfg is not None:
            self.post_norm = build_norm_layer(
                post_norm_cfg, self.embed_dims)[1] if self.pre_norm else None
        else:
            assert not self.pre_norm, f'Use prenorm in ' \
                                      f'{self.__class__.__name__},' \
                                      f'Please specify post_norm_cfg'
            self.post_norm = None

    def forward(self, *args, **kwargs):
        """Forward function for `TransformerCoder`.
        Returns:
            Tensor: forwarded results with shape [num_query, bs, embed_dims].
        """
        x = super(PETRTransformerEncoder, self).forward(*args, **kwargs)
        if self.post_norm is not None:
            x = self.post_norm(x)
        return x

@TRANSFORMER_LAYER_SEQUENCE.register_module()
class PETRTransformerDecoder(TransformerLayerSequence):
    """Implements the decoder in DETR transformer.
    Args:
        return_intermediate (bool): Whether to return intermediate outputs.
        post_norm_cfg (dict): Config of last normalization layer. Default：
            `LN`.
    """

    def __init__(self,
                 *args,
                 post_norm_cfg=dict(type='LN'),
                 return_intermediate=False,
                 **kwargs):

        super(PETRTransformerDecoder, self).__init__(*args, **kwargs)
        self.return_intermediate = return_intermediate
        if post_norm_cfg is not None:
            self.post_norm = build_norm_layer(post_norm_cfg,
                                              self.embed_dims)[1]
        else:
            self.post_norm = None

    def forward(self, query, *args, **kwargs):
        """Forward function for `TransformerDecoder`.
        Args:
            query (Tensor): Input query with shape
                `(num_query, bs, embed_dims)`.
        Returns:
            Tensor: Results with shape [1, num_query, bs, embed_dims] when
                return_intermediate is `False`, otherwise it has shape
                [num_layers, num_query, bs, embed_dims].
        """
        if not self.return_intermediate:
            x = super().forward(query, *args, **kwargs)
            if self.post_norm:
                x = self.post_norm(x)[None]
            return x

        intermediate = []
        for layer in self.layers:
            query = layer(query, *args, **kwargs)
            if self.return_intermediate:
                if self.post_norm is not None:
                    intermediate.append(self.post_norm(query))
                else:
                    intermediate.append(query)
        return torch.stack(intermediate)

@TRANSFORMER.register_module()
class PETRTemporalTransformer(BaseModule):
    """Implements the DETR transformer.
    Following the official DETR implementation, this module copy-paste
    from torch.nn.Transformer with modifications:
        * positional encodings are passed in MultiheadAttention
        * extra LN at the end of encoder is removed
        * decoder returns a stack of activations from all decoding layers
    See `paper: End-to-End Object Detection with Transformers
    <https://arxiv.org/pdf/2005.12872>`_ for details.
    Args:
        encoder (`mmcv.ConfigDict` | Dict): Config of
            TransformerEncoder. Defaults to None.
        decoder ((`mmcv.ConfigDict` | Dict)): Config of
            TransformerDecoder. Defaults to None
        init_cfg (obj:`mmcv.ConfigDict`): The Config for initialization.
            Defaults to None.
    """

    def __init__(self, encoder=None, decoder=None, init_cfg=None, cross=False):
        super(PETRTemporalTransformer, self).__init__(init_cfg=init_cfg)
        if encoder is not None:
            self.encoder = build_transformer_layer_sequence(encoder)
        else:
            self.encoder = None
        self.decoder = build_transformer_layer_sequence(decoder)
        self.embed_dims = self.decoder.embed_dims
        self.cross = cross

    def init_weights(self):
        # follow the official DETR to init parameters
        for m in self.modules():
            if hasattr(m, 'weight') and m.weight.dim() > 1:
                xavier_init(m, distribution='uniform')
        self._is_init = True

    def forward(self, memory, tgt, query_pos, pos_embed, attn_masks, temp_memory=None, temp_pos=None, mask=None, reg_branch=None):
        """Forward function for `Transformer`.
        Args:
            x (Tensor): Input query with shape [bs, c, h, w] where
                c = embed_dims.
            mask (Tensor): The key_padding_mask used for encoder and decoder,
                with shape [bs, h, w].
            query_embed (Tensor): The query embedding for decoder, with shape
                [num_query, c].
            pos_embed (Tensor): The positional encoding for encoder and
                decoder, with the same shape as `x`.
        Returns:
            tuple[Tensor]: results of decoder containing the following tensor.
                - out_dec: Output from decoder. If return_intermediate_dec \
                      is True output has shape [num_dec_layers, bs,
                      num_query, embed_dims], else has shape [1, bs, \
                      num_query, embed_dims].
                - memory: Output results from encoder, with shape \
                      [bs, embed_dims, h, w].
        """
        memory = memory.transpose(0, 1).contiguous()
        query_pos = query_pos.transpose(0, 1).contiguous()
        pos_embed = pos_embed.transpose(0, 1).contiguous()

        n, bs, c = memory.shape

        if tgt is None:
            tgt = torch.zeros_like(query_pos)
        else:
            tgt = tgt.transpose(0, 1).contiguous()

        if temp_memory is not None:
            temp_memory = temp_memory.transpose(0, 1).contiguous()
            temp_pos =  temp_pos.transpose(0, 1).contiguous()

        # out_dec: [num_layers, num_query, bs, dim]
        out_dec = self.decoder(
            query=tgt,
            key=memory,
            value=memory,
            key_pos=pos_embed,
            query_pos=query_pos,
            temp_memory=temp_memory,
            temp_pos=temp_pos,
            key_padding_mask=mask,
            attn_masks=[attn_masks, None],
            reg_branch=reg_branch,
            )
        out_dec = out_dec.transpose(1, 2).contiguous()
        memory = memory.reshape(-1, bs, c).transpose(0, 1).contiguous()
        return  out_dec, memory

@TRANSFORMER_LAYER.register_module()
class PETRTemporalDecoderLayer(BaseModule):
    """Base `TransformerLayer` for vision transformer.

    It can be built from `mmcv.ConfigDict` and support more flexible
    customization, for example, using any number of `FFN or LN ` and
    use different kinds of `attention` by specifying a list of `ConfigDict`
    named `attn_cfgs`. It is worth mentioning that it supports `prenorm`
    when you specifying `norm` as the first element of `operation_order`.
    More details about the `prenorm`: `On Layer Normalization in the
    Transformer Architecture <https://arxiv.org/abs/2002.04745>`_ .

    Args:
        attn_cfgs (list[`mmcv.ConfigDict`] | obj:`mmcv.ConfigDict` | None )):
            Configs for `self_attention` or `cross_attention` modules,
            The order of the configs in the list should be consistent with
            corresponding attentions in operation_order.
            If it is a dict, all of the attention modules in operation_order
            will be built with this config. Default: None.
        ffn_cfgs (list[`mmcv.ConfigDict`] | obj:`mmcv.ConfigDict` | None )):
            Configs for FFN, The order of the configs in the list should be
            consistent with corresponding ffn in operation_order.
            If it is a dict, all of the attention modules in operation_order
            will be built with this config.
        operation_order (tuple[str]): The execution order of operation
            in transformer. Such as ('self_attn', 'norm', 'ffn', 'norm').
            Support `prenorm` when you specifying first element as `norm`.
            Default：None.
        norm_cfg (dict): Config dict for normalization layer.
            Default: dict(type='LN').
        init_cfg (obj:`mmcv.ConfigDict`): The Config for initialization.
            Default: None.
        batch_first (bool): Key, Query and Value are shape
            of (batch, n, embed_dim)
            or (n, batch, embed_dim). Default to False.
    """

    def __init__(self,
                 attn_cfgs=None,
                 ffn_cfgs=dict(
                     type='FFN',
                     embed_dims=256,
                     feedforward_channels=1024,
                     num_fcs=2,
                     ffn_drop=0.,
                     act_cfg=dict(type='ReLU', inplace=True),
                 ),
                 operation_order=None,
                 norm_cfg=dict(type='LN'),
                 init_cfg=None,
                 batch_first=False,
                 with_cp=True,
                 **kwargs):

        deprecated_args = dict(
            feedforward_channels='feedforward_channels',
            ffn_dropout='ffn_drop',
            ffn_num_fcs='num_fcs')
        for ori_name, new_name in deprecated_args.items():
            if ori_name in kwargs:
                warnings.warn(
                    f'The arguments `{ori_name}` in BaseTransformerLayer '
                    f'has been deprecated, now you should set `{new_name}` '
                    f'and other FFN related arguments '
                    f'to a dict named `ffn_cfgs`. ', DeprecationWarning)
                ffn_cfgs[new_name] = kwargs[ori_name]

        super().__init__(init_cfg)

        self.batch_first = batch_first

        assert set(operation_order) & {
            'self_attn', 'norm', 'ffn', 'cross_attn'} == \
            set(operation_order), f'The operation_order of' \
            f' {self.__class__.__name__} should ' \
            f'contains all four operation type ' \
            f"{['self_attn', 'norm', 'ffn', 'cross_attn']}"

        num_attn = operation_order.count('self_attn') + operation_order.count(
            'cross_attn')
        if isinstance(attn_cfgs, dict):
            attn_cfgs = [copy.deepcopy(attn_cfgs) for _ in range(num_attn)]
        else:
            assert num_attn == len(attn_cfgs), f'The length ' \
                f'of attn_cfg {num_attn} is ' \
                f'not consistent with the number of attention' \
                f'in operation_order {operation_order}.'

        self.num_attn = num_attn
        self.operation_order = operation_order
        self.norm_cfg = norm_cfg
        self.pre_norm = operation_order[0] == 'norm'
        self.attentions = ModuleList()

        index = 0
        for operation_name in operation_order:
            if operation_name in ['self_attn', 'cross_attn']:
                if 'batch_first' in attn_cfgs[index]:
                    assert self.batch_first == attn_cfgs[index]['batch_first']
                else:
                    attn_cfgs[index]['batch_first'] = self.batch_first
                attention = build_attention(attn_cfgs[index])
                # Some custom attentions used as `self_attn`
                # or `cross_attn` can have different behavior.
                attention.operation_name = operation_name
                self.attentions.append(attention)
                index += 1

        self.embed_dims = self.attentions[0].embed_dims

        self.ffns = ModuleList()
        num_ffns = operation_order.count('ffn')
        if isinstance(ffn_cfgs, dict):
            ffn_cfgs = ConfigDict(ffn_cfgs)
        if isinstance(ffn_cfgs, dict):
            ffn_cfgs = [copy.deepcopy(ffn_cfgs) for _ in range(num_ffns)]
        assert len(ffn_cfgs) == num_ffns
        for ffn_index in range(num_ffns):
            if 'embed_dims' not in ffn_cfgs[ffn_index]:
                ffn_cfgs[ffn_index]['embed_dims'] = self.embed_dims
            else:
                assert ffn_cfgs[ffn_index]['embed_dims'] == self.embed_dims
            self.ffns.append(
                build_feedforward_network(ffn_cfgs[ffn_index],
                                          dict(type='FFN')))

        self.norms = ModuleList()
        num_norms = operation_order.count('norm')
        for _ in range(num_norms):
            self.norms.append(build_norm_layer(norm_cfg, self.embed_dims)[1])

        self.use_checkpoint = with_cp

    def _forward(self,
                query,
                key=None,
                value=None,
                query_pos=None,
                key_pos=None,
                temp_memory=None,
                temp_pos=None,
                attn_masks=None,
                query_key_padding_mask=None,
                key_padding_mask=None,
                **kwargs):
        """Forward function for `TransformerDecoderLayer`.

        **kwargs contains some specific arguments of attentions.

        Args:
            query (Tensor): The input query with shape
                [num_queries, bs, embed_dims] if
                self.batch_first is False, else
                [bs, num_queries embed_dims].
            key (Tensor): The key tensor with shape [num_keys, bs,
                embed_dims] if self.batch_first is False, else
                [bs, num_keys, embed_dims] .
            value (Tensor): The value tensor with same shape as `key`.
            query_pos (Tensor): The positional encoding for `query`.
                Default: None.
            key_pos (Tensor): The positional encoding for `key`.
                Default: None.
            attn_masks (List[Tensor] | None): 2D Tensor used in
                calculation of corresponding attention. The length of
                it should equal to the number of `attention` in
                `operation_order`. Default: None.
            query_key_padding_mask (Tensor): ByteTensor for `query`, with
                shape [bs, num_queries]. Only used in `self_attn` layer.
                Defaults to None.
            key_padding_mask (Tensor): ByteTensor for `query`, with
                shape [bs, num_keys]. Default: None.

        Returns:
            Tensor: forwarded results with shape [num_queries, bs, embed_dims].
        """

        norm_index = 0
        attn_index = 0
        ffn_index = 0
        identity = query
        if attn_masks is None:
            attn_masks = [None for _ in range(self.num_attn)]
        elif isinstance(attn_masks, torch.Tensor):
            attn_masks = [
                copy.deepcopy(attn_masks) for _ in range(self.num_attn)
            ]
            warnings.warn(f'Use same attn_mask in all attentions in '
                          f'{self.__class__.__name__} ')
        else:
            assert len(attn_masks) == self.num_attn, f'The length of ' \
                        f'attn_masks {len(attn_masks)} must be equal ' \
                        f'to the number of attention in ' \
                        f'operation_order {self.num_attn}'

        for layer in self.operation_order:
            if layer == 'self_attn':
                if temp_memory is not None:
                    temp_key = temp_value = torch.cat([query, temp_memory], dim=0)
                    temp_pos = torch.cat([query_pos, temp_pos], dim=0)
                else:
                    temp_key = temp_value = query
                    temp_pos = query_pos
                query = self.attentions[attn_index](
                    query,
                    temp_key,
                    temp_value,
                    identity if self.pre_norm else None,
                    query_pos=query_pos,
                    key_pos=temp_pos,
                    attn_mask=attn_masks[attn_index],
                    key_padding_mask=query_key_padding_mask,
                    **kwargs)
                attn_index += 1
                identity = query

            elif layer == 'norm':
                query = self.norms[norm_index](query)
                norm_index += 1

            elif layer == 'cross_attn':
                query = self.attentions[attn_index](
                    query,
                    key,
                    value,
                    identity if self.pre_norm else None,
                    query_pos=query_pos,
                    key_pos=key_pos,
                    attn_mask=attn_masks[attn_index],
                    key_padding_mask=key_padding_mask,
                    **kwargs)
                attn_index += 1
                identity = query

            elif layer == 'ffn':
                query = self.ffns[ffn_index](
                    query, identity if self.pre_norm else None)
                ffn_index += 1

        return query

    def forward(self,
                query,
                key=None,
                value=None,
                query_pos=None,
                key_pos=None,
                temp_memory=None,
                temp_pos=None,
                attn_masks=None,
                query_key_padding_mask=None,
                key_padding_mask=None,
                **kwargs
                ):
        """Forward function for `TransformerCoder`.
        Returns:
            Tensor: forwarded results with shape [num_query, bs, embed_dims].
        """

        if self.use_checkpoint and self.training:
            x = cp.checkpoint(
                self._forward,
                query,
                key,
                value,
                query_pos,
                key_pos,
                temp_memory,
                temp_pos,
                attn_masks,
                query_key_padding_mask,
                key_padding_mask,
                )
        else:
            x = self._forward(
            query,
            key,
            value,
            query_pos,
            key_pos,
            temp_memory,
            temp_pos,
            attn_masks,
            query_key_padding_mask,
            key_padding_mask,
        )
        return x


import json
import unittest
import sys
sys.path.append('/app')

from tipi_backend.app import create_app
from tipi_backend.settings import Config

class TestLimit(unittest.TestCase):

    def setUp(self):
        Config.TESTING = True
        Config.USE_ALERTS = True
        app = create_app(config=Config)
        self.client = app.test_client()

    def test_limit_each_call(self):
        x = 0
        data = json.dumps({'email': 'foo@bar.com', 'search': '{"topic": "bar"}'})
        while x < 10:
            # labeling
            res = self.client.post('/labels/extract', data={'text': 'example'})
            self.assertEqual(res.status_code, 200)
            # alerts
            res = self.client.post('/alerts', data=data, content_type='application/json')
            self.assertEqual(res.status_code, 200)
            x += 1
        res = self.client.post('/alerts', data=data, content_type='application/json')
        self.assertEqual(res.status_code, 429)

if __name__ == '__main__':
    unittest.main()


import pandas as pd
import sys
import os
import time
sys.path.append(os.path.abspath(os.path.join(sys.path[0], '..', 'src', 'lib')))
import market_maker_training
import market_maker_scoring

def test_training():
    coin_pair_dict = {'target':'btcusdt',
                    'alt':'ethusdt',
                    'through':'trxeth',
                    'excharb':'ethusdt'}

    mm_training = market_maker_training.CobinhoodTraining(coin_pair_dict, [5], [10])
    mm_training.training_data_sql

def test_scoring():
    mm_scoring = market_maker_scoring.CobinhoodScoring()
    # Set scoring data and retrieve the most recent minutes features
    mm_scoring.set_scoring_data(in_parallel=False)
    recent_df = mm_scoring.scoring_features_df.sort_values('open_time')
    # Check time diff
    latest_timestamp = recent_df.iloc[-1:]['close_time_x'].item() / 1000
    print(f"Last timestamp in scoring data: {latest_timestamp} compared to current: {time.time()} with {time.time() - latest_timestamp}")

if __name__ == '__main__':
    test_scoring()


import inspect
import os
import sys
from abc import ABC, abstractmethod
from datetime import datetime
from itertools import chain
from typing import List

from celery import shared_task

class AbstractDeployer(ABC):

    @abstractmethod
    def deployer_pre_exec_jobs(self) -> List[str]:
        pass

    @abstractmethod
    def deployer_exec_jobs(self) -> List[str]:
        pass

    @abstractmethod
    def deployer_post_exec_jobs(self) -> List[str]:
        pass

    @abstractmethod
    def check_prerequisite(self) -> bool:
        print("requirements check!")

    @staticmethod
    @abstractmethod
    def deploy(self):
        print("deploy!")

class PreExecContext:
    def __init__(self, step):
        self.step = step
        self.started = datetime.now()

class ExecContext:
    def __init__(self, step):
        self.step = step
        self.started = datetime.now()

class PostExecContext:
    def __init__(self, step):
        self.step = step
        self.started = datetime.now()

class BaseDeployer(AbstractDeployer, ABC):

    def deployer_pre_exec_jobs(self) -> List[str]:
        return []

    def deployer_exec_jobs(self) -> List[str]:
        return []

    def deployer_post_exec_jobs(self) -> List[str]:
        return []

    @staticmethod
    def check_prerequisite() -> bool:
        print("requirements check!")
        return True

    @staticmethod
    @shared_task(bind=True, autoretry_for=(Exception,), retry_backoff=True, retry_kwargs={"max_retries": 5},
                 name='deployers:deploy')
    def deploy(self, deployer_id: str, deployment_id: str):
        dir_path = os.path.dirname(os.path.realpath(__file__)) + os.sep
        for module in os.listdir(os.path.dirname(dir_path)):
            if module == '__init__.py' or module[-3:] != '.py' or module == str(__file__).removeprefix(dir_path):
                continue
            __import__(f'app.core.deployers.{module[:-3]}', locals(), globals())
        del module

        deployers_modules = [m for _, m in inspect.getmembers(sys.modules['app.core.deployers'], inspect.ismodule)]
        cls_members = list(chain.from_iterable([inspect.getmembers(m, inspect.isclass) for m in deployers_modules]))
        deployer = next(filter(lambda c: c[0] == deployer_id, cls_members))[1]
        check_method = getattr(deployer, 'check_prerequisite')
        check_prerequisite_res = check_method()
        if check_prerequisite_res:
            pre_jobs = getattr(deployer, 'deployer_pre_exec_jobs')
            pre_exec_context = PreExecContext(0)
            for job in pre_jobs():
                pre_exec_job = getattr(deployer, job)
                pre_exec_job(pre_exec_context)
                pre_exec_context.step += 1

            exec_jobs = getattr(deployer, 'deployer_exec_jobs')
            exec_context = ExecContext(0)
            for job in exec_jobs():
                exec_job = getattr(deployer, job)
                exec_job(exec_context)
                exec_context.step += 1

            post_jobs = getattr(deployer, 'deployer_post_exec_jobs')
            post_exec_context = PostExecContext(0)
            for job in post_jobs():
                post_exec_job = getattr(deployer, job)
                post_exec_job(post_exec_context)
                post_exec_context.step += 1
        else:
            raise RuntimeError(f'{type(deployer).__name__} deploy failed! Check Prerequisites Failed!')