# vim: set expandtab:

import typing
import subprocess
import os
import pwd
import sys
import docker
import re
import socket
import psutil
import io
import pickle
import math
import threading
import time
import shutil
import uuid

from .imageTransient import TransientImageSlurmBackend, list_instances, get_gce_client
from ..utils import get_default_gcp_zone, get_default_gcp_project, gcp_hourly_cost, isatty, canine_logging

from requests.exceptions import ConnectionError as RConnectionError, ReadTimeout
from urllib3.exceptions import ProtocolError

import pandas as pd

import threading

gce_lock = threading.Lock()

class DockerTransientImageSlurmBackend(TransientImageSlurmBackend): # {{{
    def __init__(
        self, cluster_name, *,
        action_on_stop = "delete",
        image_family = "slurm-gcp-docker-v1",
        image_project = "broad-getzlab-workflows",
        image = None,
        storage_namespace = "workspace", storage_bucket = None, storage_disk = None, storage_disk_size = "100",
        clust_frac = 1.0, user = os.environ["USER"], shutdown_on_exit = False, **kwargs
    ):
        if user is None:
            # IE: USER was not set
            raise ValueError("USER not set in environment. Must explicitly pass user argument")

        if storage_bucket is not None and storage_disk is not None:
            canine_logging.warning("You specified both a persistent disk and cloud bucket to store workflow outputs; will only store to bucket!")

        if "image" not in kwargs:
            kwargs["image"] = image

        super().__init__(**{**kwargs, **{ "slurm_conf_path" : "" }})

        self.config = {
          "cluster_name" : cluster_name,
          "worker_prefix" : socket.gethostname(),
          "action_on_stop" : action_on_stop,
          "image_family" : image_family,
          "image_project" : image_project,
          "clust_frac" : 1.0,
          "user" : user,
          "storage_namespace" : storage_namespace,
          "storage_bucket" : storage_bucket,
          "storage_disk" : storage_disk,
          "storage_disk_size" : storage_disk_size,
          "storage_uuid" : str(uuid.uuid4().hex[0:4]),
          **{ k : v for k, v in self.config.items() if k not in { "worker_prefix", "user", "action_on_stop" } }
        }
        self.config["image"] = self.get_latest_image(
          image_family = self.config["image_family"],
          project = self.config["image_project"],
        )["name"] if image is None else image

        # placeholder for Docker API
        self.dkr = None

        # placeholder for Docker container object
        self.container = None

        # flag to indicate whether the Docker was already running
        self.preexisting_container = False

        # flag to indicate whether we shutdown the container once the backend is exited
        self.shutdown_on_exit = shutdown_on_exit

        # placeholder for node list (loaded from lookup table)
        self.nodes = pd.DataFrame()

    def init_slurm(self):
        # query /etc/passwd for UID/GID information if we are running as a different user
        # FIXME: how should this work for OS Login/LDAP/etc.?
        uid = None; gid = None
        if self.config["user"] != os.environ["USER"]:
            uinfo = pwd.getpwnam(self.config["user"])
            uid = uinfo.pw_uid; gid = uinfo.pw_gid
        else:
            uid = os.getuid(); gid = os.getgid()

        # initialize Docker API
        self.dkr = docker.from_env()

        #
        # check if image exists
        try:
            image = self.dkr.images.get(f'gcr.io/{self.config["image_project"]}/slurm_gcp_docker:latest')
        except docker.errors.ImageNotFound:
            raise Exception("You have not yet built or pulled the Slurm Docker image!")
        except RConnectionError as e:
            if isinstance(e.args[0], ProtocolError):
                if isinstance(e.args[0].args[1], PermissionError):
                    raise PermissionError("You do not have permission to run Docker!")
                elif isinstance(e.args[0].args[1], ConnectionRefusedError):
                    raise ConnectionRefusedError("The Docker daemon does not appear to be running on this machine. Please start it.")
            raise Exception("Unknown problem connecting to the Docker daemon")
        except Exception as e:
            raise Exception("Problem starting Slurm Docker: {}: {}".format(
              type(e).__name__, e
            ))

        #
        # create NFS/rclone bucket mountpoints
        self.create_filesystem(uid, gid)

        # copy credential files to NFS
        self.copy_cloud_credentials()

        #
        # ensure that Docker can start (no Slurm processes outside of Docker already running)
        try:
            ready_for_docker()
        except:
            canine_logging.error("Docker host is not ready to start container!")
            raise

        #
        # create the Slurm container if it's not already present
        canine_logging.info1("Starting Slurm controller ...")
        if self.config["cluster_name"] not in [x.name for x in self.dkr.containers.list()]:
            self.dkr.containers.run(
              image = image.tags[0], detach = True, network_mode = "host",
              mounts = [
                docker.types.Mount(
                  target = "/mnt", source = "/mnt", type = "bind", propagation = "rshared"
                ),
                docker.types.Mount(
                  target = "/dev", source = "/dev", type = "bind", propagation = "rshared"
                )
              ],
              name = self.config["cluster_name"], command = "/bin/bash",
              stdin_open = True, remove = True, privileged = True,
              environment = { "HOST_USER" : self.config["user"], "HOST_UID" : uid, "HOST_GID" : gid }
            )
            self.container = self._get_container(self.config["cluster_name"])

        # otherwise, try and start it if it's stopped
        else:
            self.preexisting_container = True
            self.container = self._get_container(self.config["cluster_name"])
            if self.container().status == "exited":
                self.container().start()

        # TODO: should we restart slurmctld in the container here?

        #
        # wait until the container is fully started, or error out if it failed
        # to start
        self.wait_for_container_to_be_ready(timeout = 60)

        #
        # initialize storage
        self.init_storage()

        #
        # save the configuration to disk so that Slurm knows how to configure
        # the nodes it creates
        subprocess.check_call("""
          [ ! -d /mnt/nfs/clust_conf/canine ] && mkdir -p /mnt/nfs/clust_conf/canine ||
            echo -n
          """, shell = True, executable = '/bin/bash')
        with open("/mnt/nfs/clust_conf/canine/backend_conf.pickle", "wb") as f:
            pickle.dump(self.config, f)

    def init_nodes(self):
        self.wait_for_cluster_ready(elastic = True, timeout=60)

        # list all the nodes that Slurm is aware of; this may be useful subsequently?
        #allnodes = pd.read_pickle("/mnt/nfs/clust_conf/slurm/host_LuT.pickle")

    def stop(self):
        # remove any bucket mount commands created by this instance
        if os.path.exists(f"/mnt/nfs/.rclone_mounts_{self.config['storage_uuid']}.sh"):
            os.remove(f"/mnt/nfs/.rclone_mounts_{self.config['storage_uuid']}.sh")

        # if the Docker was not spun up by this context manager, do not tear
        # anything down -- we don't want to clobber an already running cluster
        if self.shutdown_on_exit and not self.preexisting_container:
            # delete node configuration file
            try:
                subprocess.check_call(
                  "rm -f /mnt/nfs/clust_conf/canine/backend_conf.pickle",
                  shell = True,
                  timeout = 10
                )
            except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
                canine_logging.error("Couldn't delete node configuration file:")
                canine_logging.error(e)

            #
            # shutdown nodes that are still running (except NFS)
            allnodes = self.nodes

            # if we're aborting before the NFS has even been started, there are no
            # nodes to shutdown.
            if not allnodes.empty:
                # sometimes the Google API will spectacularly fail; in that case, we
                # just try to shutdown everything in the node list, regardless of whether
                # it exists.
                try:
                    extant_nodes = self.list_instances_all_zones()
                    self.nodes = allnodes.loc[allnodes.index.isin(extant_nodes["name"]) &
                                   (allnodes["machine_type"] != "nfs")]
                except:
                    self.nodes = allnodes.loc[allnodes["machine_type"] != "nfs"]

                # superclass method will stop/delete/leave these running, depending on how
                # self.config["action_on_stop"] is set
                super().stop()

            #
            # stop the Docker

            # this needs to happen after super().stop() is invoked, since that
            # calls scancel, which in turn requires a running Slurm controller Docker
            if self.container is not None:
                self.container().stop()

    def _get_container(self, container_name):
        def closure():
            backoff_factor = 1
            while True:
                try:
                    container = self.dkr.containers.get(container_name)
                    break
                except ReadTimeout:
                    canine_logging.warning(f"Request to controller Docker timed out; retrying in {int(10*backoff_factor)} seconds ...")
                    time.sleep(10*backoff_factor)
                    backoff_factor *= 1.1
            return container

        return closure

    def reenter(self):
        # docker object cannot be reused between processes due to socket connections
        self.dkr = docker.from_env()

    def create_filesystem(self, uid, gid):
        ## create main NFS mountpoint
        if not os.path.exists("/mnt/nfs"):
            try:
                subprocess.check_call("sudo mkdir /mnt/nfs", shell=True)
                subprocess.check_call(f"sudo chown {uid}:{gid} /mnt/nfs", shell=True)
            except:
                # TODO: be more specific about exception catching
                canine_logging.error("Could not create NFS mountpoint; see stack trace for details")
                raise

        ## create root mountpoint for rclone bucket filesystems (will be bind mounted
        #  into main mountpoint; for some reason, these cannot be nested in /mnt/nfs)
        if not os.path.exists("/mnt/rclone"):
            try:
                subprocess.check_call("sudo mkdir /mnt/rclone", shell=True)
                subprocess.check_call(f"sudo chown {uid}:{gid} /mnt/rclone", shell=True)
            except:
                # TODO: be more specific about exception catching
                canine_logging.error("Could not create rclone; see stack trace for details")
                raise

    def init_storage(self):
        ## make /mnt/nfs NFS its own virtual filesystem
        # this is so that Canine's system for detecting whether files to be
        # localized won't symlink things that reside outside /mnt/nfs but on the same
        # actual filesystem as /mnt/nfs
        subprocess.check_call("""[ $(df -P /mnt/nfs/ | awk 'NR > 1 { print $6 }') == '/mnt/nfs' ] || \
          sudo mount --bind /mnt/nfs /mnt/nfs""", shell=True, executable="/bin/bash")

        ## mount bucket via rclone (unstable!)
        if self.config["storage_bucket"] is not None:
            canine_logging.info1(f"Saving workflow results to bucket {self.config['storage_bucket']} mounted at /mnt/nfs/{self.config['storage_namespace']} ...")

            # TODO: check if bucket exists; create it if not

            # generate cache disk if it doesn't exist
            rc, stdout, stderr = self.invoke(
              "gcloud_make_rwdisk {disk_name} {disk_size} {mount_prefix} false {node_name} {node_zone}".format(
                disk_name = f'cache-disk-{self.config["worker_prefix"]}',
                disk_size = 100,
                mount_prefix = "/tmp/rclone_cache",
                node_name = self.config["worker_prefix"],
                node_zone = get_default_gcp_zone()
              )
            )
            if rc != 0:
                canine_logging.error(f"Could not generate bucket mountpoint; see error log for details:")
                canine_logging.error(stderr.read().decode())
                raise RuntimeError()

            # invoke rclone inside docker to create bucket-backed FUSE filesystem
            # this is complicated, because NFS cannot export a nested FUSE filesystem.
            # we thus mount it outside of the NFS to /mnt/rclone/<bucket> (which will be exported to workers),
            # and bind mount /mnt/rclone/<bucket> to /mnt/nfs/<bucket> (in order to access it on the controller)
            # workers will NFS mount /mnt/rclone/<bucket> to /mnt/nfs/<bucket> for consistent paths.
            rc, stdout, stderr = self.invoke(
              """bash -c \
                'set -e; export GOOGLE_APPLICATION_CREDENTIALS=$CLOUDSDK_CONFIG/application_default_credentials.json; \
                 [ ! -d {mountpoint} ] && mkdir {mountpoint}; \
                 [ ! -d {bind_mountpoint} ] && mkdir {bind_mountpoint}; \
                 df -t fuse.rclone {mountpoint} || \
                  rclone mount gcs:{bucket_name} {mountpoint} --daemon --links \
                   --uid $HOST_UID --gid $HOST_GID --file-perms 0755 --allow-other \
                   --vfs-cache-mode full --cache-dir /tmp/rclone_cache/ --vfs-cache-max-size 95Gi \
                   --vfs-write-back 10m --vfs-cache-poll-interval 10m --vfs-fast-fingerprint \
                   --vfs-cache-max-age 48h --dir-cache-time 2h --poll-interval 10m --no-modtime \
                   --config /sgcpd/conf/rclone.conf; \
                 df -t fuse.rclone {bind_mountpoint} || \
                  mount --bind {mountpoint} {bind_mountpoint}'""".format(
                bucket_name = self.config["storage_bucket"][5:] if self.config["storage_bucket"].startswith("gs://") else self.config["storage_bucket"],
                mountpoint = f'/mnt/rclone/{self.config["storage_namespace"]}',
                bind_mountpoint = f'/mnt/nfs/{self.config["storage_namespace"]}'
              ),
              user = "root"
            )
            if rc != 0:
                canine_logging.error(f"Could not mount bucket; see error log for details:")
                canine_logging.error(stderr.read().decode())
                raise RuntimeError()

            # export mount
            subprocess.check_call(f"sudo exportfs -o fsid=1,rw,async,no_subtree_check,insecure,no_root_squash *.internal:/mnt/rclone/{self.config['storage_namespace']}", shell=True)

            # save rclone mountpoint list (worker nodes will subsequently read this in to know what directories to mount after starting up)
            # this file will be erased when backend is torn down
            # TODO: use flock to remove file if backend crashes; when starting backend, check for unlocked files and remove them
            # will need try/except on worker nodes to avoid them freezing up if they inadvertently attempt to mount non-exported buckets
            with open(f"/mnt/nfs/.rclone_mounts_{self.config['storage_uuid']}.sh", "w") as f:
                f.write("if ! mountpoint -q /mnt/nfs/{mount_dir}; then sudo timeout -k 30 30 mount -o defaults,hard,intr ${{CONTROLLER_NAME}}:/mnt/rclone/{mount_dir} /mnt/nfs/{mount_dir} || echo 'Could not mount rclone mountpoint /mnt/rclone/{mount_dir}'; fi".format(mount_dir = self.config['storage_namespace']))

        ## create disk
        # note that bucket takes priority if both are specified
        elif self.config["storage_disk"] is not None:
            # GCP disks must match regex '^[a-z]([a-z0-9-]*[a-z0-9])?'; raise error if not
            if re.match('^[a-z]([a-z0-9-]*[a-z0-9])?$', self.config["storage_disk"]) is None:
                raise ValueError(f"\"{self.config['storage_disk']}\" is an invalid storage namespace. Storage namespaces can only contain lowercase letters, numbers, and dashes, must start with a letter, and must end with a letter/number.")

            canine_logging.info1(f"Saving workflow results to persistent disk {self.config['storage_disk']} ({self.config['storage_disk_size']}GB) mounted at /mnt/nfs/{self.config['storage_namespace']} ...")
            # use procedure to create RW disk from localization, inside docker
            rc, stdout, stderr = self.invoke(
              "gcloud_make_rwdisk {disk_name} {disk_size} {mount_prefix} false {node_name} {node_zone}".format(
                disk_name = self.config["storage_disk"],
                disk_size = f"{self.config['storage_disk_size']}",
                mount_prefix = f"/mnt/nfs/{self.config['storage_namespace']}",
                node_name = self.config["worker_prefix"],
                node_zone = get_default_gcp_zone()
              )
            )
            if rc != 0:
                canine_logging.error("Error attaching workflow results disk; see error log for details:")
                canine_logging.error(stderr.read().decode())
                raise RuntimeError()

        ## Check disk usage and warn user if it is small
        free_space_gb = int(shutil.disk_usage(f"/mnt/nfs/{self.config['storage_namespace']}").free/(1024**3))
        if free_space_gb < 300:
            canine_logging.warning(
                f"Workflow results disk low on space ({free_space_gb} GB remaining)"
            )

        # TODO: add warnings if overall disk is small (bad disk IO) or node core
        #       count is low (bad network IO)

        ## export NFS
        subprocess.check_call("sudo exportfs -o fsid=0,rw,async,no_subtree_check,insecure,no_root_squash,crossmnt *.internal:/mnt/nfs", shell=True)

    def copy_cloud_credentials(self):
        ## gcloud
        # TODO: check that we are properly authenticated
        # TODO: check $CLOUDSDK_CONFIG environment variable
        gcloud_conf_dir = subprocess.check_output("echo -n ~/.config/gcloud", shell = True).decode()
        if os.path.isdir(gcloud_conf_dir):
            if not os.path.isdir("/mnt/nfs/credentials/gcloud"):
                os.makedirs("/mnt/nfs/credentials/gcloud")
            subprocess.run(f'cp -rf $(find {gcloud_conf_dir} -mindepth 1 -maxdepth 1 ! -name "logs") /mnt/nfs/credentials/gcloud', shell = True)

    def get_latest_image(self, image_family = None, project = None):
        image_family = self.config["image_family"] if image_family is None else image_family
        project = self.config["project"] if project is None else project
        with gce_lock:
            ans = get_gce_client().images().getFromFamily(family = image_family, project = project).execute()
        return ans

    def invoke(self, command, interactive = False, bypass_docker = False, user = None):
        """
        Set bypass_docker to True to execute the command directly on the host,
        rather than in the controller container. Useful for debugging.
        """
        if not isatty(sys.stdout, sys.stdin):
            interactive = False

        user = self.config["user"] if user is None else user

        # re-purpose LocalSlurmBackend's invoke
        local_invoke = super(TransientImageSlurmBackend, self).invoke
        if self.container is not None and self.container().status == "running":
            if not bypass_docker:
                cmd = "docker exec --user {user} {ti_flag} {container} {command}".format(
                  user = user,
                  ti_flag = "-ti" if interactive else "",
                  container = self.config["cluster_name"],
                  command = command
                )
                # if command fails for a recoverable reason, retry up to 7 times
                # with exponential backoff (max ~2 minute wait)
                timeout = 8
                tries = 0
                while tries < 7:
                    ret, stdout, stderr = local_invoke(cmd, interactive)
                    stderr_str = stderr.read().decode().rstrip()
                    stderr.seek(0)

                    # no stderr; assume command ran successfully
                    if stderr_str == "":
                        break

                    # stderr corresponds to a known Docker failure mode than can
                    # be recovered from
                    if any([reason in stderr_str for reason in [
                      "Error response from daemon: No such exec instance",
                      "OCI runtime exec failed: exec failed",
                      "Unable to contact slurm controller (connect failure)",
                      "Socket timed out on send/recv operation"
                    ]]):
                        canine_logging.warning(
                          'Command {cmd} failed with known recoverable error reason "{err}"; retrying in {timeout} seconds up to {tries} more times'.format(
                            cmd = command,
                            err = stderr_str,
                            timeout = timeout,
                            tries = 7 - tries
                          )
                        )
                        time.sleep(timeout)
                        timeout *= 2
                        tries += 1

                    # warn the user that the command had stuff written to stderr,
                    # since this may indicate something is wrong
                    else:
                        canine_logging.debug(
                          'Command {cmd} returned stderr "{err}"'.format(
                            cmd = command,
                            err = stderr_str
                          )
                        )
                        break
            else:
                cmd = command
                ret, stdout, stderr = local_invoke(cmd, interactive)
            return (ret, stdout, stderr)
        else:
            return (1, io.BytesIO(), io.BytesIO(b"Container is not running!"))

    def wait_for_container_to_be_ready(self, timeout = 3000):
        canine_logging.info1("Waiting up to {} seconds for Slurm controller to start ...".format(timeout))
        (rc, _, _) = self.invoke(
          "timeout {} bash -c 'while [ ! -f /.started ]; do sleep 1; done'".format(timeout),
          interactive = True,
          user = "root"
        )
        if rc == 124:
            raise TimeoutError("Slurm controller did not start within {} seconds!".format(timeout))
        canine_logging.info1("Started Slurm controller.")

# }}}

class LocalDockerSlurmBackend(DockerTransientImageSlurmBackend):
    def __enter__(self):
        self.dkr = docker.from_env()
        self.container = self._get_container(self.config["cluster_name"])
        return self
    def __exit__(self, *args):
        pass

# Python version of checks in docker_run.sh
def ready_for_docker():
    #
    # check if Slurm/mysql/Munge are already running
    already_running = [["slurmctld", "A Slurm controller"],
                       ["slurmdbd", "The Slurm database daemon"],
                       ["mysqld", "mysql"],
                       ["munged", "Munge"]]

    all_procs = { x.name() : x.pid for x in psutil.process_iter() }

    for proc, desc in already_running:
        # is the process is running at all?
        if proc in all_procs:
            # iterate through parents to see if it was launched in a Docker (containerd)
            for parent_proc in psutil.Process(all_procs[proc]).parents():
                if parent_proc.name().startswith("containerd"):
                    break
                if parent_proc.pid == 1:
                    raise Exception("{desc} is already running on this machine (outside of a Docker container). Please run `[sudo] killall {proc}' and try again.".format(desc = desc, proc = proc))


#!/usr/bin/env python
# coding: utf-8

# # COURSE: Master statistics and machine learning: Intuition, Math, code
# ##### COURSE URL: udemy.com/course/statsml_x/?couponCode=202006
# ## SECTION: The t-test family
# ### VIDEO: Permutation testing
# #### TEACHER: Mike X Cohen, sincxpress.com

# In[ ]:

# import libraries
import matplotlib.pyplot as plt
import numpy as np
import scipy.stats as stats

# In[ ]:

## simulate two distributions

# number of trials
N = 100

# dataset "A"
r = np.random.randn(N)
r[r>0] = np.log(1+r[r>0])
dataA = 26-r*10

# get histogram values for later comparison
yA,xA = np.histogram(dataA,20)
xA = (xA[:-1]+xA[1:])/2

# dataset "B"
r = np.random.randn(N)
r[r>0] = np.log(1+r[r>0])
dataB = 30-r*10

#get histogram values for later comparison
yB,xB = np.histogram(dataB,20)
xB = (xB[:-1]+xB[1:])/2

plt.stem(xA,yA,'b',markerfmt='bo',basefmt=' ',label='Data"A"')
plt.stem(xB,yB,'r',markerfmt='ro',basefmt=' ',label='Data"B"')
plt.legend()
plt.show()

# In[ ]:

## mix trials together

# concatenate trials
alldata = np.hstack((dataA,dataB))

# condition labels
conds = np.hstack((np.ones(N),2*np.ones(N)))

# In[ ]:

## generate one null hypothesis scenario

# random permutation
fakeconds = np.random.permutation(N*2)

# shuffled condition labels
fakeconds[fakeconds<N] = 1
fakeconds[fakeconds>1] = 2

# these two means should be different.
print([np.mean(alldata[conds==1]), np.mean(alldata[conds==2])])

# should these two be different?
print([np.mean(alldata[fakeconds==1]), np.mean(alldata[fakeconds==2])])

# In[ ]:

## and now a distribution of null hypothesis values

nPerms = 1000
permdiffs = np.zeros(nPerms)

for permi in range(nPerms):
    fconds = np.random.permutation(N*2)
    fconds[fconds<N] = 1
    fconds[fconds>1] = 2
    permdiffs[permi] = np.mean(alldata[fconds==2]) - np.mean(alldata[fconds==1])

# plot the distribution of H0 values
plt.hist(permdiffs,50)

# and plot the observed value on top
obsval = np.mean(alldata[conds==2]) - np.mean(alldata[conds==1])
plt.plot([obsval, obsval],[0, 50],'m',linewidth=10)
plt.xlabel('Value')
plt.ylabel('Count')
plt.show()

# In[ ]:

## two methods of evaluating statistical significance

# Z-value
zVal = ( obsval-np.mean(permdiffs) ) / np.std(permdiffs,ddof=1)
p = 1-stats.norm.cdf(abs(zVal))

# p-value count
pCount = sum(permdiffs>obsval)/nPerms

print(p,pCount)


import math
from torch import nn
import torch.nn.init as init
from .common import AdaptiveFM
from .dgr import DGR
import torch
import torch.nn.functional as F

def make_model(args, parent=False):
    return ESPCN(args)

def get_valid_padding(kernel_size, dilation):
    kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1)
    padding = (kernel_size - 1) // 2
    return padding

def espcn_conv(input,prefix, weights_dict):
    x = F.conv2d(input, weights_dict[prefix + '.weight'], weights_dict[prefix + '.bias'], stride=1, padding=1)
    return x

class AdaptiveFM(nn.Module):
    # hello ckx
    def __init__(self, in_channel, kernel_size):

        super(AdaptiveFM, self).__init__()
        padding = get_valid_padding(kernel_size, 1)
        self.transformer = nn.Conv2d(in_channel, in_channel, kernel_size,
                                     padding=padding, groups=in_channel//2)
        self.gamma = nn.Parameter(torch.zeros(1), requires_grad=True)
    def forward(self, x):
        return self.transformer(x) * self.gamma + x

class ESPCN(nn.Module):
    def __init__(self, args):
        super(ESPCN, self).__init__()
        # self.act_func = nn.LeakyReLU(negative_slope=0.2)
        self.act_func = nn.ReLU(inplace=True)
        self.scale = int(args.scale[0])  # use scale[0]
        self.n_colors = args.n_colors
        self.args = args
        if self.args.new_espcn:
            self.conv1 = nn.Conv2d(self.n_colors, 64, (3, 3), (1, 1), (1, 1))
            self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
            self.conv3 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
            self.conv4 = nn.Conv2d(64, 3 * (self.scale ** 2), (3, 3), (1, 1), (1, 1))
        else:
            self.conv1 = nn.Conv2d(self.n_colors, 64, (3, 3), (1, 1), (1, 1))
            self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
            self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
            self.conv4 = nn.Conv2d(32, 3 * (self.scale ** 2), (3, 3), (1, 1), (1, 1))
        self.pixel_shuffle = nn.PixelShuffle(self.scale)

        self._initialize_weights()

    def forward(self, x, num, weights_dict=None):
        if weights_dict!=None:
            #print(3333333333333333333333)
            out = self.act_func(espcn_conv(x, 'model.conv1', weights_dict))
            out = self.act_func(espcn_conv(out, 'model.conv2', weights_dict))
            out = self.act_func(espcn_conv(out, 'model.conv3', weights_dict))
            out = self.pixel_shuffle(espcn_conv(out, 'model.conv4', weights_dict))
            return out
        else:
            #print(1111111111111111111111)
            if self.args.new_espcn:
                res = self.act_func(self.conv1(x))
                out = self.act_func(self.conv2(res))
                out = self.conv3(out)
                out = res+out
                out = self.pixel_shuffle(self.conv4(out))
                return out
            else:
                out = self.act_func(self.conv1(x))
                out = self.act_func(self.conv2(out))
                out = self.act_func(self.conv3(out))
                out = self.pixel_shuffle(self.conv4(out))
                return out

    def _initialize_weights(self):
        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv4.weight)

    def load_state_dict(self, state_dict, strict=True):
        own_state = self.state_dict()
        for name, param in state_dict.items():
            if name in own_state:
                if isinstance(param, nn.Parameter):
                    param = param.data
                own_state[name].copy_(param)


import functools
import turtle

import hangman
import words

def write_word(word):
    writer = turtle.Turtle()
    writer.penup()
    writer.goto(100, 200)
    writer.write(word, font=('Arial', 16, 'bold'))
    writer.hideturtle()

# Иницилизация
original_word = words.get_random_word()
word = '_' * len(original_word)
write_word(word)
errors = 0

def on_key_press (key):
    global errors
    global word
    if errors <= hangman.MAX_ERRORS:
        if key in original_word:
            word = words.fill_char(key, word, original_word)
            write_word(word)
            if word == original_word:
                print('You win')
        else:
            errors += 1
            hangman.build_hangman(errors)
    else:
        print('You lost')

# Handle press keys 'a' to 'z'
for i in range(ord('a'), ord('z') + 1):
    key = chr(i)
    turtle.onkey(functools.partial(on_key_press, key), key)

turtle.listen()
turtle.mainloop()


import json
import os
import numpy as np
import sys
import copy
import random
import jsonlines
import time
import scipy.stats

task = sys.argv[1]
model = sys.argv[2]
model = f"en_dense_lm_{model}"

# !!! replace by your $base_dir/ana_rlt here
base_dir = "$base_dir/ana_rlt"
ana_rlt_dir = f"{base_dir}/{model}"

debug_scale = 1
debug_n = 100
icl_k_dict = {
    'cb': 19,
    'sst2': 32,
    'sst5': 32,
    'subj': 32,
    'mr': 32,
    'agnews': 22,
}
icl_k = icl_k_dict[task]
save_rlts = {}

def load_info(ana_setting):
    rlt_dir = f"{ana_rlt_dir}/{task}/{ana_setting}"
    info = []
    # jsonlines
    to_read_num = 100000 if ana_setting == 'ft' else debug_n
    with open(f"{rlt_dir}/record_info.jsonl", "r") as f:
        for item in jsonlines.Reader(f):
            if ana_setting == 'ft' or ana_setting == 'ftzs':
                item = item['attn_q']
            elif ana_setting == 'icl':
                item = item['attn_map_ctx']
            else:
                assert 1 == 2
            info.append(item)
            to_read_num -= 1
            if to_read_num == 0:
                break
    n_records = len(info)
    print(rlt_dir, n_records)
    info = info[:n_records//debug_scale]
    return info

stt_time = time.time()

ft_attn_q = load_info('ft')  # (n_example, n_layers, ~len, q_hidden_dim)
print(f'loading ft data costs {time.time() - stt_time} seconds')
stt_time = time.time()

ftzs_attn_q = load_info('ftzs')  # (n_example, n_layers, q_hidden_dim)
print(f'loading ftzs data costs {time.time() - stt_time} seconds')
stt_time = time.time()

icl_attn_map_ctx = load_info('icl')  # (n_example, n_layers, n_head, tot_len)
print(f'loading icl data costs {time.time() - stt_time} seconds')
stt_time = time.time()

stored_slice_range = []

def slice_icl_attn_map(icl_attn_map_ctx):
    # (n_example, n_layers, n_head, tot_len)
    n_example = len(ft_attn_q)
    example_len = []
    for i in range(n_example):
        example_len.append(len(ft_attn_q[i][0]) - 1)

    slice_idx = []
    cur_pos = 1
    for i in range(n_example):
        new_range = list(range(cur_pos, cur_pos + example_len[i]))
        slice_idx.extend(new_range)
        cur_pos += example_len[i] + 1

    cur_pos = 0
    for i in range(n_example):
        new_range = list(range(cur_pos, cur_pos + example_len[i]))
        stored_slice_range.append(new_range)
        cur_pos += example_len[i]

    icl_attn_map_ctx = icl_attn_map_ctx[:, :, slice_idx]
    return icl_attn_map_ctx

icl_attn_map_ctx = np.array(icl_attn_map_ctx)  # (n_example, n_layers, tot_len)
icl_attn_map_ctx = slice_icl_attn_map(icl_attn_map_ctx)
print('sliced icl_attn_map_ctx shape:', icl_attn_map_ctx.shape)
print(f'slicing icl_attn_map_ctx costs {time.time() - stt_time} seconds')
stt_time = time.time()

def compact_ft_attn_q(ft_attn_q):
    # (n_example, n_layers, ~len, q_hidden_dim)
    n_examples = len(ft_attn_q)
    n_layers = len(ft_attn_q[0])
    compacted_ft_attn_q = []
    for i in range(n_layers):
        layer_ft_attn_q = []
        for j in range(n_examples):
            cur_len = len(ft_attn_q[j][i])
            for k in range(1, cur_len):
                q_vector = ft_attn_q[j][i][k]
                layer_ft_attn_q.append(q_vector)
        compacted_ft_attn_q.append(layer_ft_attn_q)
    return compacted_ft_attn_q

ft_attn_q = compact_ft_attn_q(ft_attn_q)
ft_attn_q = np.array(ft_attn_q)  # (n_layers, tot_len, q_hidden_dim)
print('compacted ft_attn_q shape:', ft_attn_q.shape)
print(f'compacting ft_attn_q costs {time.time() - stt_time} seconds')
stt_time = time.time()

def cal_cos_sim(v1, v2):
    num = (v1 * v2).sum(axis=-1)  # dot product
    denom = np.linalg.norm(v1, axis=-1) * np.linalg.norm(v2, axis=-1) + 1e-20  # length
    res = num / denom
    return res

def norm_hidden(hidden):
    norm = np.linalg.norm(hidden, axis=-1) + 1e-20  # length
    norm = norm[:, :, np.newaxis]
    hidden = hidden / norm
    return hidden

def np_softmax(x, axis=-1):
    x -= np.max(x, axis=axis, keepdims=True)
    x = np.exp(x) / np.sum(np.exp(x), axis=axis, keepdims=True)
    return x

def compute_training_attn(level='token'):
    expand_ftzs_attn_q = np.array(ftzs_attn_q)[:, :, np.newaxis, :]  # (n_example, n_layers, 1, q_hidden_dim)
    expand_ft_attn_q = ft_attn_q[np.newaxis, :, :, :]  # (1, n_layers, tot_len, q_hidden_dim)

    ftzs_attn_map = (expand_ft_attn_q * expand_ftzs_attn_q).sum(axis=-1)  # (n_example, n_layers, tot_len)
    del expand_ftzs_attn_q
    del expand_ft_attn_q
    icl_attn_map = icl_attn_map_ctx  # (n_example, n_layers, tot_len)

    icl_attn_map = np_softmax(icl_attn_map, axis=-1)
    ftzs_attn_map = np_softmax(ftzs_attn_map, axis=-1)

    if level == 'example':
        n_example = len(stored_slice_range)
        tmp_ftzs_attn_map = np.random.random([ftzs_attn_map.shape[0], ftzs_attn_map.shape[1], n_example])
        tmp_icl_attn_map = np.random.random([ftzs_attn_map.shape[0], ftzs_attn_map.shape[1], n_example])
        for i, slice_range in enumerate(stored_slice_range):
            ftzs_attn_map_example = ftzs_attn_map[:, :, slice_range].sum(axis=-1)  # (n_example, n_layers)
            icl_attn_map_example = icl_attn_map[:, :, slice_range].sum(axis=-1)  # (n_example, n_layers)
            tmp_ftzs_attn_map[:, :, i] = ftzs_attn_map_example
            tmp_icl_attn_map[:, :, i] = icl_attn_map_example
        ftzs_attn_map = tmp_ftzs_attn_map
        icl_attn_map = tmp_icl_attn_map

    # =================== analyze recall of attention top-k ========================
    topk_k = 5 if level == 'example' else 100
    recall_topk = [0 for i in range(icl_attn_map.shape[1])]
    random_recall_topk = [0 for i in range(icl_attn_map.shape[1])]
    total_topk = 0
    for i in range(icl_attn_map.shape[0]):
        for j in range(icl_attn_map.shape[1]):
            ftzs_topk = np.argsort(-ftzs_attn_map[i][j])[:topk_k].tolist()
            icl_topk = np.argsort(-icl_attn_map[i][j])[:topk_k].tolist()
            for l in range(topk_k):
                if icl_topk[l] in ftzs_topk:
                    recall_topk[j] += 1
                random_idx = random.randint(0, icl_attn_map.shape[2] - 1)
                if random_idx in ftzs_topk:
                    random_recall_topk[j] += 1
            total_topk += topk_k
    total_topk /= icl_attn_map.shape[1]
    recall_topk = [recall_topk[i] / total_topk for i in range(icl_attn_map.shape[1])]
    random_recall_topk = [random_recall_topk[i] / total_topk for i in range(icl_attn_map.shape[1])]
    save_rlts[f'{level}-level recall_topk'] = recall_topk
    save_rlts[f'{level}-level random_recall_topk'] = random_recall_topk
    print('=========================================')
    print(f'{level}-level recall_topk:', recall_topk, ' ||| ', np.mean(recall_topk))
    print(f'{level}-level random_recall_topk:', random_recall_topk, ' ||| ', np.mean(random_recall_topk))

    # =================== analyze kendall of attention weights ========================

    kendall = [0 for i in range(len(icl_attn_map[0]))]
    random_kendall = [0 for i in range(len(icl_attn_map[0]))]
    for i in range(icl_attn_map.shape[0]):
        for j in range(icl_attn_map.shape[1]):
            cur_icl_map = np.array(icl_attn_map[i][j])
            cur_ftzs_map = np.array(ftzs_attn_map[i][j])
            random_map = np.random.random(ftzs_attn_map[i][j].shape)

            kendall[j] += scipy.stats.kendalltau(cur_icl_map, cur_ftzs_map)[0]
            random_kendall[j] += scipy.stats.kendalltau(cur_icl_map, random_map)[0]

    kendall = [kendall[i] / icl_attn_map.shape[0] for i in range(icl_attn_map.shape[1])]
    random_kendall = [random_kendall[i] / icl_attn_map.shape[0] for i in range(icl_attn_map.shape[1])]
    save_rlts[f'{level}-level kendall'] = kendall
    save_rlts[f'{level}-level random_kendall'] = random_kendall
    print('=========================================')
    print(f'{level}-level kendall:', kendall, ' ||| ', np.mean(kendall))
    print(f'{level}-level random_kendall:', random_kendall, ' ||| ', np.mean(random_kendall))

compute_training_attn(level='example')
compute_training_attn(level='token')
print(f'computing attention costs {time.time() - stt_time} seconds')
stt_time = time.time()
with open(f'{base_dir}/rlt_json/{task}-{model}_training_attn.json', 'w') as f:
    json.dump(save_rlts, f, indent=2)
print(f'saving data costs {time.time() - stt_time} seconds')
stt_time = time.time()


import os
import click
import requests
from .utils import prepare_path, save_list_and_cache, write_to, get_result
from .help import BRANDING

prepare_path()

# @click.group()
@click.command()
@click.argument('ignore', required=False)
@click.version_option(message=BRANDING)
@click.option('-a', '--listall', help='Get all available terms', default=False, is_flag=True)
# @click.pass_context
def cli(listall, ignore):
    """ AzatAI Addignore Automation tool. Searches for a proper .gitignore file for your project and adds to the
        local .gitignore file.
        Search for operating systems, IDEs, languages and other environmental terms. Use ',' to separate multiple items.
        Usage: addignore python """
    # ensure that ctx.obj exists and is a dict (in case `cli()` is called
    # by means other than the `if` block below)
    # ctx.ensure_object(dict)
    prepare_path()
    if ignore:
        here = os.getcwd()
        file_path = os.path.join(here, '.gitignore')
        write_to(get_result(ignore), file_path, mode='file')
    elif not ignore and listall:
        click.secho(save_list_and_cache())
    else:
        with click.get_current_context() as context:
            click.echo(context.get_help())

# @cli.command()
# @click.pass_context
# def list(ctx):
#     """Prints all the available results."""
#     click.secho('Hello')

if __name__ == '__main__':
    cli()


#load matplotlib.pyplot as plt
import matplotlib.pyplot as plt

#load numpy as np
import numpy as np

#x range: [-pi, pi]
x = np.linspace(-np.pi, np.pi, 256, endpoint = True)

#y = sin(x)
y_sin = np.sin(x)
#y = cos(x)
y_cos = np.cos(x)

#Figure & subplot
fig = plt.figure(figsize = (12,8))
ax = fig.add_subplot(1,1,1)

#Graph
ax.plot(x, y_sin, 'r', linewidth = 2.0, label = 'sin')
ax.plot(x, y_cos, 'b', linewidth = 2.0, label = 'cos')

#Set limits
ax.set_xlim(-4,4)
ax.set_ylim(-1.5,1.5)

#Set ticks
ax.set_xticks(np.linspace(-np.pi, np.pi, 5, endpoint = True))
ax.set_yticks(np.linspace(-1.5, 1.5, 5, endpoint = True))

#Set spines
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')

ax.spines['bottom'].set_position(('data',0))
ax.spines['left'].set_position(('data',0))

#Set labels
ax.set_title('trigonometric function')
ax.set_xlabel('X')
ax.set_ylabel('Y')

#Set legend
ax.legend(loc = 'best')

#no margin
plt.tight_layout()
plt.show() ## Figure_1


from fastapi import FastAPI
from models import User, UserBet, VerifyRequest

app = FastAPI()

@app.post("/server_seed")
def get_hashed_server_seed(request: User):
    resp = request.get_server_seed()
    return resp

@app.post("/bet")
def bet(request: UserBet):
    resp = request.process()
    return resp

@app.post("/verify")
def verify(request: VerifyRequest):
    resp = request.process()
    return resp


import pandas as pd
import numpy as np
import datetime
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import colors
import seaborn as sns
from sklearn.preprocessing import LabelEncoder, OneHotEncoder, OrdinalEncoder
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt, numpy as np
from mpl_toolkits.mplot3d import Axes3D
from sklearn.cluster import AgglomerativeClustering
from matplotlib.colors import ListedColormap
from sklearn import metrics
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
from sklearn.svm import SVR
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error,mean_absolute_error,mean_absolute_percentage_error,r2_score
from sklearn.metrics import f1_score, accuracy_score
from sklearn.naive_bayes import GaussianNB
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.neighbors import KNeighborsClassifier
import pickle

data=pd.read_csv("data/Fraud Claims.csv")
req_col=['Month', 'WeekOfMonth', 'DayOfWeek',
       'DayOfWeekClaimed', 'MonthClaimed', 'WeekOfMonthClaimed', 'Fault',
       'FraudFound_P', 'RepNumber', 'Days_Policy_Accident',
       'Days_Policy_Claim', 'PoliceReportFiled', 'WitnessPresent', 'AgentType',
       'NumberOfSuppliments', 'AddressChange_Claim', 'BasePolicy',
       'Vehicle Class', 'AccidentArea', 'Age', 'Product Name', 'VehiclePrice',
       'DriverRating', 'AgeOfVehicle']
reduced_data=data[req_col]
def var_sep(x):
    categorical_columns=[]
    numerical_columns=[]
    data_type=pd.DataFrame(x.dtypes).reset_index()
    data_type.columns=["variable_name","data_type"]
    for i in range(0,len(data_type)):
        if data_type.data_type[i]=="object":
            categorical_columns.append(data_type.variable_name[i])
        else:
            numerical_columns.append(data_type.variable_name[i])
    return categorical_columns,numerical_columns

cat_col,num_col=var_sep(reduced_data)
outlier_cols=["Age"]
def outlier_treatment():
    for i in range(0, len(outlier_cols)):
        pecentile10 = np.percentile(data[outlier_cols[i]], 10)
        percentile90 = np.percentile(data[outlier_cols[i]], 90)
        for j in range(0, len(data[outlier_cols[i]])):
            data[outlier_cols[i]][j] = np.where(data[outlier_cols[i]][j] > percentile90, percentile90,
                                                data[outlier_cols[i]][j])
            data[outlier_cols[i]][j] = np.where(data[outlier_cols[i]][j] < pecentile10, pecentile10,
                                                data[outlier_cols[i]][j])
outlier_treatment()

def label_encoding():
    global enc_data
    le = LabelEncoder()
    enc_data = reduced_data.copy()
    for i in cat_col:
        label_encoding = le.fit_transform(enc_data[i])
        enc_data[i] = label_encoding
label_encoding()

def x_y():
    global X
    global Y
    X = enc_data.drop(columns=["FraudFound_P"])
    Y = enc_data["FraudFound_P"]
x_y()

gaussian_pipeline=Pipeline([('Gaussian Naive Bayes',GaussianNB())])
SVC_pipeline=Pipeline([('Support Vector Classifier',SVC(kernel = 'rbf', probability = True))])
DTC_pipeline=Pipeline([('Decision Tree Classifier',DecisionTreeClassifier(random_state = 0))])
RFC_pipeline=Pipeline([('Random Forest Classifier',RandomForestClassifier(n_estimators = 100, random_state = 0))])
GBC_pipeline=Pipeline([('Gradient Boosting Classifier',GradientBoostingClassifier(random_state = 0))])
KNN_pipeline=Pipeline([('K-Nearest Neighbours',KNeighborsClassifier(n_neighbors=5))])

classifier_list=[gaussian_pipeline,SVC_pipeline,DTC_pipeline,RFC_pipeline,GBC_pipeline,KNN_pipeline]
classifiers = [GaussianNB(),
               SVC(kernel = 'rbf', probability = True),
               DecisionTreeClassifier(random_state = 0),
               RandomForestClassifier(n_estimators = 100, random_state = 0),
               GradientBoostingClassifier(random_state = 0),
               KNeighborsClassifier(n_neighbors=5)
              ]
classifier_names = ["Gaussian Naive Bayes",
                    "Support Vector Classifier",
                    "Decision Tree Classifier",
                    "Random Forest Classifier",
                    "Gradient Boosting Classifier",
                   "K-Nearest Neighbours"]
accuracies = []
x_train,x_test,y_train,y_test=train_test_split(X,Y,test_size=0.2)
pred_df2=pd.DataFrame()
def fit_predict():
    for i in range(len(classifier_list)):
        classifier_list[i].fit(x_train, y_train)
        y_pred = classifier_list[i].predict(x_test)
        pred_df2[i] = y_pred

fit_predict()

def acc():
    for i in range(len(classifiers)):
        accuracy = accuracy_score(y_test, pred_df2[i]) * 100
        accuracies.append(accuracy)

acc()

def model_selection():
    for i in range(0, len(accuracies)):
        if accuracies[i] == pd.Series(accuracies).max():
            pickle.dump(classifier_list[i], open('best_fraud_classification_model.pkl', 'wb'))

model_selection()


from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline

from data_processing.transformers.CommonSimilarity import CommonSimilarity
from data_processing.transformers.ToPandas import ToPandas
from recommend.transformers.CosSimilarity import CosSimilarity
from recommend.transformers.FetchSimilarItems import FetchSimilarItems

def recommend(instance, dataset_processed, dataset_preprocessed, num_recommendations):
    """
    Provides recommendations for a specific instance based on processed Steam game dataset.

    This function takes an instance from the dataset and calculates its similarity with other instances
    in the dataset, both numerical and categorical. It then returns a set number of items most similar
    to the given instance.

    Parameters:
    instance (Series): The instance for which recommendations are to be made.
    dataset_processed (DataFrame): The processed dataframe used to calculate similarities.
    dataset_preprocessed (DataFrame): The preprocessed dataframe to fetch recommendation details.
    num_recommendations (int): The number of recommendations to be returned.

    Pipeline steps:
    1. Similarity_preprocessor: Calculates cosine similarity for numerical and categorical columns.
    2. ToPandas1: Converts the numpy array to a pandas dataframe with similarity scores.
    3. CommonSimilarity: Calculates a common similarity score based on numerical and categorical similarity scores.
    4. ToPandas2: Converts the numpy array with common similarity scores back to a pandas dataframe.
    5. FetchSimilarItems: Fetches the required number of items most similar to the given instance.

    Returns:
    DataFrame: A dataframe with the recommended items, including their URL, title, and steam_id.
    """
    print('Recommending...')

    similarity_calc_numerical_columns = ['price_eur', 'release_date', 'steam_rating', 'total_overall_reviews',
                                          'positive_overall_reviews']
    similarity_calc_categorical_columns = [col for col in dataset_processed.columns if 'app_tags' in col]
    columns_to_place_back = similarity_calc_numerical_columns.copy()
    columns_to_place_back.append('num_similarity')
    columns_to_place_back.extend(similarity_calc_categorical_columns)
    columns_to_place_back.append('cat_similarity')
    columns_to_place_back.append('com_similarity')

    index_to_place_back = dataset_processed.index

    instance_index = instance.index
    instance = dataset_processed.loc[instance_index]

    recommendation_pipeline = Pipeline([
        ('similarity_preprocessor', ColumnTransformer(transformers=[
            ('num', CosSimilarity(instance, 'num'), similarity_calc_numerical_columns),
            ('cat', CosSimilarity(instance, 'cat'), similarity_calc_categorical_columns)
        ],
            remainder='passthrough')),
        ('to_pandas1', ToPandas(columns=columns_to_place_back, index=index_to_place_back)),
        ('common_similarity', CommonSimilarity(numerical_weight=1, categorical_weight=0.6)),
        ('to_pandas2', ToPandas(columns=columns_to_place_back, index=index_to_place_back)),
        ('fetch_similar_items', FetchSimilarItems(num_recommendations=num_recommendations))
    ])

    recommended_items = recommendation_pipeline.fit_transform(dataset_processed)
    recommended_items_index = recommended_items.index
    recommended_items['url'] = dataset_preprocessed.loc[recommended_items_index]['url']
    recommended_items['title'] = dataset_preprocessed.loc[recommended_items_index]['title']
    recommended_items['steam_id'] = dataset_preprocessed.loc[recommended_items_index]['steam_id']

    print('Recommending: Done.\n')
    return recommended_items


import torch
import numpy as np
from scipy import interpolate

def load_pretrained(checkpoint_path, model, simmim):

    if not simmim:
        load_pretrained_swin(checkpoint_path, model)
    else:
        load_pretrained_simmim(checkpoint_path, model)

def load_pretrained_swin(checkpoint_path, model):
    checkpoint = torch.load(checkpoint_path, map_location='cpu')
    model_key = 'model' if 'model' in checkpoint else 'model_state'
    ckpt_state_dict = checkpoint[model_key]

    # delete relative_position_index since we always re-init it
    relative_position_index_keys = [k for k in ckpt_state_dict.keys() if "relative_position_index" in k]
    for k in relative_position_index_keys:
        del ckpt_state_dict[k]

    # delete relative_coords_table since we always re-init it
    relative_position_index_keys = [k for k in ckpt_state_dict.keys() if "relative_coords_table" in k]
    for k in relative_position_index_keys:
        del ckpt_state_dict[k]

    # delete attn_mask since we always re-init it
    attn_mask_keys = [k for k in ckpt_state_dict.keys() if "attn_mask" in k]
    for k in attn_mask_keys:
        del ckpt_state_dict[k]

    # bicubic interpolate relative_position_bias_table if not match
    relative_position_bias_table_keys = [k for k in ckpt_state_dict.keys() if "relative_position_bias_table" in k]
    for k in relative_position_bias_table_keys:
        relative_position_bias_table_pretrained = ckpt_state_dict[k]
        relative_position_bias_table_current = model.state_dict()[k]
        L1, nH1 = relative_position_bias_table_pretrained.size()
        L2, nH2 = relative_position_bias_table_current.size()

        assert nH1 == nH1, "Number of heads should be the same"
        if L1 != L2:
            # bicubic interpolate relative_position_bias_table if not match
            S1 = int(L1 ** 0.5)
            S2 = int(L2 ** 0.5)
            relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate(
                relative_position_bias_table_pretrained.permute(1, 0).view(1, nH1, S1, S1), size=(S2, S2),
                mode='bicubic')
            ckpt_state_dict[k] = relative_position_bias_table_pretrained_resized.view(nH2, L2).permute(1, 0)

    # bicubic interpolate absolute_pos_embed if not match
    absolute_pos_embed_keys = [k for k in ckpt_state_dict.keys() if "absolute_pos_embed" in k]
    assert absolute_pos_embed_keys == [], "Absolute pos embedings should be zero"

    # changing head
    head_bias_pretrained = ckpt_state_dict['head.bias']
    Nc1 = head_bias_pretrained.shape[0]
    Nc2 = model.head.bias.shape[0]
    in_features = model.head.in_features
    # changing head for loading model
    if Nc1 != Nc2:
        model.head = torch.nn.Linear(in_features, Nc1)
    model.load_state_dict(ckpt_state_dict, strict=False)

    del checkpoint
    torch.cuda.empty_cache()

def load_pretrained_simmim(checkpoint_path, model):
    checkpoint = torch.load(checkpoint_path, map_location='cpu')
    model_key = 'model' if 'model' in checkpoint else 'model_state'
    ckpt_state_dict = checkpoint[model_key]

    # remove [encoder.] prefix.
    if any([True if 'encoder.' in k else False for k in ckpt_state_dict.keys()]):
        ckpt_state_dict = {k.replace('encoder.', ''): v for k, v in ckpt_state_dict.items() if k.startswith('encoder.')}

    ckpt_state_dict = remap_pretrained_keys_swin(model, ckpt_state_dict)
    model.load_state_dict(ckpt_state_dict, strict=False)

    del checkpoint
    torch.cuda.empty_cache()

def remap_pretrained_keys_swin(model, checkpoint_model):
    state_dict = model.state_dict()

    # Geometric interpolation when pre-trained patch size mismatch with fine-tuned patch size
    all_keys = list(checkpoint_model.keys())
    for key in all_keys:
        if "relative_position_bias_table" in key:
            relative_position_bias_table_pretrained = checkpoint_model[key]
            relative_position_bias_table_current = state_dict[key]
            L1, nH1 = relative_position_bias_table_pretrained.size()
            L2, nH2 = relative_position_bias_table_current.size()

            assert nH1 == nH1, "Number of heads should be the same"
            # Interpolate relative_position_bias_table using geo.
            if L1 != L2:
                src_size = int(L1 ** 0.5)
                dst_size = int(L2 ** 0.5)

                def geometric_progression(a, r, n):
                    return a * (1.0 - r ** n) / (1.0 - r)

                left, right = 1.01, 1.5
                while right - left > 1e-6:
                    q = (left + right) / 2.0
                    gp = geometric_progression(1, q, src_size // 2)
                    if gp > dst_size // 2:
                        right = q
                    else:
                        left = q

                # if q > 1.090307:
                #     q = 1.090307

                dis = []
                cur = 1
                for i in range(src_size // 2):
                    dis.append(cur)
                    cur += q ** (i + 1)

                r_ids = [-_ for _ in reversed(dis)]

                x = r_ids + [0] + dis
                y = r_ids + [0] + dis

                t = dst_size // 2.0
                dx = np.arange(-t, t + 0.1, 1.0)
                dy = np.arange(-t, t + 0.1, 1.0)
                all_rel_pos_bias = []

                for i in range(nH1):
                    z = relative_position_bias_table_pretrained[:, i].view(src_size, src_size).float().numpy()
                    f_cubic = interpolate.interp2d(x, y, z, kind='cubic')
                    all_rel_pos_bias.append(torch.Tensor(f_cubic(dx, dy)).contiguous().view(-1, 1).to(
                        relative_position_bias_table_pretrained.device))

                new_rel_pos_bias = torch.cat(all_rel_pos_bias, dim=-1)
                checkpoint_model[key] = new_rel_pos_bias

    # delete relative_position_index since we always re-init it
    relative_position_index_keys = [k for k in checkpoint_model.keys() if "relative_position_index" in k]
    for k in relative_position_index_keys:
        del checkpoint_model[k]

    # delete relative_coords_table since we always re-init it
    relative_coords_table_keys = [k for k in checkpoint_model.keys() if "relative_coords_table" in k]
    for k in relative_coords_table_keys:
        del checkpoint_model[k]

    # re-map keys due to name change
    rpe_mlp_keys = [k for k in checkpoint_model.keys() if "rpe_mlp" in k]
    for k in rpe_mlp_keys:
        checkpoint_model[k.replace('rpe_mlp', 'cpb_mlp')] = checkpoint_model.pop(k)

    # delete attn_mask since we always re-init it
    attn_mask_keys = [k for k in checkpoint_model.keys() if "attn_mask" in k]
    for k in attn_mask_keys:
        del checkpoint_model[k]

    return checkpoint_model


from __future__ import absolute_import
from __future__ import print_function
import os
import sys
import re
from builtins import str as text
import wx

# -----------------------------------------------------------------------------
# Global variables
# ----------------
#

__author__ = "Pierre Rouleau"
__version__ = "$Revision: 1.5 $"

# -----------------------------------------------------------------------------

def getSupportedLanguageDict(appname):
    """
    Returns dictionary with languages already supported
    by given application

    param: appname:
        name of application
    """
    languageDict = {}
    ext = '.po'
    files = [x for x in os.listdir('.')
             if x.startswith(appname) and x.endswith(ext)]

    langs = [x.split(appname + '_')[1].split(ext)[0] for x in files]
    for lang in langs:
        languageDict[lang] = lang

    return languageDict

def getlanguageDict():
    languageDict = {}
    getSupportedLanguageDict('Beremiz')
    if wx.VERSION >= (3, 0, 0):
        _app = wx.App()
    else:
        _app = wx.PySimpleApp()

    for lang in [x for x in dir(wx) if x.startswith("LANGUAGE")]:
        i = wx.Locale(wx.LANGUAGE_DEFAULT).GetLanguageInfo(getattr(wx, lang))
        if i:
            languageDict[i.CanonicalName] = i.Description

    return languageDict

def verbosePrint(verbose, str):
    if verbose:
        print(str)

def processCustomFiles(filein, fileout, regexp, prefix=''):
    appfil_file = open(filein, 'r')
    messages_file = open(fileout, 'r')
    messages = messages_file.read()
    messages_file.close()
    messages_file = open(fileout, 'a')
    messages_file.write('\n')
    messages_file.write('#: %s\n' % prefix)
    messages_file.write('\n')

    words_found = {}
    for filepath in appfil_file.readlines():
        code_file = open(filepath.strip(), 'r')
        for match in regexp.finditer(code_file.read()):
            word = match.group(1)
            if not words_found.get(word, False) and messages.find("msgid \"%s\"\nmsgstr \"\"" % word) == -1:
                words_found[word] = True
                messages_file.write('\n')
                messages_file.write("msgid \"%s\"\n" % word)
                messages_file.write("msgstr \"\"\n")
        code_file.close()

    messages_file.close()
    appfil_file.close()

# -----------------------------------------------------------------------------
# m a k e P O ( )         -- Build the Portable Object file for the application --
# ^^^^^^^^^^^^^^^
#
def makePO(applicationDirectoryPath,  applicationDomain=None, verbose=0):
    """Build the Portable Object Template file for the application.

    makePO builds the .pot file for the application stored inside
    a specified directory by running xgettext for all application source
    files.  It finds the name of all files by looking for a file called 'app.fil'.
    If this file does not exists, makePo raises an IOError exception.
    By default the application domain (the application
    name) is the same as the directory name but it can be overridden by the
    'applicationDomain' argument.

    makePO always creates a new file called messages.pot.  If it finds files
    of the form app_xx.po where 'app' is the application name and 'xx' is one
    of the ISO 639 two-letter language codes, makePO resynchronizes those
    files with the latest extracted strings (now contained in messages.pot).
    This process updates all line location number in the language-specific
    .po files and may also create new entries for translation (or comment out
    some).  The .po file is not changed, instead a new file is created with
    the .new extension appended to the name of the .po file.

    By default the function does not display what it is doing.  Set the
    verbose argument to 1 to force it to print its commands.
    """

    if applicationDomain is None:
        applicationName = fileBaseOf(applicationDirectoryPath, withPath=0)
    else:
        applicationName = applicationDomain
    currentDir = os.getcwd()
    os.chdir(applicationDirectoryPath)
    filelist = 'app.fil'
    if not os.path.exists(filelist):
        raise IOError(2, 'No module file: %s' % filelist)

    fileout = 'messages.pot'
    # Steps:
    #  Use xgettext to parse all application modules
    #  The following switches are used:
    #
    #   -s                          : sort output by string content (easier to use when we need to merge several .po files)
    #   --files-from=app.fil        : The list of files is taken from the file: app.fil
    #   --output=                   : specifies the name of the output file (using a .pot extension)
    cmd = 'xgettext -s --no-wrap --language=Python --files-from=' + filelist + ' --output=' + fileout + ' --package-name ' + applicationName
    verbosePrint(verbose, cmd)
    os.system(cmd)

    XSD_STRING_MODEL = re.compile(r"<xsd\:(?:element|attribute) name=\"([^\"]*)\"[^\>]*\>")
    processCustomFiles(filelist, fileout, XSD_STRING_MODEL, 'Extra XSD strings')

    XML_TC6_STRING_MODEL = re.compile(r"<documentation>\s*<xhtml\:p><!\[CDATA\[([^\]]*)\]\]></xhtml\:p>\s*</documentation>", re.MULTILINE | re.DOTALL)
    processCustomFiles(filelist, fileout, XML_TC6_STRING_MODEL, 'Extra TC6 documentation strings')

    # generate messages.po
    cmd = 'msginit --no-wrap --no-translator -i %s -l en_US.UTF-8 -o messages.po' % (fileout)
    verbosePrint(verbose, cmd)
    os.system(cmd)

    languageDict = getSupportedLanguageDict(applicationName)

    for langCode in languageDict.keys():
        if langCode == 'en':
            pass
        else:
            langPOfileName = "%s_%s.po" % (applicationName, langCode)
            if os.path.exists(langPOfileName):
                cmd = 'msgmerge -s --no-wrap "%s" %s > "%s.new"' % (langPOfileName, fileout, langPOfileName)
                verbosePrint(verbose, cmd)
                os.system(cmd)
    os.chdir(currentDir)

def catPO(applicationDirectoryPath, listOf_extraPo, applicationDomain=None, targetDir=None, verbose=0):
    """Concatenate one or several PO files with the application domain files.
    """

    if applicationDomain is None:
        applicationName = fileBaseOf(applicationDirectoryPath, withPath=0)
    else:
        applicationName = applicationDomain
    currentDir = os.getcwd()
    os.chdir(applicationDirectoryPath)

    languageDict = getSupportedLanguageDict(applicationName)

    for langCode in languageDict.keys():
        if langCode == 'en':
            pass
        else:
            langPOfileName = "%s_%s.po" % (applicationName, langCode)
            if os.path.exists(langPOfileName):
                fileList = ''
                for fileName in listOf_extraPo:
                    fileList += ("%s_%s.po " % (fileName, langCode))
                cmd = "msgcat -s --no-wrap %s %s > %s.cat" % (langPOfileName, fileList, langPOfileName)
                verbosePrint(verbose, cmd)
                os.system(cmd)
                if targetDir is None:
                    pass
                else:
                    mo_targetDir = "%s/%s/LC_MESSAGES" % (targetDir, langCode)
                    cmd = "msgfmt --output-file=%s/%s.mo %s_%s.po.cat" % (mo_targetDir, applicationName, applicationName, langCode)
                    verbosePrint(verbose, cmd)
                    os.system(cmd)
    os.chdir(currentDir)

def makeMO(applicationDirectoryPath, targetDir='./locale', applicationDomain=None, verbose=0, forceEnglish=0):
    """Compile the Portable Object files into the Machine Object stored in the right location.

    makeMO converts all translated language-specific PO files located inside
    the  application directory into the binary .MO files stored inside the
    LC_MESSAGES sub-directory for the found locale files.

    makeMO searches for all files that have a name of the form 'app_xx.po'
    inside the application directory specified by the first argument.  The
    'app' is the application domain name (that can be specified by the
    applicationDomain argument or is taken from the directory name). The 'xx'
    corresponds to one of the ISO 639 two-letter language codes.

    makeMo stores the resulting files inside a sub-directory of `targetDir`
    called xx/LC_MESSAGES where 'xx' corresponds to the 2-letter language
    code.
    """

    if targetDir is None:
        targetDir = './locale'

    verbosePrint(verbose, "Target directory for .mo files is: %s" % targetDir)

    if applicationDomain is None:
        applicationName = fileBaseOf(applicationDirectoryPath, withPath=0)
    else:
        applicationName = applicationDomain
    currentDir = os.getcwd()
    os.chdir(applicationDirectoryPath)

    languageDict = getSupportedLanguageDict(applicationName)

    for langCode in languageDict.keys():
        if (langCode == 'en') and (forceEnglish == 0):
            pass
        else:
            langPOfileName = "%s_%s.po" % (applicationName, langCode)
            if os.path.exists(langPOfileName):
                mo_targetDir = "%s/%s/LC_MESSAGES" % (targetDir, langCode)
                if not os.path.exists(mo_targetDir):
                    mkdir(mo_targetDir)
                cmd = 'msgfmt --output-file="%s/%s.mo" "%s_%s.po"' % (mo_targetDir, applicationName, applicationName, langCode)
                verbosePrint(verbose, cmd)
                os.system(cmd)
    os.chdir(currentDir)

def printUsage(errorMsg=None):
    """Displays how to use this script from the command line."""
    print("""
    ##################################################################################
    #   mki18n :   Make internationalization files.                                  #
    #              Uses the GNU gettext system to create PO (Portable Object) files  #
    #              from source code, coimpile PO into MO (Machine Object) files.     #
    #              Supports C,C++,Python source files.                               #
    #                                                                                #
    #   Usage: mki18n {OPTION} [appDirPath]                                          #
    #                                                                                #
    #   Options:                                                                     #
    #     -e               : When -m is used, forces English .mo file creation       #
    #     -h               : prints this help                                        #
    #     -m               : make MO from existing PO files                          #
    #     -p               : make PO, update PO files: Creates a new messages.pot    #
    #                        file. Creates a dom_xx.po.new for every existing        #
    #                        language specific .po file. ('xx' stands for the ISO639 #
    #                        two-letter language code and 'dom' stands for the       #
    #                        application domain name).  mki18n requires that you     #
    #                        write a 'app.fil' file  which contains the list of all  #
    #                        source code to parse.                                   #
    #     -v               : verbose (prints comments while running)                 #
    #     --domain=appName : specifies the application domain name.  By default      #
    #                        the directory name is used.                             #
    #     --moTarget=dir : specifies the directory where .mo files are stored.       #
    #                      If not specified, the target is './locale'                #
    #                                                                                #
    #   You must specify one of the -p or -m option to perform the work.  You can    #
    #   specify the path of the target application.  If you leave it out mki18n      #
    #   will use the current directory as the application main directory.            #
    #                                                                                #
    ##################################################################################""")
    if errorMsg:
        print("\n   ERROR: %s" % errorMsg)

def fileBaseOf(filename, withPath=0):
    """fileBaseOf(filename,withPath) ---> string

    Return base name of filename.  The returned string never includes the extension.
    Use os.path.basename() to return the basename with the extension.  The
    second argument is optional.  If specified and if set to 'true' (non zero)
    the string returned contains the full path of the file name.  Otherwise the
    path is excluded.

    [Example]
    >>> fn = 'd:/dev/telepath/tvapp/code/test.html'
    >>> fileBaseOf(fn)
    'test'
    >>> fileBaseOf(fn)
    'test'
    >>> fileBaseOf(fn,1)
    'd:/dev/telepath/tvapp/code/test'
    >>> fileBaseOf(fn,0)
    'test'
    >>> fn = 'abcdef'
    >>> fileBaseOf(fn)
    'abcdef'
    >>> fileBaseOf(fn,1)
    'abcdef'
    >>> fn = "abcdef."
    >>> fileBaseOf(fn)
    'abcdef'
    >>> fileBaseOf(fn,1)
    'abcdef'
    """
    pos = filename.rfind('.')
    if pos > 0:
        filename = filename[:pos]
    if withPath:
        return filename
    else:
        return os.path.basename(filename)

def mkdir(directory):
    """Create a directory (and possibly the entire tree).

    The os.mkdir() will fail to create a directory if one of the
    directory in the specified path does not exist.  mkdir()
    solves this problem.  It creates every intermediate directory
    required to create the final path. Under Unix, the function
    only supports forward slash separator, but under Windows and MacOS
    the function supports the forward slash and the OS separator (backslash
    under windows).
    """

    # translate the path separators
    directory = unixpath(directory)
    # build a list of all directory elements
    aList = filter(lambda x: len(x) > 0, directory.split('/'))
    theLen = len(aList)
    # if the first element is a Windows-style disk drive
    # concatenate it with the first directory
    if aList[0].endswith(':'):
        if theLen > 1:
            aList[1] = aList[0] + '/' + aList[1]
            del aList[0]
            theLen -= 1
    # if the original directory starts at root,
    # make sure the first element of the list
    # starts at root too
    if directory[0] == '/':
        aList[0] = '/' + aList[0]
    # Now iterate through the list, check if the
    # directory exists and if not create it
    theDir = ''
    for i in range(theLen):
        theDir += aList[i]
        if not os.path.exists(theDir):
            os.mkdir(theDir)
        theDir += '/'

def unixpath(thePath):
    r"""Return a path name that contains Unix separator.

    [Example]
    >>> unixpath(r"d:\test")
    'd:/test'
    >>> unixpath("d:/test/file.txt")
    'd:/test/file.txt'
    >>>
    """
    thePath = os.path.normpath(thePath)
    if os.sep == '/':
        return thePath
    else:
        return thePath.replace(os.sep, '/')

# -----------------------------------------------------------------------------

# S c r i p t   e x e c u t i o n               -- Runs when invoked from the command line --
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
if __name__ == "__main__":
    # command line parsing
    import getopt    # pylint: disable=wrong-import-order,wrong-import-position
    argc = len(sys.argv)
    if argc == 1:
        printUsage('Missing argument: specify at least one of -m or -p (or both).')
        sys.exit(1)
    # If there is some arguments, parse the command line
    validOptions = "ehmpv"
    validLongOptions = ['domain=', 'moTarget=']
    option = {}
    option['forceEnglish'] = 0
    option['mo'] = 0
    option['po'] = 0
    option['verbose'] = 0
    option['domain'] = None
    option['moTarget'] = None
    optionKey = {
        '-e':         'forceEnglish',
        '-m':         'mo',
        '-p':         'po',
        '-v':         'verbose',
        '--domain':   'domain',
        '--moTarget': 'moTarget',
    }
    exit_code = 1
    try:
        optionList, pargs = getopt.getopt(sys.argv[1:], validOptions, validLongOptions)
    except getopt.GetoptError as e:
        printUsage(e)
        sys.exit(1)
    for (opt, val) in optionList:
        if opt == '-h':
            printUsage()
            sys.exit(0)
        option[optionKey[opt]] = 1 if val == '' else val
    if len(pargs) == 0:
        appDirPath = os.getcwd()
        if option['verbose']:
            print("No project directory given. Using current one:  %s" % appDirPath)
    elif len(pargs) == 1:
        appDirPath = pargs[0]
    else:
        printUsage('Too many arguments (%u).  Use double quotes if you have space in directory name' % len(pargs))
        sys.exit(1)
    if option['domain'] is None:
        # If no domain specified, use the name of the target directory
        option['domain'] = fileBaseOf(appDirPath)
    if option['verbose']:
        print("Application domain used is: '%s'" % option['domain'])
    if option['po']:
        try:
            makePO(appDirPath, option['domain'], option['verbose'])
            exit_code = 0
        except IOError as e:
            printUsage(text(e) + '\n   You must write a file app.fil that contains the list of all files to parse.')
    if option['mo']:
        makeMO(appDirPath, option['moTarget'], option['domain'], option['verbose'], option['forceEnglish'])
        exit_code = 0
    sys.exit(exit_code)

# -----------------------------------------------------------------------------


import cPickle as pickle

import scipy.io

import numpy as np

import theano.tensor as T
import theano
from theano.sandbox.rng_mrg import MRG_RandomStreams as RandomStreams

from utils import *

#this scrip is designed to produce prediction results from the pickled
#models generated by the utils and cnnlearning scripts

def load_dataset_test():
    """Open the file located at ../train/test.mat and extracts a tuple
    CNNDat,HogDat,FullDat"""
    train = scipy.io.loadmat('../train/test.mat')
    data = train['test'][0,0]
    xcnndat = data[0]
    xhogdat = data[1]
    xfulldat = np.concatenate((xhogdat, xcnndat),axis=1)
    return xcnndat,xhogdat,xfulldat

def model(X, w_h, b_h, w_h2, b_h2, w_o, b_o, p1, p2):
    """Define the output tensor of the model given its inputs and
    parameters."""
    X = dropout(X, srng, p1)
    h = rectify(T.dot(X, w_h)+b_h)

    h = dropout(h, srng, p2)
    h2 = rectify(T.dot(h, w_h2)+b_h2)

    h2 = dropout(h2, srng, p2)
    py_x = softmax(T.dot(h2, w_o)+b_o)
    return py_x

if __name__=='__main__':
    srng = RandomStreams()
    #switch between mult and bin to have results for multiclass or binary
    #classification.
    tog = 'mult'
    #this variable gives the version of the pickled parameters you want to
    #use
    suff = '900'
    if tog == 'mult':
        input_dir = 'models/'
        dmap = {3:np.array([1,0,0,0]),4:np.array([0,1,0,0]),
                2:np.array([0,0,1,0]),1:np.array([0,0,0,1])}
        rmap = {0:3,1:4,2:2,3:1}
        filename = 'pred_multiclass'
    else:
        input_dir = 'models_bin/'
        dmap = {1:np.array([1,0]),2:np.array([1,0]),3:np.array([1,0]),
                4:np.array([0,1])}
        filename = 'pred_binary'
    _,_,x = load_dataset_test()

    #usual theano stuffs
    teX = theano.shared(floatX(x))

    w_h = theano.shared(floatX(pickle.load(open(input_dir+'w_h'+suff,'rb'))))
    w_h2 = theano.shared(floatX(pickle.load(open(input_dir+'w_h2'+suff,'rb'))))
    b_h = theano.shared(floatX(pickle.load(open(input_dir+'b_h'+suff,'rb'))))
    b_h2 = theano.shared(floatX(pickle.load(open(input_dir+'b_h2'+suff,'rb'))))
    w_o = theano.shared(floatX(pickle.load(open(input_dir+'w_o'+suff,'rb'))))
    b_o = theano.shared(floatX(pickle.load(open(input_dir+'b_o'+suff,'rb'))))

    X = T.dmatrix()

    py_x = model(X,w_h,b_h,w_h2,b_h2,w_o,b_o,0.,0.)
    y_x = T.argmax(py_x,axis=1)

    predict = theano.function(inputs=[],outputs=y_x,
                              givens={
                                  X:teX
                              })
    tres = predict()
    res = np.zeros((tres.shape[0],1))
    if tog=='mult':
        for i,it in enumerate(tres):
            res[i] = rmap[tres[i]]
    else:
        for i,it in enumerate(tres):
            res[i] = 1-tres[i]
    scipy.io.savemat(filename+'.mat',mdict={'Ytest':res})


import torch
import torch.nn as nn
import numpy as np
import pickle

from utils.utils import NeighborSampler

class MTL(nn.Module):
    def __init__(self, base_encoder_k, encoder, view_learner, edge_rnn, sample_time_encoder, len_full_edge,
                 train_e_idx_l, train_node_set, train_ts_l, e_feat, device, dim=172, K=600, m=0.999, tau=0.1, gtau=1.0,
                 ratio=0.9, can_nn=20, rnn_nn=20, can_type='3rd'):
        super(MTL, self).__init__()
        self.K = K
        self.m = m
        self.tau = tau
        self.gtau = gtau
        self.ratio = ratio
        self.can_nn = can_nn
        self.rnn_nn = rnn_nn
        self.can_type = can_type

        self.encoder_k = base_encoder_k
        self.encoder = encoder
        self.view_learner = view_learner
        self.edge_rnn = edge_rnn
        self.sample_time_encoder = sample_time_encoder

        self.e_feat = e_feat
        self.len_full_edge = len_full_edge
        self.train_e_idx_l = train_e_idx_l
        self.train_node_set = np.array(list(train_node_set))
        self.train_ts_l = train_ts_l
        self.max_train_ts_l = max(train_ts_l)
        self.device = device

        for param_q, param_k in zip(
            self.encoder.parameters(), self.encoder_k.parameters()
        ):
            param_k.data.copy_(param_q.data)  # initialize
            param_k.requires_grad = False  # not update by gradient

        self.register_buffer("queue", torch.randn(dim, K))
        self.queue = nn.functional.normalize(self.queue, dim=0)
        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))

    @torch.no_grad()
    def _momentum_update_key_encoder(self):
        for param_q, param_k in zip(
            self.encoder.parameters(), self.encoder_k.parameters()
        ):
            param_k.data = param_k.data * self.m + param_q.data * (1.0 - self.m)

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys):
        batch_size = keys.shape[0]
        ptr = int(self.queue_ptr)
        if ptr + batch_size > self.K:
            keys_prev = keys[:self.K - ptr]
            keys_aft = keys[self.K - ptr:]
            self.queue[:, ptr: self.K] = keys_prev.T
            self.queue[:, : len(keys_aft)] = keys_aft.T
            ptr = len(keys_aft)
            self.queue_ptr[0] = ptr
        else:
            self.queue[:, ptr : ptr + batch_size] = keys.T
            ptr = (ptr + batch_size) % self.K  # move pointer
            self.queue_ptr[0] = ptr

    def forward(self, src_l_cut, dst_l_cut, src_l_fake, dst_l_fake, ts_l_cut, eid_l_cut, NUM_NEIGHBORS, time_gap, g, adj_list, adj_list_pickle):
        batch_src_node_embeddings, batch_dst_node_embeddings = self.encoder[0].compute_src_dst_node_temporal_embeddings(src_node_ids=src_l_cut, dst_node_ids=dst_l_cut, node_interact_times=ts_l_cut, num_neighbors=NUM_NEIGHBORS, time_gap=time_gap)
        batch_neg_src_node_embeddings, batch_neg_dst_node_embeddings = self.encoder[0].compute_src_dst_node_temporal_embeddings(src_node_ids=src_l_fake, dst_node_ids=dst_l_fake, node_interact_times=ts_l_cut, num_neighbors=NUM_NEIGHBORS, time_gap=time_gap)
        pos_score = self.encoder[1](input_1=batch_src_node_embeddings, input_2=batch_dst_node_embeddings).squeeze(dim=-1).sigmoid()
        neg_score = self.encoder[1](input_1=batch_neg_src_node_embeddings, input_2=batch_neg_dst_node_embeddings).squeeze(dim=-1).sigmoid()

        train_node_cut = np.array(list(set(np.append(src_l_cut, dst_l_cut))))
        max_ts = np.array([self.max_train_ts_l + 1] * len(train_node_cut))

        neighbor_finder = NeighborSampler(adj_list, sample_neighbor_strategy='uniform')
        neighbor_node_idx, neighbor_edge_idx, neighbor_ts = neighbor_finder.get_historical_neighbors(train_node_cut, max_ts, num_neighbors=self.rnn_nn)

        train_edge_feat = self.view_learner(g)
        full_edge_feat = torch.zeros(self.len_full_edge + 1, train_edge_feat.shape[1], device=self.device)
        full_edge_feat[self.train_e_idx_l - 1] = train_edge_feat

        neighbor_edge_idx = neighbor_edge_idx.reshape(-1) - 1
        neighbor_edge_feat = full_edge_feat[neighbor_edge_idx]  # [bs * 20, 172]
        neighbor_edge_feat = neighbor_edge_feat.reshape(neighbor_node_idx.shape[0], neighbor_node_idx.shape[1], -1)  # [bs, 20, 172]

        neighbor_edge_feat = neighbor_edge_feat.transpose(0, 1)
        _, (h_n, _) = self.edge_rnn(neighbor_edge_feat)
        context_vec = h_n[-1]  # [bs, 172]

        neighbor_finder.sample_neighbor_strategy = 'uniform'
        if self.can_type == '1st':
            candidate_node_idx, candidate_edge_idx, candidate_ts = neighbor_finder.get_historical_neighbors(train_node_cut, max_ts, num_neighbors=self.can_nn)

            src_node_idx_aug = np.repeat(train_node_cut.reshape(train_node_cut.shape[0], 1), candidate_node_idx.shape[1], axis=1)  # [bs, 20]
            dst_node_idx_aug = candidate_node_idx  # [bs, 20]
        elif self.can_type == '3rd':
            candidate_node_idx, candidate_edge_idx, candidate_ts = neighbor_finder.get_multi_hop_neighbors(3, train_node_cut, max_ts, num_neighbors=self.can_nn)
            candidate_node_idx = candidate_node_idx[-1].reshape(train_node_cut.shape[0], -1)
            candidate_edge_idx = candidate_edge_idx[-1].reshape(train_node_cut.shape[0], -1)
            candidate_ts = candidate_ts[-1].reshape(train_node_cut.shape[0], -1)

            src_node_idx_aug = np.repeat(train_node_cut.reshape(train_node_cut.shape[0], 1), candidate_node_idx.shape[1], axis=1)  # [bs, 20]
            dst_node_idx_aug = candidate_node_idx  # [bs, 20]
        elif self.can_type == 'random':
            candidate_node_idx = np.random.choice(self.train_node_set, size=train_node_cut.shape[0] * self.can_nn, replace=True).reshape(train_node_cut.shape[0], -1)  # [bs, 20]
            candidate_edge_idx = np.array([0] * (train_node_cut.shape[0] * self.can_nn)).reshape(train_node_cut.shape[0], -1)
            candidate_ts = np.random.rand(train_node_cut.shape[0], self.can_nn) * self.max_train_ts_l

            src_node_idx_aug = np.repeat(train_node_cut.reshape(train_node_cut.shape[0], 1), candidate_node_idx.shape[1], axis=1)  # [bs, 20]
            dst_node_idx_aug = candidate_node_idx  # [bs, 20]
        elif self.can_type == 'mix':
            candidate_node_idx_1st, candidate_edge_idx_1st, candidate_ts_1st = neighbor_finder.get_historical_neighbors(train_node_cut, max_ts, num_neighbors=self.can_nn)
            candidate_node_idx_3rd, candidate_edge_idx_3rd, candidate_ts_3rd = neighbor_finder.get_multi_hop_neighbors(3, train_node_cut, max_ts, num_neighbors=self.can_nn)
            candidate_node_idx_3rd = candidate_node_idx_3rd[-1].reshape(train_node_cut.shape[0], -1)
            candidate_edge_idx_3rd = candidate_edge_idx_3rd[-1].reshape(train_node_cut.shape[0], -1)
            candidate_ts_3rd = candidate_ts_3rd[-1].reshape(train_node_cut.shape[0], -1)

            candidate_node_idx = np.concatenate((candidate_node_idx_1st, candidate_node_idx_3rd), axis=-1)
            candidate_edge_idx = np.concatenate((candidate_edge_idx_1st, candidate_edge_idx_3rd), axis=-1)
            candidate_ts = np.concatenate((candidate_ts_1st, candidate_ts_3rd), axis=-1)

            src_node_idx_aug = np.repeat(train_node_cut.reshape(train_node_cut.shape[0], 1), candidate_node_idx.shape[1], axis=1)  # [bs, 20]
            dst_node_idx_aug = candidate_node_idx  # [bs, 20]
        else:
            pass

        candidate_edge_idx = candidate_edge_idx.reshape(-1) - 1
        candidate_edge_feat = full_edge_feat[candidate_edge_idx]  # [bs * 20, 172]
        candidate_edge_feat = candidate_edge_feat.reshape(candidate_node_idx.shape[0], candidate_node_idx.shape[1], -1)  # [bs, 20, 172]

        ts_aug = np.random.rand(candidate_ts.shape[0], candidate_ts.shape[1]) * self.max_train_ts_l
        delta_ts_sample = ts_aug - candidate_ts
        delta_ts_sample_context = ts_aug - np.ones_like(candidate_ts) * self.max_train_ts_l
        delta_ts_sample_embedding = self.sample_time_encoder(torch.tensor(delta_ts_sample.reshape(-1, 1), dtype=torch.float32).to(self.device)).reshape(ts_aug.shape[0], ts_aug.shape[1], -1)
        delta_ts_sample_context_embedding = self.sample_time_encoder(torch.tensor(delta_ts_sample_context.reshape(-1, 1), dtype=torch.float32).to(self.device)).reshape(ts_aug.shape[0], ts_aug.shape[1], -1)

        context_vec = context_vec.unsqueeze(1).expand_as(candidate_edge_feat)
        context_vec = context_vec * delta_ts_sample_context_embedding
        candidate_edge_feat = candidate_edge_feat * delta_ts_sample_embedding
        aug_edge_logits = torch.sum(context_vec * candidate_edge_feat, dim=-1)  # [bs, 20, 1]

        # Gumble-Top-K
        bias = 0.0 + 0.0001  # If bias is 0, we run into problems
        eps = (bias - (1 - bias)) * torch.rand(aug_edge_logits.size()) + (1 - bias)
        gate_inputs = torch.log(eps) - torch.log(1 - eps)
        gate_inputs = gate_inputs.to(aug_edge_logits.device)
        gate_inputs = (gate_inputs + aug_edge_logits) / self.gtau
        z = torch.sigmoid(gate_inputs).squeeze()  # [bs, 20]
        __, sorted_idx = z.sort(dim=-1, descending=True)
        k = int(self.ratio * z.size(1))
        keep = sorted_idx[:, :k]  # [bs, k]

        aug_edge_logits = torch.sigmoid(gate_inputs).squeeze()  # [bs, 20]
        aug_edge_weight = torch.gather(aug_edge_logits, dim=1, index=keep)  # [bs, k]
        ts_aug = torch.gather(torch.tensor(ts_aug, device=self.device), dim=1, index=keep).detach().cpu().numpy()  # [bs, k]
        src_node_idx_aug = torch.gather(torch.tensor(src_node_idx_aug, device=self.device), dim=1, index=keep).detach().cpu().numpy()  # [bs, k]
        dst_node_idx_aug = torch.gather(torch.tensor(dst_node_idx_aug, device=self.device), dim=1, index=keep).detach().cpu().numpy()  # [bs, k]
        candidate_edge_feat = torch.gather(candidate_edge_feat, dim=1, index=keep.unsqueeze(2).repeat(1, 1, candidate_edge_feat.shape[2]))  # [bs, k, 172]

        aug_edge_weight = aug_edge_weight.reshape(-1)
        ts_aug = ts_aug.reshape(-1)
        src_node_idx_aug = src_node_idx_aug.reshape(-1)
        dst_node_idx_aug = dst_node_idx_aug.reshape(-1)

        temp_eid = self.len_full_edge
        new_eid_list = []
        adj_list_aug = pickle.loads(adj_list_pickle)
        for src, dst, ts in zip(src_node_idx_aug, dst_node_idx_aug, ts_aug):
            adj_list_aug[src].append((dst, temp_eid, ts))
            adj_list_aug[dst].append((src, temp_eid, ts))
            new_eid_list.append(temp_eid)
            temp_eid += 1

        train_ngh_finder_aug = NeighborSampler(adj_list_aug, sample_neighbor_strategy=self.encoder[0].neighbor_sampler.sample_neighbor_strategy, seed=self.encoder[0].neighbor_sampler.seed)

        new_eid_list = np.array(new_eid_list)
        full_aug_edge_weight = torch.ones(temp_eid, device=self.device)
        full_aug_edge_weight[new_eid_list - 1] = aug_edge_weight

        candidate_edge_feat = candidate_edge_feat.reshape(-1, candidate_edge_feat.shape[2]).detach().cpu().numpy()  # [bs * k, 172]
        e_feat_aug = np.concatenate((self.e_feat, candidate_edge_feat), axis=0)
        edge_raw_embed_aug = torch.from_numpy(e_feat_aug.astype(np.float32)).to(self.device)

        ngh_finder_ori = self.encoder[0].neighbor_sampler
        self.encoder[0].set_neighbor_sampler(train_ngh_finder_aug)
        edge_raw_embed_ori = self.encoder[0].edge_raw_features
        self.encoder[0].edge_raw_features = edge_raw_embed_aug

        max_ts = np.array([self.max_train_ts_l + 1] * len(train_node_cut))

        batch_src_node_embeddings, batch_dst_node_embeddings = self.encoder[0].compute_src_dst_node_temporal_embeddings(src_node_ids=src_l_cut, dst_node_ids=dst_l_cut, node_interact_times=ts_l_cut, num_neighbors=NUM_NEIGHBORS, time_gap=time_gap, full_aug_edge_weight=full_aug_edge_weight)
        batch_neg_src_node_embeddings, batch_neg_dst_node_embeddings = self.encoder[0].compute_src_dst_node_temporal_embeddings(src_node_ids=src_l_fake, dst_node_ids=dst_l_fake, node_interact_times=ts_l_cut, num_neighbors=NUM_NEIGHBORS, time_gap=time_gap, full_aug_edge_weight=full_aug_edge_weight)
        pos_score_ed = self.encoder[1](input_1=batch_src_node_embeddings, input_2=batch_dst_node_embeddings).squeeze(dim=-1).sigmoid()
        neg_score_ed = self.encoder[1](input_1=batch_neg_src_node_embeddings, input_2=batch_neg_dst_node_embeddings).squeeze(dim=-1).sigmoid()

        q = self.encoder[0].compute_node_temporal_embeddings(train_node_cut, max_ts, NUM_NEIGHBORS, time_gap, full_aug_edge_weight=full_aug_edge_weight)
        q = nn.functional.normalize(q, dim=1)

        # Recover
        self.encoder[0].set_neighbor_sampler(ngh_finder_ori)
        self.encoder[0].edge_raw_features = edge_raw_embed_ori

        with torch.no_grad():  # no gradient to keys
            self._momentum_update_key_encoder()  # update the key encoder
            k = self.encoder_k[0].compute_node_temporal_embeddings(train_node_cut, max_ts, NUM_NEIGHBORS, time_gap)
            k = nn.functional.normalize(k, dim=1)

        l_pos = torch.einsum("nc,nc->n", [q, k]).unsqueeze(-1)
        l_neg = torch.einsum("nc,ck->nk", [q, self.queue.clone().detach()])
        logits = torch.cat([l_pos, l_neg], dim=1)
        logits /= self.tau
        labels = torch.zeros(logits.shape[0], dtype=torch.long, device=self.device)
        self._dequeue_and_enqueue(k)

        return pos_score, neg_score, pos_score_ed, neg_score_ed, logits, labels


from sklearn import preprocessing
import pickle
import joblib
import nltk
import pandas as pd
from flask import Flask,request
"""
string = "infection including flu pneumonia immunization diphtheria tetanus child teething \
    infant inflammatory disease including rheumatoid arthritis ra crohn disease blood \
        clot extreme sunburn food poisoning medication including antibiotic"
        """

label_encoder = preprocessing.LabelEncoder()
stopword = nltk.corpus.stopwords.words('english')

def clean(text):
    wn = nltk.WordNetLemmatizer()
    tokens = nltk.word_tokenize(text)
    lower = [word.lower() for word in tokens]
    no_stopwords = [word for word in lower if word not in stopword]
    no_alpha = [word for word in no_stopwords if word.isalpha()]
    lemm_text = [wn.lemmatize(word) for word in no_alpha]
    return lemm_text

def vectorize(data, tfidf_vect_fit):
    X_tfidf = tfidf_vect_fit.transform(data)
    words = tfidf_vect_fit.get_feature_names()
    X_tfidf_df = pd.DataFrame(X_tfidf.toarray())
    X_tfidf_df.columns = words
    return(X_tfidf_df)

def predictions(string):

    strngs = clean(string)
    x_string = ""

    for stng in strngs:
        x_string = x_string+" "+stng

    with open('vector.pkl', 'rb') as vector:
        model_vector = pickle.load(vector)
        vector.close()

    with open('model.pickle', 'rb') as mlmodel:
        model1 = pickle.load(mlmodel)
        mlmodel.close()

    with open('labels.txt', 'rb') as resltdict:
        labels = dict(pickle.load(resltdict))

    pred = model1.predict(model_vector.transform([x_string]))

    return labels[int(pred)]

app = Flask(__name__)

@app.route('/predict', methods=['GET', 'POST'])
def model_code():
    if request.method=="POST":
        strings=request.form['str']
        return predictions(strings)

if __name__=='__main__':
    app.run()