# 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))


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

from absl import flags
import tensorflow as tf

import dataloader
import retinanet_model
from tensorflow.contrib.tpu.python.tpu import tpu_config
from tensorflow.contrib.tpu.python.tpu import tpu_estimator
from tensorflow.contrib.training.python.training import evaluation

# Cloud TPU Cluster Resolvers
flags.DEFINE_string(
    'tpu', default=None,
    help='The Cloud TPU to use for training. This should be either the name '
    'used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 url.')
flags.DEFINE_string(
    'gcp_project', default=None,
    help='Project name for the Cloud TPU-enabled project. If not specified, we '
    'will attempt to automatically detect the GCE project from metadata.')
flags.DEFINE_string(
    'tpu_zone', default=None,
    help='GCE zone where the Cloud TPU is located in. If not specified, we '
    'will attempt to automatically detect the GCE project from metadata.')

# Model specific paramenters
flags.DEFINE_string(
    'eval_master', default='',
    help='GRPC URL of the eval master. Set to an appropiate value when running '
    'on CPU/GPU')
flags.DEFINE_bool('use_tpu', True, 'Use TPUs rather than CPUs')
flags.DEFINE_bool(
    'use_xla', False,
    'Use XLA even if use_tpu is false.  If use_tpu is true, we always use XLA, '
    'and this flag has no effect.')
flags.DEFINE_string('model_dir', None, 'Location of model_dir')
flags.DEFINE_string('resnet_checkpoint', '',
                    'Location of the ResNet50 checkpoint to use for model '
                    'initialization.')
flags.DEFINE_string('hparams', '',
                    'Comma separated k=v pairs of hyperparameters.')
flags.DEFINE_integer(
    'num_shards', default=8, help='Number of shards (TPU cores)')
flags.DEFINE_integer('train_batch_size', 64, 'training batch size')
flags.DEFINE_integer('eval_steps', 5000, 'evaluation steps')
flags.DEFINE_integer(
    'iterations_per_loop', 100, 'Number of iterations per TPU training loop')
flags.DEFINE_string(
    'training_file_pattern', None,
    'Glob for training data files (e.g., COCO train - minival set)')
flags.DEFINE_string(
    'validation_file_pattern', None,
    'Glob for evaluation tfrecords (e.g., COCO val2017 set)')
flags.DEFINE_string(
    'val_json_file',
    '',
    'COCO validation JSON containing golden bounding boxes.')
flags.DEFINE_integer('num_examples_per_epoch', 120000,
                     'Number of examples in one epoch')
flags.DEFINE_integer('num_epochs', 15, 'Number of epochs for training')
flags.DEFINE_string('mode', 'train',
                    'Mode to run: train or eval (default: train)')
flags.DEFINE_bool('eval_after_training', False, 'Run one eval after the '
                  'training finishes.')
# For Eval mode
flags.DEFINE_integer('min_eval_interval', 180,
                     'Minimum seconds between evaluations.')
flags.DEFINE_integer(
    'eval_timeout', None,
    'Maximum seconds between checkpoints before evaluation terminates.')

FLAGS = flags.FLAGS

def main(argv):
  del argv  # Unused.

  if FLAGS.use_tpu:
    tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
        FLAGS.tpu,
        zone=FLAGS.tpu_zone,
        project=FLAGS.gcp_project)
    tpu_grpc_url = tpu_cluster_resolver.get_master()
    tf.Session.reset(tpu_grpc_url)
  else:
    tpu_cluster_resolver = None

  if FLAGS.mode is 'train' and FLAGS.training_file_pattern is None:
    raise RuntimeError('You must specify --training_file_pattern for training.')
  if FLAGS.mode is 'eval':
    if FLAGS.valid_data_dir is None:
      raise RuntimeError('You must specify --valid_data_dir for evaluation.')
    if FLAGS.val_json_file is None:
      raise RuntimeError('You must specify --val_json_file for evaluation.')

  # Parse hparams
  hparams = retinanet_model.default_hparams()
  hparams.parse(FLAGS.hparams)

  params = dict(
      hparams.values(),
      num_shards=FLAGS.num_shards,
      num_examples_per_epoch=FLAGS.num_examples_per_epoch,
      use_tpu=FLAGS.use_tpu,
      resnet_checkpoint=FLAGS.resnet_checkpoint,
      val_json_file=FLAGS.val_json_file,
      mode=FLAGS.mode,
  )
  config_proto = tf.ConfigProto(
      allow_soft_placement=True, log_device_placement=False)
  if FLAGS.use_xla and not FLAGS.use_tpu:
    config_proto.graph_options.optimizer_options.global_jit_level = (
        tf.OptimizerOptions.ON_1)

  run_config = tpu_config.RunConfig(
      cluster=tpu_cluster_resolver,
      evaluation_master=FLAGS.eval_master,
      model_dir=FLAGS.model_dir,
      log_step_count_steps=FLAGS.iterations_per_loop,
      session_config=config_proto,
      tpu_config=tpu_config.TPUConfig(FLAGS.iterations_per_loop,
                                      FLAGS.num_shards))

  # TPU Estimator
  if FLAGS.mode == 'train':
    train_estimator = tpu_estimator.TPUEstimator(
        model_fn=retinanet_model.retinanet_model_fn,
        use_tpu=FLAGS.use_tpu,
        train_batch_size=FLAGS.train_batch_size,
        config=run_config,
        params=params)
    train_estimator.train(
        input_fn=dataloader.InputReader(FLAGS.training_file_pattern,
                                        is_training=True),
        max_steps=int((FLAGS.num_epochs * FLAGS.num_examples_per_epoch) /
                      FLAGS.train_batch_size))

    if FLAGS.eval_after_training:
      # Run evaluation after training finishes.
      eval_params = dict(
          params,
          use_tpu=False,
          input_rand_hflip=False,
          skip_crowd=False,
          resnet_checkpoint=None,
          is_training_bn=False,
          use_bfloat16=False,
      )
      eval_estimator = tpu_estimator.TPUEstimator(
          model_fn=retinanet_model.retinanet_model_fn,
          use_tpu=False,
          train_batch_size=FLAGS.train_batch_size,
          eval_batch_size=1,
          config=run_config,
          params=eval_params)
      eval_results = eval_estimator.evaluate(
          input_fn=dataloader.InputReader(FLAGS.validation_file_pattern,
                                          is_training=False),
          steps=FLAGS.eval_steps)
      tf.logging.info('Eval results: %s' % eval_results)

  elif FLAGS.mode == 'eval':
    # eval only runs on CPU or GPU host with batch_size = 1

    # Override the default options: disable randomization in the input pipeline
    # and don't run on the TPU.
    eval_params = dict(
        params,
        use_tpu=False,
        input_rand_hflip=False,
        skip_crowd=False,
        resnet_checkpoint=None,
        is_training_bn=False,
        use_bfloat16=False,
    )

    eval_estimator = tpu_estimator.TPUEstimator(
        model_fn=retinanet_model.retinanet_model_fn,
        use_tpu=False,
        eval_batch_size=1,
        train_batch_size=FLAGS.train_batch_size,
        config=run_config,
        params=eval_params)

    def terminate_eval():
      tf.logging.info('Terminating eval after %d seconds of no checkpoints' %
                      FLAGS.eval_timeout)
      return True

    # Run evaluation when there's a new checkpoint
    for ckpt in evaluation.checkpoints_iterator(
        FLAGS.model_dir,
        min_interval_secs=FLAGS.min_eval_interval,
        timeout=FLAGS.eval_timeout,
        timeout_fn=terminate_eval):

      tf.logging.info('Starting to evaluate.')
      try:
        eval_results = eval_estimator.evaluate(
            input_fn=dataloader.InputReader(FLAGS.validation_file_pattern,
                                            is_training=False),
            steps=FLAGS.eval_steps)
        tf.logging.info('Eval results: %s' % eval_results)

        # Terminate eval job when final checkpoint is reached
        current_step = int(os.path.basename(ckpt).split('-')[1])
        total_step = int((FLAGS.num_epochs * FLAGS.num_examples_per_epoch) /
                         FLAGS.train_batch_size)
        if current_step >= total_step:
          tf.logging.info('Evaluation finished after training step %d' %
                          current_step)
          break

      except tf.errors.NotFoundError:
        # Since the coordinator is on a different job than the TPU worker,
        # sometimes the TPU worker does not finish initializing until long after
        # the CPU job tells it to start evaluating. In this case, the checkpoint
        # file could have been deleted already.
        tf.logging.info('Checkpoint %s no longer exists, skipping checkpoint' %
                        ckpt)
  else:
    tf.logging.info('Mode not found.')

if __name__ == '__main__':
  tf.logging.set_verbosity(tf.logging.INFO)
  tf.app.run(main)


from sqlitedatabase import add_entry,get_entries,create_connection,create_table

menu = """ Welcome to the programming diary!

Please select one of the following options:
1) Add new entry for today.
2) View entries.
3) Exit.

Your selection:
"""
welcome = "**Welcome to the programing diary!**"

# entries = [
#     {"content": "Today I started learning programing.", "date": "01-01-2020"},
#     {"content": "I created my first SQLite database!", "date": "02-01-2020"},
#     {"content": "I finished writing my programming diary application.", "date": "03-01-2020"},
#     {"content": "Today I'm going to continue learning programming!", "date": "04-01-2020"},
# ]

def prompt_new_entry():
    entry_content = input("What have you learned today? ")
    entry_date = input("Enter the date: ")

    add_entry(entry_content, entry_date)

def view_entries(entries):
    for entry in entries:
        print(f"{entry['date']}\n{entry['content']}\n\n")

# print(welcome)
# while (user_input := input(menu)) != "3":
#     if user_input == "1":
#         prompt_new_entry()
#     elif user_input == "2":
#         view_entries(get_entries())
#     else:
#         print("Invalid option, please try again!")

def main():
    database = r"/Users/manishgarg/software/sqlite/database.db"

    # create a database connection
    connection = create_connection(database)
    with connection:
        create_table(connection)
        add_entry(connection,"Learning SQLite with Python" ,"2020-06-01")
        entries = get_entries(connection)
        for entry in entries:
            print(entry)

if __name__ == '__main__':
    main()


# n = int(input())
# times = []
#
# for _ in range(n):
#     start, end = map(int, input().split())
#     times.append([start, end])
#
# times = sorted(times, key=lambda x: x[0])
# times = sorted(times, key=lambda x: x[1])
#
# finish_time = 0
# count = 0
#
# for start, end in times:
#     if start >= finish_time:
#         count += 1
#         finish_time = end
#
# print(count)

'''
LIS - dp (dynamic programming)
백준 - 전봇대
'''

n = int(input())
times = []

for _ in range(n):
    start, end = map(int, input().split())
    times.append([start, end])

time = sorted(times, key=lambda x: (x[0], x[1]))

finish_time = 0
cnt = 0

for start, end in time:
    if start > finish_time:
        cnt += 1
        finish_time = end

print(cnt)


from datetime import datetime

import scrapy
from scrapy.http import Request

from maksavit_scrapy.items import ScrapyMaksavitItem
from maksavit_scrapy.settings import (CATEGORIES, DOMAIN, LOCATION, MAIN_URL,
                                      PROXY)

class MaksavitSpider(scrapy.Spider):
    name = 'maksavit'
    page = 1
    allowed_domains = [DOMAIN]
    start_urls = []

    def start_requests(self):
        for category in CATEGORIES:
            self.start_urls.append(f'{MAIN_URL}{category}?location={LOCATION}')

        for url in self.start_urls:
            yield Request(url, dont_filter=True, meta=PROXY)

    def parse(self, response):
        print(response.request.url)
        max_page = int(
            response.css(
                'div.app-filter__pagination a::text').getall()[-1].strip()
        )

        for product in response.css('div.product-card-block'):
            link = product.css('a.product-card-block__title::attr(href)').get()
            yield response.follow(link, callback=self.parse_item)
        current_page = int(
            response.css('a.ui-pagination__item_checked::text').get().strip()
        )
        next_page = current_page + 1

        if current_page < max_page:
            url = f'{response.request.url}'.replace(
                f'&page={current_page}', '')
            yield response.follow(
                f'{url}&page={next_page}', callback=self.parse
            )

    def parse_item(self, response):
        timestamp = int(datetime.timestamp(datetime.now()))
        code = response.request.url.split('/')[-2]
        url = response.request.url
        title = response.css('h1.product-top__title::text').get()
        badge = response.css('div.badge-discount::text').get()
        badges = [badge.strip()] if badge is not None else badge
        production = response.css('a.product-info__brand-value::text').get()
        brand = (None if production is None else ', '.join(
            production.strip().split(', ')[:-1]))
        country = (
            None if production is None else production.strip().split(', ')[-1]
        )
        sections = response.css(
            'ul.breadcrumbs li a.breadcrumbs__link span::text').getall()

        current_price = response.css('div.price-info__price span::text').get()
        current_price = (
            float(current_price.split(' ')[0])
            if current_price is not None else None)

        original_price = response.css('div.price-box__old-price::text').get()
        original_price = (current_price if original_price is None else float(
            original_price.replace('\xa0₽', '').replace(' ', '').strip()))

        discount = (((original_price - current_price) / original_price)
                    * 100 if current_price != original_price else 0)
        sale_tag = f'Скидка {discount:.0f}%' if discount != 0 else None

        status = response.css('div.available-count ::text').get()
        status = None if status is None else status[:-1]
        in_stock = True if status == 'сегодня' else False
        # Количество подсчитать невозможно.
        # Ставлю дефолтное None, как указано в ТЗ
        count = None

        img = response.css('div.product-picture img::attr(src)').get()
        main_img = response.urljoin(img) if img != '' else None
        set_img = [response.urljoin(img) if img != '' else None
                   for img in response.css(
            'div.product-picture img::attr(src)').getall()]

        # На сайте не представлено ни одного фото360/видео.
        # Т.к. в ТЗ эти элементы перечислены обязательными,
        # добавил round_view и video для демонстрации их поиска
        round_view = (
            response.css('div.product-picture')
            .xpath('//img[contains(@src, "360view")]/@src')
            .getall()
        )
        video = (
            response.css('div.product-picture')
            .xpath('//img[contains(@src, "video")]/@src')
            .getall()
        )

        info = response.css('div.product-instruction__guide')

        description = info.css('div.ph23 ::text').getall()
        description = (
            ''.join(description).replace('Описание', '')
            if description
            else None
        )

        form = info.css('div.ph14 ::text').getall()
        form = ''.join(form).replace('Форма выпуска', '') if form else None

        exp = info.css('div.ph17 ::text').getall()
        exp = ''.join(exp).replace('Срок годности', '') if exp else None

        quanity = response.css(
            'div.quantity-items-wrapper span::text').getall()
        quanity = len(quanity) if quanity else 1

        price = {
            'current': current_price,
            'original': original_price,
            'sale_tag': sale_tag,
        }
        assets = {
            'main_image': main_img,
            'set_images': set_img,
            'view360': round_view,
            'video': video,
        }
        stock = {'in_stock': in_stock, 'count': count}
        metadata = {
            'description': description,
            'code': code,
            'form': form,
            'exp': exp,
            'production': brand,
            'country': country,
        }
        data = {
            'timestamp': timestamp,
            'RPC': code,
            'url': url,
            'title': title,
            'marketing_tags': badges,
            'brand': brand,
            'section': sections,
            'price': price,
            'stock': stock,
            'assets': assets,
            'metadata': metadata,
            'variants': quanity,
        }
        yield ScrapyMaksavitItem(data)


import math
from tkinter import *
size = 600
radius1 = size / 2.9
#radius orbiti
radius2 = size / 100
#radius vrashaushegosya kruga
def coords(angle):
    x = math.cos(angle) * radius1
    y = math.sin(angle) * radius1
    return x - radius2 + size / 2, y - radius2 + size / 2, x + radius2 + size / 2, y + radius2 + size / 2

def motion(angle):
    angle = angle + speed
    c.coords(ball, coords(angle))
    root.after(int(abs(speed * 1000)), lambda: motion(angle))

angle = 0
speed = 0.1

root = Tk()
root.geometry('600x600')
root.configure(background = "grey")
c = Canvas(root, width=size, height=size)
c.pack()

c.create_oval(100, 100, 500, 500, width = 3)

ball = c.create_oval(250,50,350,150,fill = 'black')

motion(angle)

root.mainloop()

Знакомство с tkinter.py


# Square root digital expansion

# Problem 80

# It is well known that if the square root of a natural number is not an integer, then it is irrational.
# The decimal expansion of such square roots is infinite without any repeating pattern at all.
# The square root of two is 1.41421356237309504880..., and the digital sum of the first one hundred decimal digits is 475.
# For the first one hundred natural numbers, find the total of the digital sums of the first one hundred decimal digits for all the irrational square roots.

# https://projecteuler.net/problem=80

import Utilities
import decimal
decimal.getcontext().prec = 110

nums = list()
for i in range(2,101):
    if not Utilities.IsSquareNum(i):
        nums.append(i)
sm = 0
for i in nums:
    sqr = decimal.Decimal(i).sqrt()
    s1 = 0
    for c in str(sqr)[0:101]:
        if c == '.':
            continue
        s1 += int(c)
    print(i, s1, sqr)
    sm += s1
print(sm)


from scrapy.contrib.spiders import CrawlSpider, Rule
from scrapy.spider import BaseSpider
from scrapy.selector import HtmlXPathSelector

from cssspy.utils import domains_from_urls, absolute_urls
from cssspy.cssscrapy.items import CssFilesItem
from scrapy.contrib.linkextractors.htmlparser import HtmlParserLinkExtractor

class CssSpiderMixin(object):
    name = "css"

    def __init__(self, urls=None, *args, **kwargs):
        self.start_urls = urls
        self.allowed_domains = domains_from_urls(*self.start_urls)

    def parse_item(self, response):
        hxs = HtmlXPathSelector(response)
        css_links = hxs.select('//link[@rel="stylesheet"]/@href').extract()
        item = CssFilesItem()
        item['page_url'] = response.url
        item['page_body'] = response.body
        item['css_urls'] = absolute_urls(response.url, css_links)
        item['dead_selectors'] = []
        return item

class CssCrawlSpider(CssSpiderMixin, CrawlSpider):

    rules = (
        Rule(HtmlParserLinkExtractor(), callback='parse_item', follow=True
         ),
    )

    def __init__(self, *args, **kwargs):
        CrawlSpider.__init__(self, *args, **kwargs)
        CssSpiderMixin.__init__(self, *args, **kwargs)

class CssBaseSpider(CssSpiderMixin, BaseSpider):

    def __init__(self, *args, **kwargs):
        BaseSpider.__init__(self, *args, **kwargs)
        CssSpiderMixin.__init__(self, *args, **kwargs)

    def parse(self, response):
        return self.parse_item(response)


from typing import Dict, List

import asana

from giges.models.team import Team
from giges.slack import SlackClient
from giges.tasks.app import app
from giges.util import validate_uuid

def _add_ds_class_item(custom_field: Dict[str, str]) -> str:
    """
    Returns the slack representation for a Data Science item.

    :param custom_field: the unprocessed value from asana API
    :return: a string with an emoji and the name
    """
    return f':large_green_square: {custom_field["name"]}'

def _add_pf_class_item(custom_field: Dict[str, str]) -> str:
    """
    Returns the slack representation for a Platform item.

    :param custom_field: the unprocessed value from asana API
    :return: a string with an emoji and the name
    """
    return f':large_blue_square: {custom_field["name"]}'

def _add_section(memberships: List[Dict], team_projects: List[str]) -> str:
    """
    Returns the name of the section of the workflow project
    the selected task is currently on.

    :param memberships: the memberships response from asana API
    :param team_projects: a list of the ids of the team projects
    :return: a string with the name of the section
    """
    section = ""
    for membership in memberships:
        if membership.get("project"):
            if membership["project"].get("gid") in team_projects:
                section = membership["section"]["name"]

    return section

def _add_class_of_service(custom_field: Dict[str, str]) -> str:
    """
    Returns the slack representation for the class of service of an item.

    :param custom_field: the unprocessed value from asana API
    :return: a string with an emoji
    """
    emojis = {
        "1200760323624685": ":rocket:",
        "1200760323624695": ":bomb:",
        "1200760323630880": ":black_square_button:",
        "1200760323631982": ":brain:",
    }

    gid = custom_field["gid"] if custom_field else ""

    return f'{emojis.get(gid, ":chipmunk:")} '

def daily_stick(team_id: str = None) -> None:
    """
    Wrap the stick inside the app context to be called from SNS
    :param team_id: the giges UUID of the team
    """
    with app.app_context():
        stick(team_id)

def stick(team_id: str = None) -> None:
    """
    For all the projects that belongs to a team:
        - list all the current tasks
        - assigned to each person in the team
        - annoy the sufficient amount

    :param team_id: the giges UUID of the team
    """
    asana_client = asana.Client.access_token(app.config["ASANA_TOKEN"])
    slack_client = SlackClient()
    workspace_id = app.config["ASANA_WORKSPACE"]

    if validate_uuid(team_id):
        team = Team.query.filter_by(id=team_id).one_or_none()
    else:
        # Search for the Tech team by default
        team = Team.query.filter_by(name="Tech").one_or_none()

    projects_ids = [p.external_id for p in team.projects]

    for human in team.tesseras:
        message = f"<@{human.slack_id}|{human.name}> :robot_face: :qoobe_hi:"
        tasks = asana_client.tasks.search_tasks_for_workspace(
            workspace_gid=workspace_id,
            params={
                "assignee.any": human.asana_id,
                "projects.any": ",".join(projects_ids),
                "completed": False,
                "is_subtask": False,
                # 1200200605652838 = In Progress
                "custom_fields.1200200605652836.value": 1200200605652838,
            },
            options={"fields": "followers,assignee", "opt_pretty": True},
        )
        for t in tasks:
            task = asana_client.tasks.find_by_id(t["gid"])
            message += "\n\t -  <https://app.asana.com/0/"
            message += f"{workspace_id}/{task['gid']}/f|{task['name']}> \t"

            fields = {
                "section": _add_section(task["memberships"], projects_ids)
            }
            for field in task["custom_fields"]:
                if field["gid"] == "1200760323623705":
                    fields["class"] = _add_class_of_service(
                        field["enum_value"]
                    )
                elif (
                    field["gid"] == "1200760343663992" and field["enum_value"]
                ):
                    fields["item"] = _add_ds_class_item(field["enum_value"])
                elif (
                    field["gid"] == "1200765868551790" and field["enum_value"]
                ):
                    fields["item"] = _add_pf_class_item(field["enum_value"])
            message += f"{fields.get('class', ':memo:')} "
            message += f"{fields.get('section', ':fairy:')} "
            message += f"{fields.get('item', ':dragon:')}"

        message += "\n:speech_balloon: -> :thread:\n"
        slack_client.send_to_road_blocks(message)


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()


#!/usr/bin/env python

import sys

from setuptools import find_packages, setup

setup_requires = []

# I only release from OS X so markdown/pypandoc isn't needed in Windows
if not sys.platform.startswith('win'):
    setup_requires.extend([
        'setuptools-markdown',
    ])

setup(
    name='serplint',
    author='Beau Gunderson',
    author_email='beau@beaugunderson.com',
    url='https://github.com/beaugunderson/serplint',
    description='A linter for the serpent language',
    long_description_markdown_filename='README.md',
    keywords=['serpent', 'ethereum'],
    version='1.4.0',
    license='MIT',
    packages=find_packages(),
    py_modules=['serplint'],
    entry_points={
        'console_scripts': [
            'serplint = serplint:serplint',
        ]
    },
    install_requires=[
        'click==6.7',
        'ethereum-serpent-augur-temp==2.0.2',
    ],
    setup_requires=setup_requires,
    classifiers=[
        'Development Status :: 5 - Production/Stable',

        'Intended Audience :: End Users/Desktop',

        'License :: OSI Approved :: MIT License',

        'Operating System :: POSIX',
        'Operating System :: Microsoft :: Windows',
        'Operating System :: MacOS :: MacOS X',

        'Topic :: Utilities',

        'Programming Language :: Python :: 2',
        'Programming Language :: Python :: 3',
    ])


import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

class TrainTestSplitter:
    def __init__(self, subjects, labels):
        """
        Initializes the TrainTestSplitter class.

        Parameters:
        - subjects (list): List of subjects.
        - labels (list): List of labels.
        """
        self.subjects = subjects
        self.labels = labels

    def train_test_losego(self, df, seg_index=0, num_seg=1, samp_freq=1):
        """
        Leave one segment out for all modalities.

        Parameters:
        - df (pd.DataFrame): Input DataFrame containing raw data.
        - seg_index (int): Index of the segment to be left for the test dataset (default: 0).
        - num_seg (int): Number of segments to include in the test dataset (default: 1).

        Returns:
        - pd.DataFrame, pd.DataFrame: Train and test DataFrames.
        """
        # Initialize empty DataFrames for train and test datasets
        df_train = pd.DataFrame()
        df_test = pd.DataFrame()

        # Loop through each subject
        for i in range(len(self.subjects)):
            # Loop through each label
            for j in range(len(self.labels)):
                # Get the data for the current subject and label
                data = df[(df['Subject'] == self.subjects[i]) & (df['Label_ori'] == self.labels[j])]['Data'].values
                data = np.array([np.array(xi) for xi in data])

                # Select the test data based on the segment index and number of segments
                if samp_freq == 1:
                    #print('Sampling frequency is 1')
                    data_test = data[seg_index:(seg_index) + num_seg]
                elif samp_freq == 8: #
                    #print('Sampling frequency is 8')
                    data_test = data[seg_index*samp_freq:(seg_index*samp_freq) + num_seg]

                # Remove the test data from the training data
                data_train = np.delete(data, seg_index, axis=0)

                # Create DataFrames for the train and test data
                df_one_train = self.create_dataframe(data_train)
                df_one_test = self.create_dataframe(data_test)

                # Concatenate the train and test DataFrames
                df_train = pd.concat([df_train, df_one_train]).reset_index(drop=True)
                df_test = pd.concat([df_test, df_one_test]).reset_index(drop=True)

        # Return the train and test DataFrames
        return df_train, df_test

    def train_test_loso(self, df, subject_index=0):
        """
        Leave one subject out for all modalities.

        Parameters:
        - df (pd.DataFrame): Input DataFrame containing raw data.
        - subject_index (int): Index of the subject to be selected as the test subject (default: 0).

        Returns:
        - pd.DataFrame, pd.DataFrame: Train and test DataFrames.
        """
        # Initialize empty DataFrames for train and test data
        df_train = pd.DataFrame()
        df_test = pd.DataFrame()

        # Get train data by excluding the test subject's data from the DataFrame
        data_train = df[(df['Subject'] != self.subjects[subject_index])]['Data'].values
        data_train = np.array([np.array(xi) for xi in data_train])

        # Get test data for the selected subject
        data_test = df[(df['Subject'] == self.subjects[subject_index])]['Data'].values
        data_test = np.array([np.array(xi) for xi in data_test])

        # Create DataFrames for train and test data
        df_one_train = self.create_dataframe(data_train)
        df_one_test = self.create_dataframe(data_test)

        # Concatenate train and test DataFrames
        df_train = pd.concat([df_train, df_one_train]).reset_index(drop=True)
        df_test = pd.concat([df_test, df_one_test]).reset_index(drop=True)

        # Return the train and test DataFrames
        return df_train, df_test

    def create_dataframe(self, data):
        """
        Convert NumPy arrays back to a DataFrame.

        Parameters:
        - data (np.ndarray): Input NumPy array.

        Returns:
        - pd.DataFrame: DataFrame created from the NumPy array.
        """
        # Create an empty DataFrame
        df_one = pd.DataFrame()

        # Extract the data from the NumPy array and assign it to the 'Data' column
        df_one['Data'] = list(data[:, :-3, :])

        # Extract the label from the NumPy array and assign it to the 'Label' column
        df_one['Label'] = data[:, -3, 0]

        # Extract the original label from the NumPy array and assign it to the 'Label_ori' column
        df_one['Label_ori'] = data[:, -2, 0]

        # Extract the subject from the NumPy array and assign it to the 'Subject' column
        df_one['Subject'] = data[:, -1, 0]

        # Return the created DataFrame
        return df_one

    def train_val_split(self, df, test_size=0.2):
        """
        Split the DataFrame into train and test sets.

        Parameters:
        - df (pd.DataFrame): Input DataFrame containing raw data.
        - test_size (float): Proportion of the dataset to include in the test split (default: 0.2).

        Returns:
        - pd.DataFrame, pd.DataFrame: Train and test DataFrames.
        """
        # Split the DataFrame into train and test sets using the specified test size
        df_train, df_val = train_test_split(df, test_size=test_size)

        # Return the train and test DataFrames
        return df_train, df_val


def solution(people, limit):
    people = sorted(people, reverse=True)
    start = 0
    end = len(people) - 1
    counter = 0
    while start <= end:
        weight_heavy = people[start]
        weight_light = people[end]
        if weight_heavy + weight_light <= limit:
            end -= 1
        start += 1
        counter += 1
    return counter

if __name__ == '__main__':
    assert solution([70, 50, 80, 50], 100) == 3
    assert solution([70, 80, 50], 100) == 3


import h2o
import numpy as np

# Start H2O on your local machine
h2o.init ( ip='localhost', port=54321, nthreads=-1, max_mem_size='25g' )
# Import the train_numeric
train_numeric = h2o.import_file ( path="/Users/avinashbarnwal/Desktop/Kaggle/Bosch/train_numeric.csv" )
# train_categorical = h2o.import_file(path = "/Users/avinashbarnwal/Desktop/Kaggle/Bosch/train_categorical.csv")
# train_date = h2o.import_file(path = "/Users/avinashbarnwal/Desktop/Kaggle/Bosch/train_date.csv")

# train_numeric.describe()
train_numeric["Response"] = train_numeric["Response"].asfactor ( )

# train = train_numeric.cbind(train_categorical)
# train=  train.cbind(train_date)

train_numeric_name = train_numeric.col_names
# train_categorical_name = train_categorical.col_names
# train_date_name = train_date.col_names

# Yielding Colnames of train_numeric
t = train_numeric.shape
# Yielding Dimension of train_numeric
removelist_train_numeric = [0, (t[1] - 1)]
# Removing Id and response
train_numeric_name = [v for i, v in enumerate ( train_numeric_name ) if i not in removelist_train_numeric]

# print train_numeric_name

# removelist_train_categorical = [0]
# train_categorical_name   =     [v for i, v in enumerate(train_categorical_name) if i not in removelist_train_categorical]

# removelist_train_date = [0]
# train_date_name = [v for i, v in enumerate(train_date_name) if i not in removelist_train_date]

# Concatenate the arrays
# train_name =sum([train_numeric_name, train_categorical_name, train_date_name],[])

# Train GBM model

gbm_model = h2o.estimators.gbm.H2OGradientBoostingEstimator ( ntrees=100,
                                                              max_depth=5,
                                                              learn_rate=0.1,
                                                              sample_rate=0.8,
                                                              col_sample_rate=0.8,
                                                              seed=2016,
                                                              nfolds=5,
                                                              fold_assignment="Stratified",
                                                              keep_cross_validation_predictions='TRUE',
                                                              keep_cross_validation_fold_assignment='TRUE',
                                                              score_tree_interval=20,
                                                              stopping_rounds=10,
                                                              stopping_metric="AUC",
                                                              stopping_tolerance=0.01 )

gbm_model.train ( x=train_numeric_name, y="Response", training_frame=train_numeric )

importance = gbm_model.varimp ( )
# print importance

# importance <- as.dataframe(importance)
# importance <- importance[scaled_importance>0]


__author__ = 'M_Nour'

import numpy as np
from scipy import stats
from sklearn.semi_supervised import label_propagation
from sklearn.metrics import classification_report, confusion_matrix,accuracy_score, f1_score, recall_score
import dataset
from collections import Counter
import matplotlib.pyplot as plt
from os.path import join
from dataset import calc_features
from imblearn.over_sampling import SMOTE
from sklearn import preprocessing
from sklearn.dummy import DummyClassifier
from sklearn.ensemble import RandomForestClassifier

def make_model3():
    model =  label_propagation.LabelSpreading(kernel='knn', n_neighbors=15)
    sensor_data = dataset.load_data()
    X, y = sensor_data.data[:200], sensor_data.target[:200]
    model.fit(X, y)
    np.savetxt("X.csv", X, delimiter=",", fmt='%10.5f')
    np.savetxt("y_train.csv", y, delimiter=",", fmt='%10.1f')
    return model

def make_model():
    model = RandomForestClassifier(n_estimators=185)#label_propagation.LabelSpreading(gamma=0.25, max_iter=5)
    sensor_data = dataset.load_data()
    X, y = sensor_data.data[:200], sensor_data.target[:200]
    model.fit(X, y)
    np.savetxt("X.csv", X, delimiter=",", fmt='%10.5f')
    np.savetxt("y_train.csv", y, delimiter=",", fmt='%10.1f')
    return model

def make_model2():
    sensor_data = dataset.load_data()
    rng = np.random.RandomState(0)
    indices = np.arange(len(sensor_data.data))
    rng.shuffle(indices)
    print(len(sensor_data.data))
    sm = SMOTE(random_state=42)
    X, y  = sm.fit_sample(sensor_data.data[indices[:2000]], sensor_data.target[indices[:2000]])

    n_total_samples = len(y)
    print(len(y))
    n_labeled_points = 200
    max_iterations = 50
    unlabeled_indices = np.arange(n_total_samples)[n_labeled_points:]
    lp_model = label_propagation.LabelSpreading(kernel='knn', n_neighbors=15)

    for i in range(max_iterations):
        if len(unlabeled_indices) == 0:
            print("No unlabeled items left to label.")
            break
        y_train = np.copy(y)
        y_train[unlabeled_indices] = -1
        lp_model.fit(X, y_train)
        p = lp_model.predict_proba(X[unlabeled_indices])
        # predicted_labels = [1 if x > 0.57 else 0 for x in p[:, 1]]
        predicted_labels = lp_model.predict(X[unlabeled_indices])

        true_labels = y[unlabeled_indices]
        # print("#"*20 + "Iteration :: " + str(i) + "#"*20)
        # print(classification_report(true_labels, predicted_labels))

        pred_entropies = stats.distributions.entropy(
            lp_model.label_distributions_.T)
        uncertainty_index = np.argsort(pred_entropies)[::-1]
        uncertainty_index = uncertainty_index[
                                np.in1d(uncertainty_index, unlabeled_indices)][:40]
        delete_indices = np.array([])
        for index in uncertainty_index:
            delete_index, = np.where(unlabeled_indices == index)
            delete_indices = np.concatenate((delete_indices, delete_index))
        unlabeled_indices = np.delete(unlabeled_indices, delete_indices)
        n_labeled_points += len(uncertainty_index)

    np.savetxt("X.csv", X, delimiter=",", fmt='%10.5f')
    np.savetxt("y_train.csv", y_train, delimiter=",", fmt='%10.1f')
    return lp_model

def update_data(mdl , x, aq, entropy, threshold):
    x = x.reshape(1, -1)
    p = mdl.predict_proba(x)
    e = stats.distributions.entropy(p.T)
    entropy =np.append(entropy, e)
    if e > threshold and aq>0:
        aq-=1
    return aq,entropy

def customize_model(mdl, W_X, W_Y):
    print("@"*100)
    print(len(W_Y))
    sm = SMOTE(random_state=42)
    # mdl = label_propagation.LabelSpreading(gamma=0.25, max_iter=5)

    X = np.genfromtxt('X.csv', delimiter=',')
    y_train = np.genfromtxt('y_train.csv', delimiter=',')
    # y_train = y_train.reshape(-1,1)
    # W_Y = W_Y.reshape(-1, 1)
    # h = np.arange(len(W_Y))[[i!= -1 for i in W_Y ]]
    # W_Y = W_Y[h]
    # W_X = W_X[h]
    y_train = np.append(y_train, W_Y, axis=0)
    X = np.append(X, W_X, axis=0)
    mdl.fit(X, y_train)

    return mdl

def update_threshold_adaptive_lambda(entropy, k , N, fold_count, AR):

    threshold = sorted(entropy)[::-1][int(AR*k*fold_count/N)]
    return threshold

def update_threshold_static_lambda():
    threshold= 0.65
    return threshold

def customize(mdl, method, name, fold_count, participant_id, AR = 60):
    f_result = []
    acc_result = []
    rec_result =[]
    f_macro_result = []
    rec_macro_result = []
    threshold = 0.691
    entropy = np.array([])

    total_query = 0
    aq = AR
    user_data = np.genfromtxt(join('C:/Users/Marjan/Desktop/Thesis/participants_wild',participant_id, 'wrist_ss.csv'), delimiter=',')
    user_data = user_data[:, 1:]
    data = calc_features(user_data)
    X = preprocessing.scale(data[:, :(user_data.shape[1] - 1) * 9])
    print("*"*20)
    print(X.shape)
    X = np.tile(X, (2, 1))
    #X = np.tile(X, 10) #repeat 10 times
    print(X.shape)
    print("*" * 20)
    y = data[:, -1]
    y = np.tile(y, 2)
    print(y.shape)
    W_Y = np.copy(y)
    print(W_Y.shape)

    N = len(X)
    unlabeled_indices = np.arange(N)
    W_Y [unlabeled_indices] = -1
    # size = int(N//fold_count)
    size = 3600 / 3  # number of instances per hour
    for k, row in enumerate(X):
        if k%size ==  size-1:
            total_query+=(AR- aq)
            aq = AR

            mdl = customize_model(mdl, X[:k+1], W_Y[:k+1])
            predicted_labels = mdl.predict(X[unlabeled_indices])
            # p = mdl.predict_proba(X[unlabeled_indices])
            # predicted_labels = [1 if x > 0.50 else 0 for x in p[:, 1]]

            true_labels = y[unlabeled_indices]
            from collections import  Counter
            print(Counter(true_labels))
            print("size is :: "+ str(size))
            print("k is :: " + str(k))
            print("Iteration :: "+ str((k+1) /size)+ " Queried :: "+ str(total_query))
            print(classification_report(true_labels, predicted_labels))
            f_macro_result.append(f1_score(true_labels, predicted_labels, average='macro'))
            acc_result.append(accuracy_score(true_labels, predicted_labels))
            rec_macro_result.append(recall_score(true_labels, predicted_labels, average='macro'))
            f_result.append(f1_score(true_labels, predicted_labels, average='binary'))
            rec_result.append(recall_score(true_labels, predicted_labels, average='binary'))

        a, entropy = update_data(mdl, row, aq, entropy,threshold)
        if a < aq:
            aq = a
            unlabeled_indices = np.delete(unlabeled_indices, np.where(unlabeled_indices == k))
            W_Y[k] = y[k]
        if method.__name__ == 'update_threshold_adaptive_lambda':
            threshold = method(entropy, k, N, fold_count, AR)
        else:
            threshold = method()
           # print("%.5f" % threshold)
    np.savetxt("en.csv", sorted(entropy), delimiter=",", fmt='%10.5f')
    with open("f1_macro_score_" + name+ '_'+ participant_id + ".csv", 'wb') as afile:
        np.savetxt(afile, f_macro_result, delimiter=",", fmt='%10.5f')
    with open("accuracy_score_" + name + '_' + participant_id + ".csv", 'wb') as afile:
        np.savetxt(afile, acc_result, delimiter=",", fmt='%10.5f')

    with open("recall_macro_score_" + name + '_' + participant_id + ".csv", 'wb') as afile:
        np.savetxt(afile, rec_macro_result, delimiter=",", fmt='%10.5f')

    with open("f1_score_" + name+ '_'+ participant_id + ".csv", 'wb') as afile:
        np.savetxt(afile, f_result, delimiter=",", fmt='%10.5f')

    with open("recall_score_" + name + '_' + participant_id + ".csv", 'wb') as afile:
        np.savetxt(afile, rec_result, delimiter=",", fmt='%10.5f')

    predicted_labels = mdl.predict(X[unlabeled_indices])
    # predicted_labels = [1 if x > 0.5 else 0 for x in p[:, 1]]

    true_labels = y[unlabeled_indices]
    print(classification_report(true_labels, predicted_labels))

def depict_uncertainty(model):
    sensor_data = dataset.load_data()
    y_validate = sensor_data.test_target
    X_validate= sensor_data.test_data
    print(len(X_validate))
    p = model.predict_proba(X_validate)
    e =np.sort(stats.distributions.entropy(p.T))[::-1]
    print(e[0])
    for i in range(10):
        print(e[(i+1)*100])

    plt.scatter(e[:100], np.zeros_like(e[:100]))
    plt.ylabel("Uncertainty")
    plt.legend(loc=2)
    plt.show()

def customize_with_perfect_lambda(mdl, name, fold_count, participant_id, AR=15):
    f_result = []
    acc_result = []

    user_data = np.genfromtxt(join('C:/Users/Marjan/Desktop/Thesis/participants_wild', participant_id,'wrist_ss.csv'), delimiter=',')
    user_data = user_data[:, 1:]
    data = calc_features(user_data)

    X_n = preprocessing.scale(data[:, :(user_data.shape[1] - 1) * 9])
    y_n = data[:, -1]
    rng = np.random.RandomState(0)
    indices = np.arange(len(X_n))
    rng.shuffle(indices)
    X,y = X_n[indices[:]],y_n[indices[:]]
    W_Y = np.copy(y)
    N = len(X)
    unlabeled_indices = np.arange(N)
    W_Y[unlabeled_indices] = -1
    size = int(N // fold_count)
    for i in range(fold_count):
        current_X = np.array(X[i*size: (i+1)*size])
        # current_Y = W_Y[i*size: (i+1)*size]
        p = mdl.predict_proba(current_X)
        e = np.sort(stats.distributions.entropy(p.T))[::-1]

        # select up to 5 digit examples that the classifier is most uncertain about
        uncertainty_index = np.argsort(e)[::-1]

        uncertainty_index = uncertainty_index[
                                np.in1d(uncertainty_index, unlabeled_indices)]

        # keep track of indices that we get labels for
        delete_indices = np.array([])
        eating_counter = 0
        non_eating_counter = 0
        for index in uncertainty_index:
            if eating_counter + non_eating_counter < AR:
                delete_index, = np.where(unlabeled_indices == index)

                is_changed = False
                if y[delete_index] ==0 and non_eating_counter < int(0.5*AR):
                    non_eating_counter += 1
                    is_changed = True
                elif y[delete_index] == 1:
                    eating_counter += 1
                    is_changed = True
                if is_changed:
                    delete_indices = np.concatenate((delete_indices, delete_index))

        unlabeled_indices = np.delete(unlabeled_indices, delete_indices)

        W_Y = np.copy(y)
        W_Y[unlabeled_indices] = -1
        mdl = customize_model(mdl, X[:(i + 1) * size], W_Y[:(i + 1) * size])
        # p = mdl.predict_proba(X[unlabeled_indices])
        # predicted_labels = [1 if x > 0.50 else 0 for x in p[:, 1]]
        predicted_labels = mdl.predict(X[unlabeled_indices])
        true_labels = y[unlabeled_indices]
        print("Iteration :: " + str((i + 1)))
        print(classification_report(true_labels, predicted_labels))
        f_result.append(f1_score(true_labels, predicted_labels, average='binary'))
        acc_result.append(accuracy_score(true_labels, predicted_labels))
    with open("f1_score_"+ name + '_' + participant_id +  ".csv", 'wb') as afile:
        np.savetxt(afile, f_result, delimiter=",", fmt='%10.5f')

if __name__ == "__main__":
    # model= make_model()
    # customize(model)
    print("finished")
    participants = [2, 4, 5, 7]
    model =make_model()

    for p in participants:
        customize(model, update_threshold_adaptive_lambda, 'spal', 10, str(p))

    #depict_uncertainty(model)


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


#!/usr/bin/env python3

"""
Sales As Code: randomly selects a sales-y buzzword from a google sheet.
"""

import argparse
import logging
import os
import random
import sys

import gspread
from cachetools import TTLCache
from dotenv import load_dotenv
from oauth2client.service_account import ServiceAccountCredentials

logging.basicConfig(level=logging.INFO)

cache = TTLCache(maxsize=1, ttl=3600)

def authenticate():
    """Authenticate with Google Sheets API and return the gsheet object."""
    if args.credentials:
        creds = args.credentials
    else:
        load_dotenv()
        creds = os.getenv("SALES_AS_CODE_CREDS")

    if not creds:
        logging.error("No credentials found. Exiting...")
        sys.exit(1)

    credential = ServiceAccountCredentials.from_json_keyfile_name(
        creds,
        [
            "https://spreadsheets.google.com/feeds",
            "https://www.googleapis.com/auth/drive",
            "https://www.googleapis.com/auth/drive.file",
            "https://www.googleapis.com/auth/spreadsheets",
        ],
    )

    client = gspread.authorize(credential)
    gsheet = client.open("Sales As Code").sheet1

    return gsheet

def fetch_sales_terms(gsheet):
    """Fetch sales terms from Google Sheets."""
    try:
        sales_terms = gsheet.col_values(1)[1:]
        return sales_terms
    except Exception as err:
        logging.error("Failed to fetch sales terms: %s", err)
        return None

def cache_sales_terms(sales_terms):
    """Cache sales terms locally."""
    try:
        cache["sales_terms"] = sales_terms
    except Exception as err:
        logging.error("Failed to cache sales terms: %s", err)

def get_sales_terms():
    """Get sales terms from Google Sheets, cache them locally, and pick one at random."""
    cached_terms = cache.get("sales_terms")
    if cached_terms:
        sales_terms = cached_terms
    else:
        gsheet = authenticate()
        if not gsheet:
            logging.error("Failed to authenticate with Google Sheets. Exiting...")
            return None
        sales_terms = fetch_sales_terms(gsheet)
        if not sales_terms:
            return None
        cache_sales_terms(sales_terms)
    return random.choice(sales_terms)

def spongecase(term):
    """SpOngECAse tHe oUtPUt."""
    return "".join(
        char.upper() if random.randint(0, 1) else char.lower() for char in term
    )

def ask_if_rep_talking():
    """A while loop asking if the rep is still talking."""
    while True:
        try:
            answer = input("Is the rep still talking? (y/n) ").lower()
            if answer != "y":
                break
            sales_term = get_sales_terms()
            if args.spongecase:
                print(spongecase(sales_term))
            else:
                print(f"{sales_term}")
        except Exception as err:
            logging.error("Failed to get sales term: %s", err)

if __name__ == "__main__":
    argparser = argparse.ArgumentParser()
    argparser.add_argument(
        "-s",
        "--spongecase",
        required=False,
        action="store_true",
        help="SpONgECasE ThE oUtPut",
    )
    argparser.add_argument(
        "-c",
        "--credentials",
        required=False,
        help="Path to credentials file",
    )
    argparser.add_argument(
        "--cachebust",
        required=False,
        action="store_true",
        help="Bust the cache",
    )
    args = argparser.parse_args()

    if not authenticate():
        logging.error("Failed to authenticate with Google Sheets. Exiting...")
        sys.exit(1)

    ask_if_rep_talking()


#!/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 rclpy
from rclpy.node import Node

class Talk(Node):
    def __init__(self,name):
        super().__init__(name)
        self.get_logger().info("Hello, I'm %s" % name)

def main(args=None):
    rclpy.init(args=args)
    node = Talk("Tom")
    rclpy.spin(node)
    rclpy.shutdown()