import argparse
import logging
import os
import re
import shutil
import subprocess
import sys
import tempfile
from datetime import datetime

try:
    import yaml
except ImportError:
    import ruamel.yaml as yaml

import json

from pkg_resources import parse_version

try:
    import markdown

    HAVE_MARKDOWN = True
except ImportError:
    HAVE_MARKDOWN = False

DEFAULT_GIT_SITE = "https://github.com"
DEFAULT_GIT_ORG = "linux-system-roles"
CHANGELOG_HEADER = "Changelog\n=========\n\n"

if HAVE_MARKDOWN:

    class LSRMarkdown(markdown.Markdown):
        def convert(self, source):
            self.lines = source.split("\n")
            return self.parser.parseDocument(self.lines).getroot()

    class ChangelogManager(object):
        def __init__(self):
            # The keys of this dict are "New Features", "Bug Fixes",
            # and "Other Changes".  The value for each of these is a
            # dict.  The dict key is the rolename.  The dict value is the
            # list of changelog entries.
            self.updates = {}

        def addRoleChangelog(self, args, rolename, commit_msgs, cur_ref, new_ref):
            """Add a changelog for the role since the given tag."""
            if comp_versions(cur_ref, new_ref) >= 0:
                return
            if commit_msgs:
                # make a fake changelog since there are no new changelog entries
                self.updates.setdefault("Bug Fixes", {}).setdefault(
                    rolename, []
                ).extend(commit_msgs)
                return
            elif cur_ref is None:
                # just new feature new role
                self.updates.setdefault("New Features", {}).setdefault(
                    rolename, []
                ).insert(0, "New Role")
                return
            changelogmd = os.path.join(args.src_path, rolename, "CHANGELOG.md")
            if not os.path.exists(changelogmd):
                return
            data = open(changelogmd, encoding="utf-8").read()
            root = LSRMarkdown().convert(data)
            cur_tag_match = "[" + cur_ref + "]"
            new_tag_match = "[" + new_ref + "]"
            gathering = False
            current_hdr = None
            for elem in root:
                if elem.tag == "h2":  # tag + date
                    if elem.text.startswith(new_tag_match):
                        gathering = True
                    if elem.text.startswith(cur_tag_match):
                        return
                elif not gathering:
                    continue
                elif elem.tag == "h3":  # New Features, etc.
                    current_hdr = elem.text
                elif elem.tag == "ul":  # The list of entries
                    for entry in elem:
                        if entry.text != "none":
                            self.updates.setdefault(current_hdr, {}).setdefault(
                                rolename, []
                            ).insert(0, entry.text)

        def formatChangelogUpdate(self, tag, date):
            """Generate a new changelog update."""
            if not self.updates:
                return None
            rv = f"{CHANGELOG_HEADER}[{tag}] - {date}\n---------------------\n"
            hdrs = [hdr for hdr in ("New Features", "Bug Fixes") if hdr in self.updates]
            for hdr in hdrs:
                rv = rv + f"\n### {hdr}\n\n"
                for rolename in sorted(self.updates[hdr]):
                    for entry in self.updates[hdr][rolename]:
                        rv = rv + f"- {rolename} - {entry}\n"
            if not hdrs:
                rv = rv + "\n### Other Changes\n\n- no user-visible changes\n"
            return rv + "\n"

def run_cmd(cmdlist, cwd=None, env=None):
    """Run the given cmdlist using subprocess.  Debug log the output.
    If check is true, this function will work like check_call.
    The return value is like subprocess.run"""
    kwargs = {"env": os.environ}
    if env:
        kwargs["env"].update(env)
    rc = subprocess.run(
        cmdlist,
        cwd=cwd,
        encoding="utf-8",
        check=False,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        **kwargs,
    )
    logging.debug(
        f"{' '.join(cmdlist)} returned {rc.returncode} stdout {rc.stdout} stderr {rc.stderr}"
    )
    if rc.returncode != 0:
        raise subprocess.CalledProcessError(
            rc.returncode, cmdlist, rc.stdout, rc.stderr
        )
    return rc

def lsr_parse_version(v_str):
    test_v = parse_version("1")  # guaranteed to work
    if v_str is None:
        v = parse_version("0.0.0")
    else:
        v = parse_version(v_str)
    if not isinstance(v, test_v.__class__):
        raise ValueError(f"Error: {v_str} is not a proper version number")
    if not hasattr(v, "release"):
        if hasattr(v, "_version"):
            setattr(v, "release", v._version.release)
    return v

def check_versions_updated(cur_ref, new_ref, versions_updated):
    """Compare versions and update list.

    If cur_ref and new_ref are both valid semantic versions,
    compare them, and indicate in versions_updated if the major,
    minor, and/or micro versions were updated.  One exception is
    when the major version is upgraded from 0 to 1 - in that case,
    do not mark this as requiring a major version change.
    Another special case is adding a new role - in that case, cur_ref
    is None - we do not want to bump the major version, only the
    minor version."""

    if cur_ref is None:
        versions_updated[1] = True
        return
    try:
        cur_v = lsr_parse_version(cur_ref)
        new_v = lsr_parse_version(new_ref)
        for idx in range(0, 3):
            if cur_v.release[idx] != new_v.release[idx]:
                if idx > 0 or cur_v.release[idx] > 0:
                    versions_updated[idx] = True
    except ValueError as exc:
        logging.debug(f"Could not compare version {cur_ref} to {new_ref}: {exc}")
        if cur_ref != new_ref:
            versions_updated[3] = True

def comp_versions(cur_ref, new_ref):
    """Compare versions.

    Return values:
    -1 - if cur_ref < new_ref
     0 - if cur_ref == new_ref
     1 - if cur_ref > new_ref"""

    try:
        v_cur_ref = lsr_parse_version(cur_ref)
        v_new_ref = lsr_parse_version(new_ref)
        if v_cur_ref < v_new_ref:
            return -1
        elif v_cur_ref == v_new_ref:
            return 0
        else:
            return 1
    except ValueError as exc:
        logging.debug(f"Could not compare version {cur_ref} to {new_ref}: {exc}")
        if cur_ref == new_ref:
            return 0
        else:
            return -1  # assume new is "newer" than cur

def get_latest_tag_hash(args, rolename, cur_ref, org, repo, use_commit_hash):
    """
    Get the latest tag, hash, and tag_is_latest from the upstream repo.

    Clone and/or update the local copies of the repos.  Get
    the latest tag and/or commit hash for each repo.  Indicate if
    the tag is the latest commit.
    """
    roledir = os.path.join(args.src_path, rolename)
    # clone and/or update role repo
    if os.path.isdir(roledir):
        _ = run_cmd(["git", "fetch"], roledir)
    else:
        _ = run_cmd(
            [
                "git",
                "-c",
                "advice.detachedHead=false",
                "clone",
                "-q",
                f"{DEFAULT_GIT_SITE}/{org}/{repo}",
                roledir,
            ],
        )
    branch_output = run_cmd(["git", "branch", "-r"], roledir)
    # determine what is the main branch, check it out, and update it
    mmatch = re.search(r"origin/HEAD -> origin/(\w+)", branch_output.stdout)
    main_branch = mmatch.group(1)
    _ = run_cmd(["bash", "-c", f"git checkout {main_branch}; git pull --tags"], roledir)
    commit_msgs = ""
    # see if there have been any commits since the last time we checked
    if cur_ref:
        count_output = run_cmd(
            ["bash", "-c", f"git log --oneline {cur_ref}.. | wc -l"],
            roledir,
        )
    else:
        count_output = run_cmd(
            ["bash", "-c", "git log --oneline | wc -l"],
            roledir,
        )
    tag, commit_hash, n_commits, prev_tag = None, None, "0", None
    if count_output.stdout == "0":
        logging.debug(f"no changes to role {rolename} since ref {cur_ref}")
    else:
        # get latest tag and commit hash
        describe_cmd = ["git", "describe", "--tags", "--long", "--abbrev=40"]
        describe_output = run_cmd(describe_cmd, roledir)
        tag, n_commits, g_hash = describe_output.stdout.strip().rsplit("-", 2)
        # commit hash - skip leading "g"
        commit_hash = g_hash[1:]
        if n_commits != "0" and use_commit_hash:
            # get commit messages to use for changelog
            log_cmd = [
                "git",
                "log",
                "--oneline",
                "--no-merges",
                "--reverse",
                "--pretty=format:%s",
            ]
            if cur_ref:
                log_cmd.append(f"{cur_ref}..")
            log_output = run_cmd(log_cmd, roledir)
            commit_msgs = log_output.stdout.replace("\\r", "")
        # get previous tag in case cur_ref is a commit hash
        if cur_ref:
            try:
                describe_cmd = [
                    "git",
                    "describe",
                    "--tags",
                    "--long",
                    "--abbrev=40",
                    cur_ref,
                ]
                describe_output = run_cmd(describe_cmd, roledir)
                prev_tag = describe_output.stdout.strip().split("-")[0]
                if prev_tag == cur_ref:
                    prev_tag = None  # cur_ref was already a valid tag
            except subprocess.CalledProcessError:
                prev_tag = None  # no previous tag
        else:
            prev_tag = None  # no previous tag
    if args.no_update and cur_ref:
        # make sure cur_ref is checked out
        _ = run_cmd(["git", "checkout", cur_ref], roledir)
    return (tag, commit_hash, n_commits == "0", commit_msgs, prev_tag)

def process_ignore_and_lint_files(args, coll_dir):
    """Create collection ignore-VER.txt and .ansible-lint files from roles."""
    ignore_file_dir = os.path.join(coll_dir, "tests", "sanity")
    # role2collection is currently done using ruamel rt with a long line length,
    # because we can't ensure converted lines will be wrapped preserving
    # Jinja and ansible syntax.  This means we trigger a lot of warnings like
    # this when importing into Galaxy:
    # roles/file.yml:61: yaml line too long (187 > 160 characters)
    # the only thing we can do is suppress these
    ansible_lint = {"skip_list": ["yaml[line-length]"]}
    for role_name in os.listdir(args.src_path):
        roledir = os.path.join(args.src_path, role_name)
        if (
            os.path.isdir(roledir)
            and os.path.isfile(os.path.join(roledir, "tasks", "main.yml"))
            and role_name in args.include
        ):
            for file_name in os.listdir(roledir):
                if file_name.startswith(".sanity-ansible-ignore-"):
                    if not os.path.isdir(ignore_file_dir):
                        os.mkdir(ignore_file_dir)
                    match = re.match(
                        r"^[.]sanity-ansible-ignore-(\d[\d.]+)[.]txt$", file_name
                    )
                    if match and match.groups() and match.group(1):
                        ver = match.group(1)
                        ignore_file = os.path.join(ignore_file_dir, f"ignore-{ver}.txt")
                        role_ignore_file = os.path.join(roledir, file_name)
                        with open(ignore_file, "a") as ign_fd:
                            with open(role_ignore_file, "r") as role_ign_fd:
                                ign_fd.write(role_ign_fd.read())
                elif file_name == ".ansible-lint":
                    role_ansible_lint = yaml.safe_load(
                        open(os.path.join(roledir, file_name))
                    )
                    # ensure we suppress line-length checks
                    for key, items in role_ansible_lint.items():
                        if key not in ansible_lint:
                            ansible_lint[key] = role_ansible_lint[key]
                            continue
                        for item in items:
                            if item not in ansible_lint[key]:
                                if isinstance(ansible_lint[key], list):
                                    ansible_lint[key].append(item)
                                elif isinstance(ansible_lint[key], dict):
                                    ansible_lint[key][item] = items[item]
    yaml.safe_dump(ansible_lint, open(os.path.join(coll_dir, ".ansible-lint"), "w"))

def role_to_collection(
    args,
    rolename,
    namespace,
    collection_name,
    collection_readme,
):
    """Convert the role to collection format."""
    subrole_prefix = f"private_{rolename}_subrole_"
    cmd = [
        sys.executable,
        "lsr_role2collection.py",
        "--src-owner",
        args.src_owner,
        "--role",
        rolename,
        "--src-path",
        args.src_path,
        "--dest-path",
        args.dest_path,
        "--namespace",
        namespace,
        "--collection",
        collection_name,
        "--subrole-prefix",
        subrole_prefix,
        "--readme",
        collection_readme,
    ]
    extra_mapping = ""
    comma = ""
    if rolename == "sshd":
        # HACK - special case for ansible-sshd - not fully qualified
        extra_mapping = f"ansible-sshd:{namespace}.{collection_name}.{rolename}"
        comma = ","
    if args.extra_mapping:
        extra_mapping = extra_mapping + comma + args.extra_mapping
    if extra_mapping:
        cmd.extend(["--extra-mapping", extra_mapping])
    _ = run_cmd(cmd)

def update_galaxy_version(args, galaxy, versions_updated):
    if args.no_update:
        return
    if (
        not versions_updated[3]
        and any(versions_updated[0:3])
        and not args.no_auto_version
    ):
        galaxy_ver = lsr_parse_version(galaxy["version"])
        major, minor, micro = galaxy_ver.release
        if versions_updated[0]:
            major = major + 1
            minor = 0
            micro = 0
        elif versions_updated[1]:
            minor = minor + 1
            micro = 0
        elif versions_updated[2]:
            micro = micro + 1
        galaxy["version"] = f"{major}.{minor}.{micro}"
    elif args.new_version:
        galaxy["version"] = args.new_version
    else:
        raise Exception(
            "No new-version given and cannot auto-update version.  "
            "Please examine the changes to the collection, determine what "
            "the new semantic version will be, and re-run this script with "
            "the --new-version argument"
        )
    # not worth trying to figure out how to do a round-trip yaml write here
    # not work using ruamel here
    # this is simple enough for re.sub
    with open(args.galaxy_yml.name) as gf:
        galaxy_str = gf.read()
        galaxy_str = re.sub(
            r"^version:.*$", f'version: "{galaxy["version"]}"', galaxy_str, flags=re.M
        )
    with open(args.galaxy_yml.name, "w") as gf:
        gf.write(galaxy_str)

def create_collection_extra_files(args, coll_dir, galaxy=None):
    collection_requirements = os.path.join(
        "lsr_role2collection", "collection_requirements.txt"
    )
    collection_requirements_dest = os.path.join(coll_dir, "requirements.txt")
    collection_bindep = os.path.join("lsr_role2collection", "collection_bindep.txt")
    collection_bindep_dest = os.path.join(coll_dir, "bindep.txt")
    collection_req_yml_dest = os.path.join(coll_dir, "requirements.yml")
    collection_aee_dest = os.path.join(coll_dir, "execution-environment.yml")
    collection_yamllint = os.path.join("lsr_role2collection", "yamllint.yml")
    collection_yamllint_dest = os.path.join(coll_dir, ".yamllint.yml")
    shutil.copy(collection_yamllint, collection_yamllint_dest)
    # If --rpm, files such as galaxy.yml in coll_dir are being used.
    process_ignore_and_lint_files(args, coll_dir)
    if not args.rpm:
        if galaxy:
            with open(os.path.join(coll_dir, "galaxy.yml"), "w") as galaxy_fd:
                yaml.safe_dump(galaxy, galaxy_fd, explicit_start=True)
        else:
            shutil.copy(args.galaxy_yml.name, coll_dir)
        shutil.copy(args.collection_release_yml.name, coll_dir)
        ee_dependencies = []
        if os.path.exists(collection_requirements):
            shutil.copy(collection_requirements, collection_requirements_dest)
            ee_dependencies.append("python: requirements.txt")
        if os.path.exists(collection_bindep):
            shutil.copy(collection_bindep, collection_bindep_dest)
            ee_dependencies.append("system: bindep.txt")
        if galaxy["dependencies"]:
            with open(collection_req_yml_dest, "w") as crf:
                crf.write("---\ncollections:\n")
                for k in galaxy["dependencies"].keys():
                    crf.write("  - name: {0}\n".format(k))
            ee_dependencies.append("galaxy: requirements.yml")
        with open(collection_aee_dest, "w") as eef:
            eef.write("---\nversion: 1\n\n")
            eef.write("build_arg_defaults:\n")
            eef.write("  EE_BASE_IMAGE: {0}\n\n".format(args.ee_base_image))
            eef.write("ansible_config: 'ansible.cfg'\n\n")
            if len(ee_dependencies):
                eef.write("dependencies:\n")
                for eed in ee_dependencies:
                    eef.write("  {0}\n".format(eed))

def build_collection(args, coll_dir, galaxy=None):
    create_collection_extra_files(args, coll_dir, galaxy)
    # removing dot files/dirs
    keep_files = set([".ansible-lint", ".yamllint.yml"])
    for file_name in os.listdir(coll_dir):
        if file_name in keep_files:
            continue
        if re.match(r"^[.][a-zA-Z]+", file_name):
            full_name = os.path.join(coll_dir, file_name)
            if os.path.isdir(full_name):
                shutil.rmtree(full_name)
            else:
                os.unlink(full_name)

    if shutil.which("ansible-galaxy"):
        build_args = ["ansible-galaxy", "collection", "build", "-v"]
        if args.force:
            build_args.append("-f")
        build_args.append(coll_dir)
        _ = run_cmd(build_args, args.dest_path)
    else:
        logging.info("ansible-galaxy is skipped since it is not available.")

def conv_md2rest(changelog_md_path, changelog_rest_path):
    """
    Convert the md format to the reStructuredText format
    """
    rval = ""
    try:
        import pypandoc

        rval = pypandoc.convert_file(changelog_md_path, "rst", format="md")
    except (IOError, ImportError):
        raise Exception("Failed to import pypandoc module")
    with open(changelog_rest_path, "w", encoding="utf-8") as f:
        f.write(rval)
    # Verify the converted reST
    try:
        from rst2html import rst2html

        with open(changelog_rest_path, "r") as f:
            rst_text = f.read()
        html, warning = rst2html(rst_text, report_level=3)
        if len(warning) > 0:
            raise Exception(
                "Converted CHANGELOG.rst has a problem as html - {}".format(warning)
            )
    except (IOError, ImportError):
        logging.warning("Failed to import rst2html module")
        pass

def convert_md2rst(coll_dir):
    """
    Convert CHANGELOG.md to CHANGELOG.rst
    """
    coll_changelog_path = os.path.join(coll_dir, "docs", "CHANGELOG.md")
    if os.path.exists(coll_changelog_path):
        coll_changelog_rest_path = os.path.join(coll_dir, "CHANGELOG.rst")
        conv_md2rest(coll_changelog_path, coll_changelog_rest_path)
    else:
        logging.warning(
            "Missing collection CHANGELOG.md {0}".format(coll_changelog_path)
        )

def update_collection(args, galaxy, coll_rel):
    """
    Update refs in collection_release.yml.

    Use the latest tag for the ref.  If use_commit_hash is True
    and the latest commit is not tagged, use the commit hash of
    the latest commit for the ref.  Or if the role is specified
    using use_commit_hash_role.
    """
    coll_dir = os.path.join(
        args.dest_path, "ansible_collections", galaxy["namespace"], galaxy["name"]
    )
    if os.path.isdir(coll_dir):
        if args.keep:
            pass  # do nothing - keep whatever is there
        elif args.force:
            shutil.rmtree(coll_dir)
        else:
            raise Exception(
                "collection dest_path {} already exists - remove it, use --force "
                "to remove, or use --keep".format(coll_dir)
            )
    os.makedirs(coll_dir, exist_ok=True)
    collection_readme = os.path.join("lsr_role2collection", "collection_readme.md")
    if not args.skip_changelog:
        cl_manager = ChangelogManager()
    # major, minor, micro, hash
    versions_updated = [False, False, False, False]
    for rolename in args.include:
        if not args.skip_git:
            if args.use_commit_hash and rolename not in args.use_commit_hash_role:
                args.use_commit_hash_role.append(rolename)
            cur_ref = coll_rel[rolename]["ref"]
            tag, cm_hash, tag_is_latest, commit_msgs, prev_tag = get_latest_tag_hash(
                args,
                rolename,
                coll_rel[rolename]["ref"],
                coll_rel[rolename].get("org", args.src_owner),
                coll_rel[rolename].get("repo", rolename),
                rolename in args.use_commit_hash_role,
            )
            if tag or cm_hash:
                if tag_is_latest or rolename not in args.use_commit_hash_role:
                    coll_rel[rolename]["ref"] = tag
                else:
                    coll_rel[rolename]["ref"] = cm_hash
                if not tag_is_latest and not args.no_auto_version:
                    logging.debug(
                        f"role {rolename} tag {tag} is not the latest commit {cm_hash}"
                    )
                check_versions_updated(
                    cur_ref, coll_rel[rolename]["ref"], versions_updated
                )
            # If a version update is detected, retrieve the new changelogs from CHANGELOG.md
            if not args.skip_changelog:
                new_ref = coll_rel[rolename]["ref"]
                if comp_versions(cur_ref, new_ref) < 0:
                    if prev_tag is None:
                        prev_tag = cur_ref
                    logging.info(
                        "The role [%s] is updated. Getting changelog entries from [%s] to [%s].",
                        rolename,
                        prev_tag,
                        new_ref,
                    )
                    cl_manager.addRoleChangelog(
                        args, rolename, commit_msgs, prev_tag, new_ref
                    )
                else:
                    logging.info(
                        "No updates for role [%s] since [%s]", rolename, cur_ref
                    )
        role_to_collection(
            args,
            rolename,
            galaxy["namespace"],
            galaxy["name"],
            collection_readme,
        )
        this_collection = galaxy["namespace"] + "." + galaxy["name"]
        legacy_rqf = "requirements.yml"
        coll_rqf = "collection-requirements.yml"
        for rqf in [legacy_rqf, coll_rqf]:
            req_yml = os.path.join(args.src_path, rolename, "meta", rqf)
            if os.path.isfile(req_yml):
                req_yml_hsh = yaml.safe_load(open(req_yml))
                if isinstance(req_yml_hsh, list):
                    continue  # legacy role format
                if rqf == legacy_rqf:
                    logging.warning(
                        "The role %s is still using %s - please convert to %s instead",
                        rolename,
                        rqf,
                        coll_rqf,
                    )
                for coll in req_yml_hsh.get("collections", []):
                    if isinstance(coll, dict):
                        coll_name = coll["name"]
                        coll_ver = coll.get("version", "*")
                    else:
                        coll_name = coll
                        coll_ver = "*"
                    if coll_name != this_collection:
                        galaxy_deps = galaxy.setdefault("dependencies", {})
                        galaxy_deps[coll_name] = coll_ver
    # Existing changelogs
    orig_cl_file = "lsr_role2collection/COLLECTION_CHANGELOG.md"
    # Collection changelog file path
    coll_changelog_path = os.path.join(coll_dir, "docs", "CHANGELOG.md")
    if not args.no_update:
        if not any(versions_updated):
            logging.info(
                "Nothing in the collection was changed - current collection is up-to-date"
            )
            if not args.skip_changelog:
                shutil.copy(orig_cl_file, coll_changelog_path)
                # Convert CHANGELOG.md to CHANGELOG.rst
                if args.changelog_rst:
                    convert_md2rst(coll_dir)
            create_collection_extra_files(args, coll_dir, galaxy)
            return
        update_galaxy_version(args, galaxy, versions_updated)
        if not args.no_update:
            with open(args.collection_release_yml.name, "w") as crf:
                yaml.safe_dump(coll_rel, crf, sort_keys=False)

    # If to-be-appended changelogs are found, update COLLECTION_CHANGELOG.md
    # and copy it to the collection docs dir.
    if not args.skip_changelog:
        new_changelog_str = cl_manager.formatChangelogUpdate(
            galaxy["version"], datetime.now().date()
        )
        if new_changelog_str:
            changelog_str = open(orig_cl_file).read()
            changelog_str = changelog_str.replace(CHANGELOG_HEADER, new_changelog_str)
            open(orig_cl_file, "w").write(changelog_str)
            if args.save_current_changelog:
                # used for the release notes for the github release
                changelog_str = new_changelog_str.replace(CHANGELOG_HEADER, "")
                open("CURRENT_VER_CHANGELOG.md", "w").write(changelog_str)
            # Copy the new changelog to the docs dir in collection
            shutil.copy(orig_cl_file, coll_changelog_path)
        else:
            shutil.copy(orig_cl_file, coll_changelog_path)

        # Convert CHANGELOG.md to CHANGELOG.rst
        if args.changelog_rst:
            convert_md2rst(coll_dir)
    build_collection(args, coll_dir, galaxy)

def find(path, name):
    """Find a file 'name' in or under the directory 'path'."""
    for root, dirs, files in os.walk(path):
        if name in files:
            return os.path.join(root, name)

def process_rpm(args, default_galaxy, coll_rel):
    """
    Extract the contents of rpm.
    If the rpm contains galaxy.yml, use the file.
    Else if MANIFEST.json is found, use the info.
    Otherwise, it fails.
    Generate the collection artifact from the collection part
    of the extracted files.
    """
    workdir = tempfile.mkdtemp(suffix=".lsr", prefix="collection")
    cmdlist0 = ["rpm2cpio", args.rpm]
    p0 = subprocess.Popen(cmdlist0, cwd=workdir, stdout=subprocess.PIPE)
    cmdlist1 = ["cpio", "-id"]
    p1 = subprocess.Popen(cmdlist1, cwd=workdir, stdin=p0.stdout)
    p1.communicate()
    p0.communicate()
    logging.debug(f"{' '.join(cmdlist0)} returned {p0.returncode}")
    logging.debug(f"{' '.join(cmdlist1)} returned {p1.returncode}")
    if p0.returncode != 0:
        raise subprocess.CalledProcessError(
            p0.returncode, cmdlist0, p0.stdout, p0.stderr
        )
    if p1.returncode != 0:
        raise subprocess.CalledProcessError(
            p1.returncode, cmdlist1, p1.stdout, p1.stderr
        )

    tmp_coll = "{}/usr/share/ansible/collections/ansible_collections".format(workdir)
    if not os.path.exists(tmp_coll):
        raise Exception("Failed to extract {} from {}".format(tmp_coll, args.rpm))
    shutil.move(tmp_coll, workdir)
    shutil.rmtree("{}/usr".format(workdir))

    # If exists, use galaxy.yml in rpm
    galaxy_yml = find(workdir, "galaxy.yml")
    if galaxy_yml:
        args.galaxy_yml.close()
        args.galaxy_yml = open(galaxy_yml, "r")
        coll_dir = os.path.dirname(galaxy_yml)
        # override galaxy with the one in rpm
        galaxy = yaml.safe_load(args.galaxy_yml)
    # Otherwise, use the values in MANIFEST.json, if any.
    else:
        manifest_json = find(workdir, "MANIFEST.json")
        if manifest_json:
            with open(manifest_json) as mj:
                mj_contents = json.load(mj)
            galaxy = default_galaxy
            galaxy["namespace"] = mj_contents["collection_info"]["namespace"]
            galaxy["name"] = mj_contents["collection_info"]["name"]
            galaxy["version"] = mj_contents["collection_info"]["version"]
            galaxy["authors"] = mj_contents["collection_info"]["authors"]
            galaxy["description"] = mj_contents["collection_info"]["description"]
            coll_dir = os.path.dirname(manifest_json)
            new_galaxy_yml = "{}/galaxy.yml".format(coll_dir)
            args.galaxy_yml.close()
            args.galaxy_yml = open(new_galaxy_yml, "w")
            yaml.dump(galaxy, args.galaxy_yml)
            os.remove(manifest_json)
            files_json = "{}/FILES.json".format(coll_dir)
            if os.path.exists(files_json):
                os.remove(files_json)
        else:
            raise Exception("No galaxy.yml nor MANIFEST.json in {}".format(args.rpm))

    args.src_path = coll_dir
    build_collection(args, coll_dir)
    shutil.rmtree(workdir)

    return galaxy

def check_collection(args, galaxy):
    coll_file = (
        f"{args.dest_path}/{galaxy['namespace']}-{galaxy['name']}-"
        f"{galaxy['version']}.tar.gz"
    )
    gi_config = "lsr_role2collection/galaxy-importer.cfg"
    if os.path.exists(coll_file) and os.path.exists(gi_config):
        _ = run_cmd(
            [sys.executable, "-m", "galaxy_importer.main", coll_file],
            args.dest_path,
            {"GALAXY_IMPORTER_CONFIG": gi_config},
        )

def publish_collection(args, galaxy):
    coll_file = (
        f"{args.dest_path}/{galaxy['namespace']}-{galaxy['name']}-"
        f"{galaxy['version']}.tar.gz"
    )
    if os.path.exists(coll_file):
        cmd = ["ansible-galaxy", "collection", "publish"]
        if args.no_wait:
            cmd.append("--no-wait")
        else:
            cmd.append("-vv")
        cmd.append(coll_file)
        _ = run_cmd(cmd)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--galaxy-yml",
        type=open,
        default=os.environ.get("GALAXY_YML", "galaxy.yml"),
        help="Path/filename for galaxy.yml",
    )
    parser.add_argument(
        "--collection-release-yml",
        type=open,
        default=os.environ.get("COLLECTION_RELEASE_YML", "collection_release.yml"),
        help="Path/filename for collection_release.yml",
    )
    parser.add_argument(
        "--src-path",
        type=str,
        default=os.environ.get("COLLECTION_SRC_PATH"),
        help="Path to the directory containing the local clone of the role repos",
    )
    parser.add_argument(
        "--src-owner",
        type=str,
        default=os.environ.get("COLLECTION_SRC_OWNER", DEFAULT_GIT_ORG),
        help="Name of the role owner, in legacy role format OWNER.ROLENAME",
    )
    parser.add_argument(
        "--dest-path",
        type=str,
        default=os.environ.get(
            "COLLECTION_DEST_PATH", os.environ.get("HOME") + "/.ansible/collections"
        ),
        help=(
            "Path to collection base directory; collection directory structure "
            "will be created in DIR/ansible_collection/NAMESPACE/COLLECTION_NAME; "
            "collection package file will be created in DIR; "
            "default to ${HOME}/.ansible/collections"
        ),
    )
    parser.add_argument(
        "--force",
        default=False,
        action="store_true",
        help="Remove collection destination dir and file before creating",
    )
    parser.add_argument(
        "--keep",
        default=False,
        action="store_true",
        help="Keep collection destination dir and file before creating",
    )
    parser.add_argument(
        "--no-auto-version",
        default=False,
        action="store_true",
        help=(
            "The script will try to automatically determine the new collection version based "
            "on the semantic version changes in the versions of the tags in the "
            "role repos.  If you do not want this, add this parameter."
        ),
    )
    parser.add_argument(
        "--new-version",
        default=None,
        type=str,
        help="The new semantic version to use for the collection.",
    )
    parser.add_argument(
        "--no-update",
        default=False,
        action="store_true",
        help=(
            "By default, this script will update collection-release-yml "
            "to the latest tags or commit hashes in the upstream role.  Set "
            "this flag if you want to manually update collection-release-yml "
            "and build the collection from that file."
        ),
    )
    parser.add_argument(
        "--publish",
        default=False,
        action="store_true",
        help="Publish the collection to Galaxy.",
    )
    parser.add_argument(
        "--no-wait",
        default=False,
        action="store_true",
        help="Do not wait for the collection to be published.  Default is to wait.",
    )
    parser.add_argument(
        "--use-commit-hash",
        default=False,
        action="store_true",
        help=(
            "By default, only version tags are used for refs - if True, use latest HEAD"
            " commit hash if not tagged"
        ),
    )
    parser.add_argument(
        "--debug",
        default=False,
        action="store_true",
        help="Turn on debug logging.",
    )
    parser.add_argument(
        "--exclude",
        default=[],
        action="append",
        help="Roles to exclude",
    )
    parser.add_argument(
        "--include",
        default=[],
        action="append",
        help="Roles to include",
    )
    parser.add_argument(
        "--rpm",
        default=None,
        type=str,
        help="Path to the rpm file for the input.",
    )
    parser.add_argument(
        "--skip-git",
        default=False,
        action="store_true",
        help="True when work with local src.",
    )
    parser.add_argument(
        "--skip-check",
        default=False,
        action="store_true",
        help="True when skip check with galaxy-importer.",
    )
    parser.add_argument(
        "--use-commit-hash-role",
        default=[],
        action="append",
        help=(
            "Use the latest commit hash instead of the tag for these roles."
            "  Use this option when you want to use the tag for every role "
            "except the named roles e.g. --use-commit-hash-role sshd "
            "--use-commit-hash-role network will use the tag for every role"
            " except sshd and network, which will use the latest commit hash."
            "  Using --use-commit-hash is the same as using --use-commit-hash-role"
            " and specifying every role."
        ),
    )
    parser.add_argument(
        "--ee-base-image",
        default=os.environ.get(
            "EE_BASE_IMAGE", "quay.io/ansible/ansible-runner:latest"
        ),
        help="Base image for Ansible Execution Environment",
    )
    parser.add_argument(
        "--skip-changelog",
        default=False,
        action="store_true",
        help="If true, do not generate changelog or convert to rst.",
    )
    parser.add_argument(
        "--extra-mapping",
        type=str,
        default="",
        help=(
            "This is a comma delimited list of extra mappings to apply when "
            "converting roles to collection - this replaces the given role "
            "name with collection format with the optional given namespace "
            "and collection as well as the given FQCN with other FQCN."
        ),
    )
    parser.add_argument(
        "--changelog-rst",
        default=False,
        action="store_true",
        help="If true, generate CHANGELOG.rst in collection root directory.",
    )
    parser.add_argument(
        "--save-current-changelog",
        default=False,
        action="store_true",
        help="If true, save the changelog for the current collection version as CURRENT_VER_CHANGELOG.md",
    )
    args = parser.parse_args()

    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)
    else:
        logging.getLogger().setLevel(logging.INFO)

    galaxy = yaml.safe_load(args.galaxy_yml)
    args.galaxy_yml.close()
    coll_rel = yaml.safe_load(args.collection_release_yml)
    args.collection_release_yml.close()
    if not args.include:
        args.include = sorted(list(coll_rel.keys()))
    for item in args.exclude:
        del args.include[item]

    # --rpm rpm_file is given.
    if args.rpm:
        # Check rpm_file exists.
        if not os.path.exists(args.rpm):
            logging.error("{}: No such file given with --rpm".format(args.rpm))
            return
        # If rpm_file is relative path, convert it to absolute path.
        if not os.path.isabs(args.rpm):
            rpm_dir = os.path.abspath(os.getcwd())
            args.rpm = "{0}/{1}".format(rpm_dir, args.rpm)

    workdir = None
    if not args.rpm and not args.src_path:
        workdir = tempfile.mkdtemp(suffix=".lsr", prefix="collection")
        args.src_path = workdir
    try:
        if args.rpm:
            galaxy = process_rpm(args, galaxy, coll_rel)
        else:
            update_collection(args, galaxy, coll_rel)
        if not args.skip_check:
            check_collection(args, galaxy)
        else:
            logging.debug("check_collection is skipped.")
        if args.publish:
            publish_collection(args, galaxy)
        logging.info("Done.")
    finally:
        if workdir:
            shutil.rmtree(workdir)

if __name__ == "__main__":
    sys.exit(main())


import math
from typing import Optional, Tuple

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import CrossEntropyLoss

from transformers import PretrainedConfig
import torch.nn.functional as F
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import logging

from flash_attn.flash_attn_interface import flash_attn_varlen_func
from flash_attn.bert_padding import pad_input

logger = logging.get_logger(__name__)

class QWenConfig(PretrainedConfig):
    model_type = "qwen"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(
        self,
        vocab_size=151936,
        hidden_size=4096,
        num_hidden_layers=32,
        num_attention_heads=32,
        emb_dropout_prob=0.0,
        attn_dropout_prob=0.0,
        layer_norm_epsilon=1e-6,
        initializer_range=0.02,
        max_position_embeddings=4096,
        scale_attn_weights=True,
        use_cache=True,
        bf16=False,
        fp16=False,
        fp32=False,
        kv_channels=128,
        rotary_pct=1.0,
        rotary_emb_base=10000,
        use_dynamic_ntk=True,
        use_logn_attn=True,
        use_flash_attn="auto",
        intermediate_size=22016,
        no_bias=True,
        tie_word_embeddings=False,
        shift_attention=False,
        shift_attention_group_size=8192,
        shift_attention_group_size_ratio=1 / 4,
        neftune=False,
        neftune_alpha=5,
        rope_scaling=None,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.emb_dropout_prob = emb_dropout_prob
        self.attn_dropout_prob = attn_dropout_prob
        self.layer_norm_epsilon = layer_norm_epsilon
        self.initializer_range = initializer_range
        self.scale_attn_weights = scale_attn_weights
        self.use_cache = use_cache
        self.max_position_embeddings = max_position_embeddings
        self.bf16 = bf16
        self.fp16 = fp16
        self.fp32 = fp32
        self.kv_channels = kv_channels
        self.rotary_pct = rotary_pct
        self.rotary_emb_base = rotary_emb_base
        self.use_dynamic_ntk = use_dynamic_ntk
        self.use_logn_attn = use_logn_attn
        self.use_flash_attn = use_flash_attn
        self.no_bias = no_bias
        self.shift_attention = shift_attention
        self.shift_attention_group_size = shift_attention_group_size
        self.shift_attention_group_size_ratio = shift_attention_group_size_ratio
        self.neftune = neftune
        self.neftune_alpha = neftune_alpha
        """
        "rope_scaling": {
           "factor": 8.0,
           "type": "linear"/"static"
        }
        """
        self.rope_scaling = rope_scaling if rope_scaling is not None else {}
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)

@torch.jit.script  # type: ignore
def rms_norm(
    hidden_states: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
):
    input_dtype = hidden_states.dtype
    hidden_states = hidden_states.to(torch.float32)

    variance = (hidden_states * hidden_states).mean(-1, keepdim=True)
    hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
    return weight * hidden_states.to(input_dtype)

def rotate_half(x: torch.Tensor):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    position_ids: torch.Tensor,
):
    # q, k:     [nnz, num_heads, head_dim]
    # position_ids: [nnz]
    # cos, sin: [max_seq_len, head_dim]
    cos = cos[position_ids].unsqueeze(-2)  # [nnz, 1, head_dim]
    sin = sin[position_ids].unsqueeze(-2)  # [nnz, 1, head_dim]
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

class UnpaddedQWenRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps):
        super().__init__()
        self.variance_epsilon = eps
        self.weight = nn.Parameter(torch.ones(hidden_size))

    def forward(self, hidden_states):
        return rms_norm(hidden_states, self.weight, self.variance_epsilon)

class UnpaddedQWenRotaryEmbedding(torch.nn.Module):
    def __init__(self, dim, max_position_embeddings, base, device=None):
        super().__init__()

        # RoPE
        inv_freq = 1.0 / (
            base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)
        )
        t = torch.arange(max_position_embeddings, dtype=torch.float32, device=device)
        freqs = torch.outer(t, inv_freq)

        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        dtype = torch.get_default_dtype()
        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

    def forward(self):
        return self.cos_cached, self.sin_cached

class UnpaddedQWenMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.w1 = nn.Linear(
            config.hidden_size, config.intermediate_size // 2, bias=not config.no_bias
        )
        self.w2 = nn.Linear(
            config.hidden_size, config.intermediate_size // 2, bias=not config.no_bias
        )
        ff_dim_in = config.intermediate_size // 2
        self.c_proj = nn.Linear(ff_dim_in, config.hidden_size, bias=not config.no_bias)

    def forward(self, hidden_states):
        a1 = self.w1(hidden_states)
        a2 = self.w2(hidden_states)
        intermediate_parallel = a1 * F.silu(a2)
        output = self.c_proj(intermediate_parallel)
        return output

class UnpaddedQWenAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: QWenConfig):
        super().__init__()

        self.config = config
        self.hidden_size = config.hidden_size
        self.split_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.projection_size = config.kv_channels * config.num_attention_heads

        if (self.head_dim * self.num_heads) != self.hidden_size:
            raise ValueError(
                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
                f" and `num_heads`: {self.num_heads})."
            )

        self.c_attn = nn.Linear(config.hidden_size, 3 * self.projection_size)
        self.c_proj = nn.Linear(
            config.hidden_size, self.projection_size, bias=not config.no_bias
        )

    def forward(
        self,
        cos_sin: Tuple[torch.Tensor, torch.Tensor],
        # Unpadded inputs
        nz_hidden_states: torch.Tensor,
        position_ids: torch.LongTensor,
        cu_seqlens: torch.Tensor,
        max_seqlen: int,
    ) -> torch.Tensor:
        # nz_hidden_states: [nnz, num_heads, head_dim]
        # position_ids:  [nnz]
        # cu_seqlens:       [bs + 1]
        qkv_states = self.c_attn(nz_hidden_states)
        query, key, value = qkv_states.split(self.split_size, dim=2)
        query_states = query.view(-1, self.num_heads, self.head_dim)
        key_states = key.view(-1, self.num_heads, self.head_dim)
        value_states = value.view(-1, self.num_heads, self.head_dim)

        # RoPE
        cos, sin = cos_sin
        query_states, key_states = apply_rotary_pos_emb(
            query_states, key_states, cos, sin, position_ids
        )

        # flash attn
        attn_output = flash_attn_varlen_func(
            q=query_states,
            k=key_states,
            v=value_states,
            cu_seqlens_q=cu_seqlens,
            cu_seqlens_k=cu_seqlens,
            max_seqlen_q=max_seqlen,
            max_seqlen_k=max_seqlen,
            dropout_p=0.0,
            causal=True,
        )

        # attn_output: [total_nnz, num_heads, head_dim]
        attn_output = attn_output.view(-1, self.hidden_size)  # type: ignore
        return self.c_proj(attn_output)

class UnpaddedQWenDecoderLayer(nn.Module):
    def __init__(self, config: QWenConfig):
        super().__init__()
        self.hidden_size = config.hidden_size

        self.ln_1 = UnpaddedQWenRMSNorm(
            config.hidden_size, eps=config.layer_norm_epsilon
        )
        self.attn = UnpaddedQWenAttention(config=config)
        self.ln_2 = UnpaddedQWenRMSNorm(
            config.hidden_size, eps=config.layer_norm_epsilon
        )
        self.mlp = UnpaddedQWenMLP(config)

    def forward(
        self,
        cos_sin: Tuple[torch.Tensor, torch.Tensor],
        # Unpadded inputs
        nz_hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
        cu_seqlens: torch.Tensor,
        max_seqlen: int,
    ) -> torch.Tensor:
        # Self Attention
        residual = nz_hidden_states

        nz_hidden_states = self.ln_1(nz_hidden_states)
        nz_hidden_states = self.attn(
            cos_sin=cos_sin,
            nz_hidden_states=nz_hidden_states,
            position_ids=position_ids,
            cu_seqlens=cu_seqlens,
            max_seqlen=max_seqlen,
        )
        nz_hidden_states = residual + nz_hidden_states

        # Fully Connected
        residual = nz_hidden_states

        nz_hidden_states = self.ln_2(nz_hidden_states)
        nz_hidden_states = self.mlp(nz_hidden_states)
        nz_hidden_states = residual + nz_hidden_states

        return nz_hidden_states

class UnpaddedQWenPreTrainedModel(PreTrainedModel):
    config_class = QWenConfig
    base_model_prefix = "transformer"
    is_parallelizable = False
    supports_gradient_checkpointing = True
    _no_split_modules = ["UnpaddedQWenDecoderLayer"]

    def _init_weights(self, module):
        """Initialize the weights."""
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, UnpaddedQWenRMSNorm):
            module.weight.data.fill_(1.0)

        for name, p in module.named_parameters():
            if name == "c_proj.weight":
                p.data.normal_(
                    mean=0.0,
                    std=(
                        self.config.initializer_range
                        / math.sqrt(2 * self.config.num_hidden_layers)
                    ),
                )

    # def _set_gradient_checkpointing(self, module, value=False):
    #     if isinstance(module, UnpaddedQWenModel):
    #         module.gradient_checkpointing = value

class UnpaddedQWenModel(UnpaddedQWenPreTrainedModel):
    """
    Args:
        config: QWenConfig
    """

    def __init__(self, config: QWenConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.wte = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.rotary_emb = UnpaddedQWenRotaryEmbedding(
            dim=config.kv_channels,
            max_position_embeddings=config.max_position_embeddings,
            base=config.rotary_emb_base,
        )

        self.h = nn.ModuleList(
            [UnpaddedQWenDecoderLayer(config) for _ in range(config.num_hidden_layers)]
        )
        self.ln_f = UnpaddedQWenRMSNorm(
            config.hidden_size, eps=config.layer_norm_epsilon
        )

        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.wte

    def set_input_embeddings(self, value):
        self.wte = value

    def forward(
        self,
        # Unpadded inputs
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
        cu_seqlens: torch.Tensor,
        max_seqlen: int,
    ) -> torch.Tensor:
        nz_hidden_states = self.wte(input_ids)
        cos_sin = self.rotary_emb()

        # decoder layers
        for decoder_layer in self.h:
            if self.gradient_checkpointing and self.training:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        # None for past_key_value
                        return module(*inputs)

                    return custom_forward

                nz_hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(decoder_layer),
                    cos_sin,
                    nz_hidden_states,
                    position_ids,
                    cu_seqlens,
                    max_seqlen,
                )
            else:
                nz_hidden_states = decoder_layer(
                    cos_sin=cos_sin,
                    nz_hidden_states=nz_hidden_states,
                    position_ids=position_ids,
                    cu_seqlens=cu_seqlens,
                    max_seqlen=max_seqlen,
                )

        nz_hidden_states = self.ln_f(nz_hidden_states)

        return nz_hidden_states

class QWenLMHeadModel(UnpaddedQWenPreTrainedModel):
    # Ignore rotary emb inv_freq on load, as they will be calculated on creation
    _keys_to_ignore_on_load_unexpected = [
        r"transformer\.h\.\d+\.self_attn\.rotary_emb\.inv_freq"
    ]

    def __init__(self, config):
        super().__init__(config)
        self.transformer = UnpaddedQWenModel(config)

        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.transformer.wte

    def set_input_embeddings(self, value):
        self.transformer.wte = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def set_decoder(self, decoder):
        self.transformer = decoder

    def get_decoder(self):
        return self.transformer

    def forward(
        self,
        # Unpadded inputs
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
        cu_seqlens: torch.Tensor,
        max_seqlen: int,
        # Unpadded labels
        labels: Optional[torch.Tensor] = None,
    ) -> CausalLMOutputWithPast:
        # Model logits
        # print(
        #     f"input_ids: {input_ids.shape}, labels: {labels.shape}, position_ids: {position_ids.shape}, cu_seqlens: {cu_seqlens}, max_seqlen: {max_seqlen}"
        # )
        hidden_states = self.transformer(
            input_ids=input_ids,
            position_ids=position_ids,
            cu_seqlens=cu_seqlens,
            max_seqlen=max_seqlen,
        )
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            labels = labels.to(logits.device)
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(
                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
            )

        return CausalLMOutputWithPast(loss=loss, logits=logits)

class PaddedQWenForCausalLM(QWenLMHeadModel):
    """Compat layer for padded inputs"""

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        position_ids: Optional[torch.Tensor] = None,
        # unused
        return_dict: bool = True,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
    ):
        batch_size, seq_len = input_ids.shape
        if position_ids is None:
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 0)

        # get indices
        seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
        indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
        max_seqlen_in_batch = int(seqlens_in_batch.max().item())
        cu_seqlens = torch.nn.functional.pad(
            torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)
        )

        # Unpad inputs
        input_ids = torch.take_along_dim(input_ids, indices)
        position_ids = torch.take_along_dim(position_ids, indices)

        # Unpadded forward
        logits = (
            super()
            .forward(
                input_ids=input_ids,
                position_ids=position_ids,
                cu_seqlens=cu_seqlens,
                max_seqlen=max_seqlen_in_batch,
            )
            .logits
        )

        # Pad logits
        logits = pad_input(logits, indices, batch_size, seq_len)

        return CausalLMOutputWithPast(logits=logits)

    def prepare_inputs_for_generation(self, input_ids: torch.Tensor, **kwargs):
        return {
            "input_ids": input_ids,
            "attention_mask": kwargs.get("attention_mask"),
            "position_ids": kwargs.get("position_ids"),
        }

if __name__ == "__main__":
    model = QWenForCausalLM.from_pretrained(
        "/mnt/.cache/huggingface/hub/models--Qwen--Qwen-7B-Chat/snapshots/013d71a2b7be6824faf45c30534fbea570970468/"
    )


from typing import List, Optional, Union

import hvplot
import pandas

from .plot import Plot

class ParetoFront(Plot):
    """
    This class is used to create a pareto frontier plot of inference time v/s winrate, by framework.

    Attributes
    ----------
    plot_title: str,
        title of the plot on the dashboard website
    df_process: hvplot.Interactive,
        interactive pandas dataframe that is used to create the plot
        this should be the dataset of all frameworks aggregated across all datasets
    plot_type: str,
        type of hvplot. hvplot, table, pareto
    x_axis: Optional[Union[str, List[str]]],
        values to plot on x-axis
    y_axis: Optional[Union[str, List[str]]],
        values to plot on y-axis
    xlabel: str,
        label of x-axis
    ylabel: str,
        label of y-axis
    label_rot: int,
        rotation value of labels on axes

    Usage
    ------
    >>> pareto_front = ParetoFront("pareto plot", all_frameworks_df, "pareto", x_axis=TIME_INFER_S_RESCALED, y_axis=WINRATE)

    You can now call the `.plot()` method on this object to render the table as a Panel object on the dashboard website.

    """

    def __init__(
        self,
        plot_title: str,
        df_process: hvplot.Interactive,
        plot_type: str,
        x_axis: Optional[Union[str, List[str]]] = None,
        y_axis: Optional[Union[str, List[str]]] = None,
        graph_type: str = "bar",
        xlabel: str = "",
        ylabel: str = "",
        label_rot: int = 90,
    ) -> None:
        super().__init__(
            plot_title,
            df_process,
            plot_type,
            x_axis,
            y_axis,
            graph_type,
            xlabel,
            ylabel,
            label_rot,
        )


import argparse
import os
from sys import exit
from config import Config
from mangatracker import main

# Allow __main__ to run CLI without PySide2 installed
try:
    from mangatracker_gui import gui_main
    pyside2_installed = True
except ImportError:
    pyside2_installed = False

    def gui_main():
        print("PySide2 not installed, starting CLI")
        main()

def start_cli():
    main()
    exit()

def start_gui():
    gui_main()
    exit()

def start_default():
    """Starts default user interface as defined in config.ini"""
    if Config().default_to_gui and pyside2_installed:
        start_gui()
    else:
        start_cli()

if __name__ == "__main__":
    os.chdir(os.path.dirname(os.path.realpath(__file__)))

    parser = argparse.ArgumentParser(prog="MangaTracker",
                                     description="Track a manga collection.")
    group = parser.add_mutually_exclusive_group()
    group.add_argument("-g",
                       "--gui",
                       action="store_true",
                       help="Start GUI (Requires PySide2)")
    group.add_argument("-c",
                       "--cli",
                       action="store_true",
                       help="Start CLI")

    args = parser.parse_args()

    if args.gui:
        start_gui()
    if args.cli:
        start_cli()

    start_default()


import os

import torch
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter

from src.models.network import Net
from src.data.preprocess import testset
from src.constants import MODEL_DIR_PATH, CLASSES

class Evaluator:
    def __init__(self):
        self.load(os.path.join(MODEL_DIR_PATH, 'model.pt'))

        self.testloader = torch.utils.data.DataLoader(
            testset, batch_size=4, shuffle=False
        )
        self.writer = SummaryWriter('runs/evaluate')

    def load(self, checkpoint_path):
        self.net = Net()
        self.net.load_state_dict(torch.load(checkpoint_path, map_location=lambda storage, loc: storage))

    def evaluate(self):
        class_probs = []
        class_label = []
        with torch.no_grad():
            for data in self.testloader:
                images, labels = data
                output = self.net(images)
                class_probs_batch = [F.softmax(el, dim=0) for el in output]

                class_probs.append(class_probs_batch)
                class_label.append(labels)

        test_probs = torch.cat([torch.stack(batch) for batch in class_probs])
        test_label = torch.cat(class_label)

        return test_label, test_probs

    def add_pr_curve_tensorboard(self, global_step=0):
        test_label, test_probs = self.evaluate()

        for class_index in range(len(CLASSES)):
            tensorboard_truth = test_label == class_index
            tensorboard_probs = test_probs[:, class_index]

            self.writer.add_pr_curve(
                CLASSES[class_index],
                tensorboard_truth,
                tensorboard_probs,
                global_step=global_step
            )
            self.writer.close()


"""
Created on Mon Sep 16 18:31:35 2013
PHYS 613, Assignment 3
Nick Crump
"""

# Problem 2 (EBB2)
"""
Generating the density of states (DOS) by finding energy eigenvalues of
a particle in a 2D infinite square well.
"""

import numpy as np
from math import exp
from LeastSquares import gaussNewton3
import matplotlib.pyplot as plt

# Part A
# generate the lowest 500,000 energy eigenvalues from E(nx,ny) = nx**2 + ny**2
# ----------------------------------------------------------------------------
nstates = 500000  # keep this many states
nxy = 1000        # iterate over this many to make sure to get 500,000 lowest
states = []       # stores all E(nx,ny)

# iterate to get all E(nx,ny)
for i in range(1,nxy):
    for j in range(1,nxy):
        s = i**2 + j**2
        states.append(s)

# sort E(nx,ny) list and keep only lowest 500,000
states = np.sort(states)
states = states[0:nstates]

# plot histogram of density of states (DOS) vs energy
plt.figure(1)
plt.hist(states,bins=20,label='2D Infinite Well')
plt.xlabel('Energy')
plt.ylabel('Density of States')
plt.legend()

# Part B & C
# find energy differences E(i+1)-E(i) and E(i+2)-E(i) up to i = 100000
# ----------------------------------------------------------------------------
# keep duplicates (degeneracies) from array of energy states
num = 100000

# get E(i+1) - E(i)
delta1 = [states[i+1] - states[i] for i in range(num)]
# get E(i+2) - E(i)
delta2 = [states[i+2] - states[i] for i in range(num)]

# plot histogram of energy differences vs energy
plt.figure(2)
plt.subplot(211)
# gets the histogram data and plots histogram
n1, bins1, patches1 = plt.hist(delta1,bins=20,label='First Neighboring Levels')
plt.xlim(0,20)
plt.ylabel(r'Frequency of $\Delta E_{1}$')
plt.legend()

plt.subplot(212)
# gets the histogram data and plots histogram
n2, bins2, patches2 = plt.hist(delta2,bins=20,label='Second Neighboring Levels')
plt.xlim(0,20)
plt.xlabel(r'$\Delta\ $ Energy')
plt.ylabel(r'Frequency of $\Delta E_{2}$')
plt.legend()

# Part D
# fitting parameters of a known function to distributions
# ----------------------------------------------------------------------------
# guess at parameters for exponential decay f = A*exp(-Bx) to fit dataset 1
A1 = 75600
B1 = 1.1
x1 = np.linspace(min(bins1),max(bins1),1000)
exp1 = [A1*exp(-B1*i) for i in x1]

# guess at parameters for exponential decay f = A*exp(-Bx) to fit dataset 1
A2 = 52100
B2 = 0.7
x2 = np.linspace(min(bins2),max(bins2),1000)
exp2 = [A2*exp(-B2*i) for i in x2]

# plot histogram of energy differences vs energy with exponential fits
plt.figure(3)
plt.subplot(211)
plt.hist(delta1,bins=20,label='First Neighboring Levels')
plt.plot(x1+0.7,exp1,'r-',linewidth=3,label='Exponential Fit')
plt.xlim(0,20)
plt.ylabel(r'Frequency of $\Delta E_{1}$')
plt.legend()

plt.subplot(212)
plt.hist(delta2,bins=20,label='Second Neighboring Levels')
plt.plot(x2+0.7,exp2,'r-',linewidth=3,label='Exponential Fit')
plt.xlim(0,20)
plt.xlabel(r'$\Delta\ $ Energy')
plt.ylabel(r'Frequency of $\Delta E_{2}$')
plt.legend()


import pandas as pd
from settings import FEATURE_CHIN_ENG_DICT
import matplotlib.pyplot as plt
import numpy as np

path = '/media/reid/ext_disk1/blue_alage/dushu/geography_point.csv'

df = pd.read_csv(path)

lon_min, lon_max = df['lon'].min(), df['lon'].max()
lat_min, lat_max = df['lat'].min(), df['lat'].max()
range_step = 0.001

grid_lon = np.arange(lon_min, lon_max, range_step)
grid_lat = np.arange(lat_min, lat_max, range_step)

result = np.array([[(xi, yj) for yj in grid_lat] for xi in grid_lon])

whole_data = result.reshape(-1, 2)

lake_data = df[['lon', 'lat']].to_numpy()

result = np.any(np.all(whole_data[:, np.newaxis, :] == lake_data, axis=2), axis=1).astype(int)
ss = np.intersect1d(whole_data, lake_data)

fig, ax = plt.subplots()

# ax.scatter(combined_coordinates[:, 0], combined_coordinates[:, 1], color='green')

df.plot(kind='scatter', x='lon', y='lat', ax=ax)

plt.legend()

plt.show()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys

# if not calling for snakemake rule
try:
    sys.stderr = open(snakemake.log[0], "w")
except NameError:
    pass

import json

def main(mlst, metadata, assembly_qc, isolate_id, json_path):
    # load jsons
    with open(metadata, 'r') as fp:
        meta_dict = json.load(fp)
    with open(assembly_qc, 'r') as fp:
        assembly_dict = json.load(fp)
    with open(mlst, 'r') as fp:
        mlst_dict = json.load(fp)[0]  # mlst res as list of dict
    # get relevant entries from metadata
    dictout = meta_dict[isolate_id]
    # NB: fastas are renamed to match isolated_id
    dictout["fasta_name"] = f"{isolate_id}.fa"
    # reformat metadata
    dictout['isolation'] = {
        'org_name': dictout.pop('isolation_org'),
    }
    dictout['sequencing'] = {
        "org_name": dictout.pop('sequencing_org'),
        "extraction_method": dictout.pop('extraction_method'),
        "library_method": dictout.pop('library_method'),
        "sequencing_instrument": dictout.pop('sequencing_instrument'),
    }
    dictout['bioinformatics'] = {
        "org_name": dictout.pop('bioinformatics_org'),
        "assembly_method": dictout.pop('assembly_method'),
    }
    dictout['epidata'] = {
        "collection_date": dictout.pop('collection_date'),
        "municipality": dictout.pop('collection_municipality'),
        "country": dictout.pop('collection_country'),
        "cause": dictout.pop('collection_cause'),
        "collected_by": dictout.pop('collected_by'),
        "manufacturer": dictout.pop('manufacturer'),
        "designation": dictout.pop('designation'),
        "manufacturer_type": dictout.pop('manufacturer_type'),
        "sample_description": dictout.pop('sample_description'),
        "lot_number": dictout.pop('lot_number'),
    }
    # Add MLST infos
    dictout['mlst'] = {
        key: mlst_dict[key] for key in ['sequence_type', 'scheme', 'alleles']
    }
    # Add QC metrics
    dictout['qc_metrics'] = {
       key: assembly_dict[key] for key in [
            "sequencing_depth",
            "ref_coverage",
            "q30",
            "N50",
            "L50",
            "n_contigs_1kbp",
            "assembly_size",
            "GC_perc",
            "orthologs_found",
            "duplicated_orthologs",
            "majority_genus",
            "fraction_majority_genus",
            "majority_species",
            "fraction_majority_species",
        ]
    }
    # cleaning duplicated keys at top level
    for key in ["sequencing_depth", "ref_coverage", "q30"]:
        del dictout[key]
    # export to json
    with open(json_path, 'w') as fp:
        json.dump(dictout, fp, indent=4)

if __name__ == '__main__':
    main(
        snakemake.input['mlst'],
        snakemake.input['metadata'],
        snakemake.input['assembly_qc'],
        snakemake.params['isolate_id'],
        snakemake.output['json']
    )


import os, psutil, time, signal, sqlite3, config, subprocess, json
from ruamel.yaml import YAML
import girder_client
import xml.etree.ElementTree as ET
import girder_client

def check_pid(pid):
    try:
        os.kill(pid, 0)

        if (psutil.Process(pid).status() == psutil.STATUS_ZOMBIE):
            return False

        return True
    except OSError:
        return False

def get_wf_status(pid, guid):
    if (pid == None):
        return "NOT_YET_EXECUTED"
    if (pid<0):
        return "TERMINATED"

    if (check_pid(pid)):
        return "RUNNING"

    input_path = os.path.join(config.INPUTS,guid)

    dir_name = os.path.join(config.RESULTS, guid)

    for root, directories, files in os.walk(dir_name):
        directories[:] = [d for d in directories if d not in ['tmp']]
        for filename in files:
            return "FINISHED_OK"

    return "FINISHED_ERROR"

def terminate(GUID, pid, flag, timeout = 0):
    time.sleep(timeout)
    parent = psutil.Process(pid)
    for child in parent.children(recursive=True):  # or parent.children() for recursive=False
        os.kill(child.pid, signal.SIGINT)
    os.kill(pid, signal.SIGINT)
    conn = sqlite3.connect(config.DATABASE)
    c = conn.cursor()

    c.execute('UPDATE workflows SET PID='+str(flag)+' WHERE GUID="'+GUID+'"')

    conn.commit()
    conn.close()

def get_wf_pid(GUID):
    conn = sqlite3.connect(config.DATABASE)
    c = conn.cursor()
    c.execute('SELECT * FROM workflows WHERE GUID="'+GUID+'"')
    row = c.fetchone()

    conn.close()
    if (row is None):
        return None

    return row[3]

def get_wf(GUID):
    conn = sqlite3.connect(config.DATABASE)
    c = conn.cursor()
    c.execute('SELECT * FROM workflows WHERE GUID="'+GUID+'"')
    row = c.fetchone()

    conn.close()
    return row

def validate_yaml(input_dir):
    yaml = YAML()
    cwd = os.getcwd()
    os.chdir(input_dir)
    missing = list()
    with open('inputs.yaml') as file:

        yaml_in = yaml.load(file)
        for key in yaml_in:
            if (yaml_in[key] is not None and key!="girderIds"):
                if isinstance(yaml_in[key] , list):
                    for d in yaml_in[key]:
                        if (not os.path.isfile(d['path'])):
                            missing.append(d['path'])
                if isinstance(yaml_in[key],dict):
                    if (not os.path.isfile(yaml_in[key]['path'])):
                        missing.append(yaml_in[key]['path'])

    os.chdir(cwd)
    return missing

def getGirderIds(input_dir):
    yaml = YAML()
    cwd = os.getcwd()
    os.chdir(input_dir)
    with open('inputs.yaml') as file:
        os.chdir(cwd)
        yaml_in = yaml.load(file)
        if yaml_in.get('girderIds') is not None:
            return yaml_in['girderIds']
        else:
            return []

def downloadGirderFile(girderId, pathToInputs, girder_api_key):
    gc = girder_client.GirderClient(apiUrl=config.GIRDER_API)
    gc.authenticate(apiKey=girder_api_key)
    girderFile = gc.getFile(girderId)
    gc.downloadItem(girderFile['itemId'], pathToInputs)

def uploadToGirder(folderPath, girder_api_key, girder_parent_id):
    gc = girder_client.GirderClient(apiUrl=config.GIRDER_API)
    gc.authenticate(apiKey=girder_api_key)

    gc.upload(folderPath, girder_parent_id)

    ls = gc.listFolder(girder_parent_id, name=folderPath.split("/")[-1])
    return list(ls)[0]

def ifReadOnlyWf(workflow):
    jsonFile = json.load(open("read-only-wfs.json",'r'))
    wfs = jsonFile['wfs']

    if (workflow in wfs):
        return True
    return False

def getWfOutputDir(GUID):
    return os.path.join(config.RESULTS,GUID)

def getWfInputDir(GUID):
    return os.path.join(config.INPUTS,GUID)

def getDownloadOutputLink(GUID,path):
    link = "get-output-file?GUID="+GUID+"&filepath="+path
    return link

def getDownloadInputLink(GUID,path):
    link = "get-input-file?GUID="+GUID+"&filepath="+path
    return link

def getZipLink(GUID):
    link = "get-results?GUID="+GUID
    return link

def getGirderOutputDir(wf):
    if (wf[4] is None):
        return wf[0]
    try:
        tree = ET.fromstring(wf[4])
        return tree.find('output').text
    except:
        return wf[0]

def deleteFolderFromGirder(wf, girder_api_key):

    try:
        gc = girder_client.GirderClient(apiUrl=config.GIRDER_API)
        userId = gc.authenticate(apiKey=girder_api_key)['_id']
        privateFolder = gc.loadOrCreateFolder('Private', userId, parentType="user")

        outFolder = gc.loadOrCreateFolder('workflow-outputs', privateFolder['_id'], parentType="folder")
        #wf = get_wf(guid)
        wfFolder = gc.loadOrCreateFolder(wf[2], outFolder['_id'], parentType="folder")

        outputFolder = gc.loadOrCreateFolder(getGirderOutputDir(wf), wfFolder['_id'], parentType="folder")
        gc.delete("folder/"+outputFolder['_id'])
    except:
        return


# SGAM: Building a Virtual 3D World through Simultaneous Generation and Mapping
# Authored by Yuan Shen, Wei-Chiu Ma and Shenlong Wang
# University of Illinois at Urbana-Champaign and Massachusetts Institute of Technology

import json
import numpy as np
from torch.utils.data import Dataset
from pathlib import Path
import networkx as nx
from PIL import Image
import pickle
import os

class PRNGMixin(object):
    """Adds a prng property which is a numpy RandomState which gets
    reinitialized whenever the pid changes to avoid synchronized sampling
    behavior when used in conjunction with multiprocessing."""

    @property
    def prng(self, seed=None):
        currentpid = os.getpid()
        if getattr(self, "_initpid", None) != currentpid:
            self._initpid = currentpid
            self._prng = np.random.RandomState(seed=seed) if seed is not None else np.random.RandomState()
        return self._prng

class Blender3dBase(Dataset, PRNGMixin):
    def __init__(self, split, dataset_dir, n_src=2, dataset=None, image_resolution=None):
        self.split = split
        self.image_resolution = image_resolution
        self.dataset = dataset
        self.dataset_dir = dataset_dir
        self.src_num = n_src
        self.grids = []
        self.cumulative_sum = [0]
        self.K = np.load(f"{self.dataset_dir}/K.npy")
        for grid_transform_path in sorted(Path(self.dataset_dir, self.split).glob("*")):
            print(grid_transform_path)
            with open(str(grid_transform_path / "transforms.json"), 'r') as f:
                curr_transform = json.load(f)
                g = self.build_graph_from_transform(curr_transform['frames'], grid_transform_path)
                self.grids.append(g)
                self.cumulative_sum.append(len(g.nodes) + self.cumulative_sum[-1])

    def build_graph_from_transform(self, transforms, grid_transform_path):
        g_path = f"{self.dataset_dir}/cache/{grid_transform_path.name[:-4]}_graph_{self.split}.txt"
        if os.path.exists(g_path):
           return pickle.load(open(g_path, 'rb'))
        g = nx.Graph()
        for i, transform in enumerate(transforms):
            c2w = np.array(transform['transform_matrix']) @ np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
            w2c = np.linalg.inv(c2w)
            g.add_nodes_from([(i, {
                                    "R": w2c[:3, :3],
                                    "t": w2c[:3, 3],
                                    "position": c2w[:3, 3],
                                    "rgb_path": str(grid_transform_path / f"im_{i:05d}.png"),
                                    "depth_path": str(grid_transform_path / f"dm_{i:05d}.npy"),
                                    })])

        for i in range(len(transforms)-1):
            for j in range(i+1, len(transforms)):
                node_i = g.nodes[i]
                node_j = g.nodes[j]
                v = node_i['position'] - node_j['position']
                euclidean_dis = np.linalg.norm(v)
                if euclidean_dis <= 3: #TODO: to add as hyperparameters
                    g.add_edge(i, j, weight=euclidean_dis)
        pickle.dump(g, open(g_path, 'wb'))
        return g

    def __len__(self):
        return self.cumulative_sum[-1]

    def parse_idx(self, idx):
        for img_index, cur_cumsum in enumerate(self.cumulative_sum):
            if idx < self.cumulative_sum[img_index+1]:
                return img_index, idx - cur_cumsum

    def __getitem__(self, global_index):
        grid_id, idx = self.parse_idx(global_index)
        tgt_node = self.grids[grid_id].nodes[idx]
        tgt_neighbors = sorted(self.grids[grid_id][idx])
        if self.split == 'train':
            src_num = self.src_num
            src_nodes = [self.grids[grid_id].nodes[tgt_neighbors[k]] for k in self.prng.choice(len(tgt_neighbors), src_num)]
        else:
            state = np.random.RandomState(seed=global_index)
            tgt_neighbors = np.array(tgt_neighbors)
            state.shuffle(tgt_neighbors)
            src_num = self.src_num
            src_nodes = [self.grids[grid_id].nodes[k] for k in tgt_neighbors[:src_num]]
        img_dst = np.array(Image.open(tgt_node['rgb_path']))/127.5-1.0
        img_srcs = [np.array(Image.open(src_node['rgb_path']))/127.5-1.0 for src_node in src_nodes]

        dm_dst = np.load(tgt_node['depth_path'])#[..., None]
        dm_srcs = [np.load(src_node['depth_path']) for src_node in src_nodes]
        h, w = dm_dst.shape[:2]
        x = np.linspace(0, w - 1, w)
        y = np.linspace(0, h - 1, h)
        xs, ys = np.meshgrid(x, y)
        dm_dst = (dm_dst * self.K[0][0] / np.sqrt(self.K[0][0]**2 + (self.K[0][2] - ys - 0.5)**2 + (self.K[1][2] - xs - 0.5) ** 2))[..., None]
        dm_srcs = [(dm_src * self.K[0][0] / np.sqrt(self.K[0][0]**2 + (self.K[0][2] - ys - 0.5)**2 + (self.K[1][2] - xs - 0.5) ** 2))[..., None] for dm_src in dm_srcs]
        # for j in range(predicted_depth.shape[0]):
        #     for k in range(predicted_depth.shape[1]):
        #         predicted_depth[j, k] = predicted_depth[j, k] * K[0][0] / math.sqrt(K[0][0] ** 2 + (K[0][2] - j - 0.5) ** 2 + (K[1][2] - k - 0.5) ** 2)
        R_dst = tgt_node["R"]
        t_dst = tgt_node["t"]
        R_rels = []
        t_rels = []
        Ks = []
        K_invs = []
        T_tgt = np.eye(4)
        T_tgt[:3, :3] = R_dst
        T_tgt[:3, 3] = t_dst

        ## K
        h, w = img_dst.shape[:2]
        K = self.K
        K = K * self.image_resolution[1] / w
        K = K * self.image_resolution[0] / h
        for src_node in src_nodes:
            R_src = src_node["R"]
            t_src = src_node["t"]
            T_src = np.eye(4)
            T_src[:3, :3] = R_src
            T_src[:3, 3] = t_src
            T_rel = T_tgt @ np.linalg.inv(T_src)
            R_rel = T_rel[:3, :3]
            t_rel = T_rel[:3, 3]
            R_rels.append(R_rel)
            t_rels.append(t_rel)
            Ks.append(K)
            K_invs.append(np.linalg.inv(K))

        if self.image_resolution is not None and (self.image_resolution[0] != h or self.image_resolution[1] != w):
            ## img
            for i in range(len(img_srcs)):
                img_srcs[i] = img_srcs[i].resize((self.image_resolution[1], self.image_resolution[0]), resample=Image.LANCZOS)
            img_dst = img_dst.resize((self.image_resolution[1], self.image_resolution[0]), resample=Image.LANCZOS)

            ## depth
            for i in range(len(dm_srcs)):
                dm_srcs[i] = dm_srcs[i].resize((self.image_resolution[1], self.image_resolution[0]), resample=Image.LANCZOS)
            dm_dst = dm_dst.resize((self.image_resolution[1], self.image_resolution[0]), resample=Image.LANCZOS)
        mask = np.zeros(self.src_num)
        mask[:src_num] = 1
        while len(K_invs) < self.src_num:
            Ks.append(np.eye(3))
            K_invs.append(np.eye(3))
            R_rels.append(np.eye(3))
            t_rels.append(np.zeros(3))
            img_srcs.append(np.zeros_like(img_srcs[-1]))
            dm_srcs.append(np.zeros_like(dm_srcs[-1]))

        example = {
            "Ks": np.stack(Ks),
            "K_invs": np.stack(K_invs),
            "R_rels": np.stack(R_rels),
            "t_rels": np.stack(t_rels),
            "dst_img": img_dst, # rgb or rgbid
            "src_imgs": np.stack(img_srcs),
            "dst_depth": dm_dst,
            "src_depths": np.stack(dm_srcs),
            'src_masks': mask
        }

        for k in example:
            example[k] = example[k].astype(np.float32)

        return example

class Blender3dTrain(Blender3dBase):
    def __init__(self, dataset_dir=None, n_src=2, dataset=None, image_resolution=None):
        super().__init__(split='train', dataset_dir=dataset_dir, n_src=n_src, dataset=dataset, image_resolution=image_resolution)

class Blender3dValidation(Blender3dBase):
    def __init__(self, dataset_dir=None, n_src=2, dataset=None, image_resolution=None):
        super().__init__(split='val', n_src=n_src, dataset_dir=dataset_dir, dataset=dataset, image_resolution=image_resolution)

class Blender3dTest(Blender3dBase):
    def __init__(self, dataset_dir=None, n_src=2, dataset=None, image_resolution=None):
        super().__init__(split='test', n_src=n_src, dataset_dir=dataset_dir, dataset=dataset, image_resolution=image_resolution)


import dataclasses
from typing import Tuple

from acme import types
from acme.agents.jax.dqn import learning_lib
from acme.jax import networks as networks_lib
import chex
import jax
import jax.numpy as jnp
import reverb
import rlax

@dataclasses.dataclass
class PrioritizedDoubleQLearning(learning_lib.LossFn):
  """Clipped double q learning with prioritization on TD error."""
  discount: float = 0.99
  importance_sampling_exponent: float = 0.2
  max_abs_reward: float = 1.
  huber_loss_parameter: float = 1.

  def __call__(
      self,
      network: networks_lib.TypedFeedForwardNetwork,
      params: networks_lib.Params,
      target_params: networks_lib.Params,
      batch: reverb.ReplaySample,
      key: networks_lib.PRNGKey,
  ) -> Tuple[jax.Array, learning_lib.LossExtra]:
    """Calculate a loss on a single batch of data."""
    transitions: types.Transition = batch.data
    probs = batch.info.probability

    # Forward pass.
    key1, key2, key3 = jax.random.split(key, 3)
    q_tm1 = network.apply(
        params, transitions.observation, is_training=True, key=key1)
    q_t_value = network.apply(
        target_params, transitions.next_observation, is_training=True, key=key2)
    q_t_selector = network.apply(
        params, transitions.next_observation, is_training=True, key=key3)

    # Cast and clip rewards.
    d_t = (transitions.discount * self.discount).astype(jnp.float32)
    r_t = jnp.clip(transitions.reward, -self.max_abs_reward,
                   self.max_abs_reward).astype(jnp.float32)

    # Compute double Q-learning n-step TD-error.
    batch_error = jax.vmap(rlax.double_q_learning)
    td_error = batch_error(q_tm1, transitions.action, r_t, d_t, q_t_value,
                           q_t_selector)
    batch_loss = rlax.huber_loss(td_error, self.huber_loss_parameter)

    # Importance weighting.
    importance_weights = (1. / probs).astype(jnp.float32)
    importance_weights **= self.importance_sampling_exponent
    importance_weights /= jnp.max(importance_weights)

    # Reweight.
    loss = jnp.mean(importance_weights * batch_loss)  # []
    extra = learning_lib.LossExtra(
        metrics={}, reverb_priorities=jnp.abs(td_error).astype(jnp.float64))
    return loss, extra

@dataclasses.dataclass
class QrDqn(learning_lib.LossFn):
  """Quantile Regression DQN.

  https://arxiv.org/abs/1710.10044
  """
  num_atoms: int = 51
  huber_param: float = 1.0

  def __call__(
      self,
      network: networks_lib.TypedFeedForwardNetwork,
      params: networks_lib.Params,
      target_params: networks_lib.Params,
      batch: reverb.ReplaySample,
      key: networks_lib.PRNGKey,
  ) -> Tuple[jax.Array, learning_lib.LossExtra]:
    """Calculate a loss on a single batch of data."""
    transitions: types.Transition = batch.data
    key1, key2 = jax.random.split(key)
    _, dist_q_tm1 = network.apply(
        params, transitions.observation, is_training=True, key=key1)
    _, dist_q_target_t = network.apply(
        target_params, transitions.next_observation, is_training=True, key=key2)
    batch_size = len(transitions.observation)
    chex.assert_shape(
        dist_q_tm1, (
            batch_size,
            None,
            self.num_atoms,
        ),
        custom_message=f'Expected (batch_size, num_actions, num_atoms), got: {dist_q_tm1.shape}',
        include_default_message=True)
    chex.assert_shape(
        dist_q_target_t, (
            batch_size,
            None,
            self.num_atoms,
        ),
        custom_message=f'Expected (batch_size, num_actions, num_atoms), got: {dist_q_target_t.shape}',
        include_default_message=True)
    # Swap distribution and action dimension, since
    # rlax.quantile_q_learning expects it that way.
    dist_q_tm1 = jnp.swapaxes(dist_q_tm1, 1, 2)
    dist_q_target_t = jnp.swapaxes(dist_q_target_t, 1, 2)
    quantiles = (
        (jnp.arange(self.num_atoms, dtype=jnp.float32) + 0.5) / self.num_atoms)
    batch_quantile_q_learning = jax.vmap(
        rlax.quantile_q_learning, in_axes=(0, None, 0, 0, 0, 0, 0, None))
    losses = batch_quantile_q_learning(
        dist_q_tm1,
        quantiles,
        transitions.action,
        transitions.reward,
        transitions.discount,
        dist_q_target_t,  # No double Q-learning here.
        dist_q_target_t,
        self.huber_param,
    )
    loss = jnp.mean(losses)
    chex.assert_shape(losses, (batch_size,))
    extra = learning_lib.LossExtra(metrics={'mean_loss': loss})
    return loss, extra

@dataclasses.dataclass
class PrioritizedCategoricalDoubleQLearning(learning_lib.LossFn):
  """Categorical double q learning with prioritization on TD error."""
  discount: float = 0.99
  importance_sampling_exponent: float = 0.2
  max_abs_reward: float = 1.

  def __call__(
      self,
      network: networks_lib.TypedFeedForwardNetwork,
      params: networks_lib.Params,
      target_params: networks_lib.Params,
      batch: reverb.ReplaySample,
      key: networks_lib.PRNGKey,
  ) -> Tuple[jax.Array, learning_lib.LossExtra]:
    """Calculate a loss on a single batch of data."""
    transitions: types.Transition = batch.data
    probs = batch.info.probability

    # Forward pass.
    key1, key2, key3 = jax.random.split(key, 3)
    _, logits_tm1, atoms_tm1 = network.apply(
        params, transitions.observation, is_training=True, key=key1)
    _, logits_t, atoms_t = network.apply(
        target_params, transitions.next_observation, is_training=True, key=key2)
    q_t_selector, _, _ = network.apply(
        params, transitions.next_observation, is_training=True, key=key3)

    # Cast and clip rewards.
    d_t = (transitions.discount * self.discount).astype(jnp.float32)
    r_t = jnp.clip(transitions.reward, -self.max_abs_reward,
                   self.max_abs_reward).astype(jnp.float32)

    # Compute categorical double Q-learning loss.
    batch_loss_fn = jax.vmap(
        rlax.categorical_double_q_learning,
        in_axes=(None, 0, 0, 0, 0, None, 0, 0))
    batch_loss = batch_loss_fn(atoms_tm1, logits_tm1, transitions.action, r_t,
                               d_t, atoms_t, logits_t, q_t_selector)

    # Importance weighting.
    importance_weights = (1. / probs).astype(jnp.float32)
    importance_weights **= self.importance_sampling_exponent
    importance_weights /= jnp.max(importance_weights)

    # Reweight.
    loss = jnp.mean(importance_weights * batch_loss)  # []
    extra = learning_lib.LossExtra(
        metrics={}, reverb_priorities=jnp.abs(batch_loss).astype(jnp.float64))
    return loss, extra

@dataclasses.dataclass
class QLearning(learning_lib.LossFn):
  """Deep q learning.

  This matches the original DQN loss: https://arxiv.org/abs/1312.5602.
  It differs by two aspects that improve it on the optimization side
    - it uses Adam instead of RMSProp as an optimizer
    - it uses a square loss instead of the Huber one.
  """
  discount: float = 0.99
  max_abs_reward: float = 1.

  def __call__(
      self,
      network: networks_lib.TypedFeedForwardNetwork,
      params: networks_lib.Params,
      target_params: networks_lib.Params,
      batch: reverb.ReplaySample,
      key: networks_lib.PRNGKey,
  ) -> Tuple[jax.Array, learning_lib.LossExtra]:
    """Calculate a loss on a single batch of data."""
    transitions: types.Transition = batch.data

    # Forward pass.
    key1, key2 = jax.random.split(key)
    q_tm1 = network.apply(
        params, transitions.observation, is_training=True, key=key1)
    q_t = network.apply(
        target_params, transitions.next_observation, is_training=True, key=key2)

    # Cast and clip rewards.
    d_t = (transitions.discount * self.discount).astype(jnp.float32)
    r_t = jnp.clip(transitions.reward, -self.max_abs_reward,
                   self.max_abs_reward).astype(jnp.float32)

    # Compute Q-learning TD-error.
    batch_error = jax.vmap(rlax.q_learning)
    td_error = batch_error(q_tm1, transitions.action, r_t, d_t, q_t)
    batch_loss = jnp.square(td_error)

    loss = jnp.mean(batch_loss)
    extra = learning_lib.LossExtra(metrics={})
    return loss, extra

@dataclasses.dataclass
class RegularizedQLearning(learning_lib.LossFn):
  """Regularized Q-learning.

  Implements DQNReg loss function: https://arxiv.org/abs/2101.03958.
  This is almost identical to QLearning except: 1) Adds a regularization term;
  2) Uses vanilla TD error without huber loss. 3) No reward clipping.
  """
  discount: float = 0.99
  regularizer_coeff = 0.1

  def __call__(
      self,
      network: networks_lib.TypedFeedForwardNetwork,
      params: networks_lib.Params,
      target_params: networks_lib.Params,
      batch: reverb.ReplaySample,
      key: networks_lib.PRNGKey,
  ) -> Tuple[jax.Array, learning_lib.LossExtra]:
    """Calculate a loss on a single batch of data."""
    transitions: types.Transition = batch.data

    # Forward pass.
    key1, key2 = jax.random.split(key)
    q_tm1 = network.apply(
        params, transitions.observation, is_training=True, key=key1)
    q_t = network.apply(
        target_params, transitions.next_observation, is_training=True, key=key2)

    d_t = (transitions.discount * self.discount).astype(jnp.float32)

    # Compute Q-learning TD-error.
    batch_error = jax.vmap(rlax.q_learning)
    td_error = batch_error(
        q_tm1, transitions.action, transitions.reward, d_t, q_t)
    td_error = 0.5 * jnp.square(td_error)

    def select(qtm1, action):
      return qtm1[action]
    q_regularizer = jax.vmap(select)(q_tm1, transitions.action)

    loss = self.regularizer_coeff * jnp.mean(q_regularizer) + jnp.mean(td_error)
    extra = learning_lib.LossExtra(metrics={})
    return loss, extra

@dataclasses.dataclass
class MunchausenQLearning(learning_lib.LossFn):
  """Munchausen q learning.

  Implements M-DQN: https://arxiv.org/abs/2007.14430.
  """
  entropy_temperature: float = 0.03  # tau parameter
  munchausen_coefficient: float = 0.9  # alpha parameter
  clip_value_min: float = -1e3
  discount: float = 0.99
  max_abs_reward: float = 1.
  huber_loss_parameter: float = 1.

  def __call__(
      self,
      network: networks_lib.TypedFeedForwardNetwork,
      params: networks_lib.Params,
      target_params: networks_lib.Params,
      batch: reverb.ReplaySample,
      key: networks_lib.PRNGKey,
  ) -> Tuple[jax.Array, learning_lib.LossExtra]:
    """Calculate a loss on a single batch of data."""
    transitions: types.Transition = batch.data

    # Forward pass.
    key1, key2, key3 = jax.random.split(key, 3)
    q_online_s = network.apply(
        params, transitions.observation, is_training=True, key=key1)
    action_one_hot = jax.nn.one_hot(transitions.action, q_online_s.shape[-1])
    q_online_sa = jnp.sum(action_one_hot * q_online_s, axis=-1)
    q_target_s = network.apply(
        target_params, transitions.observation, is_training=True, key=key2)
    q_target_next = network.apply(
        target_params, transitions.next_observation, is_training=True, key=key3)

    # Cast and clip rewards.
    d_t = (transitions.discount * self.discount).astype(jnp.float32)
    r_t = jnp.clip(transitions.reward, -self.max_abs_reward,
                   self.max_abs_reward).astype(jnp.float32)

    # Munchausen term : tau * log_pi(a|s)
    munchausen_term = self.entropy_temperature * jax.nn.log_softmax(
        q_target_s / self.entropy_temperature, axis=-1)
    munchausen_term_a = jnp.sum(action_one_hot * munchausen_term, axis=-1)
    munchausen_term_a = jnp.clip(munchausen_term_a,
                                 a_min=self.clip_value_min,
                                 a_max=0.)

    # Soft Bellman operator applied to q
    next_v = self.entropy_temperature * jax.nn.logsumexp(
        q_target_next / self.entropy_temperature, axis=-1)
    target_q = jax.lax.stop_gradient(r_t + self.munchausen_coefficient *
                                     munchausen_term_a + d_t * next_v)

    batch_loss = rlax.huber_loss(target_q - q_online_sa,
                                 self.huber_loss_parameter)
    loss = jnp.mean(batch_loss)

    extra = learning_lib.LossExtra(metrics={})
    return loss, extra


from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.common.touch_actions import TouchActions
from lxml import html
import os
import json
import numpy as np

# Functions
def find_session_id(driver):
	# Find session_id
	outF = open('session_id.txt', 'w')
	url = driver.command_executor._url
	session_id = driver.session_id
	outF.write(session_id)
	outF.write('\n')
	outF.write(url)
	outF.close()

	print("url: ", url)
	print("Session Id: ", session_id)

def main():
  driver = webdriver.Chrome('./chromedriver')
  actions = ActionChains(driver)
  find_session_id(driver)

  driver.get('https://www.crunchbase.com/')
  elem = driver.find_element_by_id('mat-input-0')
  elem.send_keys('hello')
  #mat-input-0

  input('PRESS ENTER WHEN YOU ARE FINISHED')

main()


# encoding=utf-8
import requests
import time

class MethodRequest:

    def __init__(self):
        pass

    def request_value(self, method, url, data, headers):
        headers.update({'Connection': 'close'})
        print('请求方法: ', method, url, data, headers, type(url))
        requests.adapters.DEFAULT_RETRIES = 51
        requests.session().keep_alive = False
        try:
            if method.upper() == 'GET':
                if 'https' in url:
                    print('True')
                    result = requests.get(url, headers=headers, verify=False).text
                else:
                    result = requests.get(url, headers=headers).text
            elif method.upper() == 'POST':
                if 'https' in url:
                    result = requests.post(url, data=data, headers=headers, verify=False).text
                else:
                    result = requests.post(url, data=data, headers=headers).text
            elif method.upper() == 'PUT':
                if 'https' in url:
                    result = requests.put(url, data=data, headers=headers, verify=False).text
                else:
                    result = requests.put(url, data=data, headers=headers).text
            elif method.upper() == 'DELETE':
                if 'https' in url:
                    result = requests.delete(url, data=data, headers=headers, verify=False).text
                else:
                    result = requests.delete(url, data=data, headers=headers).text
            else:
                result = "请求方法不正确"
        except Exception as e:
            print(e)
            time.sleep(0.5)
            result = "解析请求结果失败 : %s" %e

        return result