#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Python Flask Server to automate Jupyter Lab port forward on HPC cluster.
Start the server with the 'run()' method

(C) 2020 George Mihaila
email georgemihaila@my.unt.edu
"""

import os
from datetime import datetime
import pexpect

from jupyter_lab import app, DATABASE_NAME, START_OPEN_PORT, END_OPEN_PORT
from helper_functions import logger, find_free_port
from sqlite_database import add_db

class JupyterLab(object):
    """
    Args:
        user: User id.
        passw: Password for login to HPC.
        addrs: Hostname of HPC.
        timeout: Seconds until end session.
    Return:
        is_user: True / Flase if user exists on HPC.
        running_jupyter: True / Flase if running jupyter instance or not.
        jupyter_port: Jupyter port, either last running instance or free port.
    """

    def __init__(self, user, credential, hostname, python_path, jupyter_bin_path, pid, session_length):
        self.user = user
        self.credential = credential
        self.hostname = hostname
        self.python_path = python_path
        self.jupyter_bin_path = jupyter_bin_path
        self.pid = pid
        self.session_length = session_length
        self.timeout = 5
        return

    def configure(self,):
        """Jupyter configure and checking running instances.
        This function executes following procedures:
            - Check if user exists.
            Change path to scratch2.
            Export enviroment variables (python and conda paths).
            Run jupyter config command.
        """

        logger(user=self.user, message='CHECK IF ANY RUNNING JUPYTER and HASH AND JUPYTER CONFIG', level='INFO')
        # UPDATE USER IN DATABASE
        add_db(db_name=DATABASE_NAME,
               user=self.user,
               last_login=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
               local_port=0,
               talon_port=0,
               login_node=self.hostname,
               pid_session=self.pid,
               state_session='initiated')

        child = pexpect.spawn('ssh %s@%s -o StrictHostKeyChecking=no' %
                              (self.user, self.hostname), encoding='utf-8', timeout=self.timeout, logfile=None)
        child.expect(['password: '])
        child.sendline(self.credential)

        shell_jupyter_config_path = "python -c \"import os; print('jupyter_config: ',os.path.exists('/storage/scratch2/%s/.jupyter/jupyter_notebook_config.json'))\""%(self.user)
        shell_free_port = "python -c 'import socket; s=socket.socket(); s.bind((\"\", 0)); print(\"free port:\",s.getsockname()[1])'"
        # STAUS VARIABLES
        jupyter_port = None

        # STATE BOOLEAN VARIABLES
        is_jupyter_config = False
        is_user = False
        checking_running_jupyter = False
        running_jupyter = False
        configuring_jupyter = False

        # JUPYTER CONFIG FILE
        allow_remote_access = "sed -i 's/#c.NotebookApp.allow_remote_access = False/c.NotebookApp.allow_remote_access = True/g' .jupyter/jupyter_notebook_config.py"
        cull_busy = "sed -i 's/#c.MappingKernelManager.cull_busy = False/c.MappingKernelManager.cull_busy = False/g' .jupyter/jupyter_notebook_config.py"
        cull_connected = "sed -i 's/#c.MappingKernelManager.cull_connected = False/c.MappingKernelManager.cull_connected = True/g' .jupyter/jupyter_notebook_config.py"
        cull_idle_timeout = "sed -i 's/#c.MappingKernelManager.cull_idle_timeout = 0/c.MappingKernelManager.cull_idle_timeout = %s/g' .jupyter/jupyter_notebook_config.py"%self.session_length
        cull_interval = "sed -i 's/#c.MappingKernelManager.cull_interval = 300/c.MappingKernelManager.cull_interval = 5/g' .jupyter/jupyter_notebook_config.py"
        kernel_info_timeout = "sed -i 's/#c.MappingKernelManager.kernel_info_timeout = 60/c.MappingKernelManager.kernel_info_timeout = 60/g' .jupyter/jupyter_notebook_config.py"
        shutdown_no_activity_timeout = "sed -i 's/#c.NotebookApp.shutdown_no_activity_timeout = 0/c.NotebookApp.shutdown_no_activity_timeout = %s/g' .jupyter/jupyter_notebook_config.py"%self.session_length
        notebook_dir = "sed -i \"s/#c.NotebookApp.notebook_dir = ''/c.NotebookApp.notebook_dir = '\\/storage\\/scratch2\\/%s\\/'/g\" .jupyter/jupyter_notebook_config.py"%(self.user)
        map_root_dir = "sed -i \"s/#c.MappingKernelManager.root_dir = ''/c.MappingKernelManager.root_dir = '\\/storage\\/scratch2\\/%s'/g\" .jupyter/jupyter_notebook_config.py"%(self.user)
        content_root_dir = "sed -i \"s/#c.ContentsManager.root_dir = '\\/'/c.ContentsManager.root_dir = '\\/storage\\/scratch2\\/%s'/g\" .jupyter/jupyter_notebook_config.py"%(self.user)
        iopub_data_rate_limit = "sed -i 's/#c.NotebookApp.iopub_data_rate_limit = 1000000/c.NotebookApp.iopub_data_rate_limit = 1e10/g' .jupyter/jupyter_notebook_config.py"
        mathjax_config = "sed -i \"s/#c.NotebookApp.mathjax_config = 'TeX-AMS-MML_HTMLorMML-full,Safe'/c.NotebookApp.mathjax_config = 'TeX-AMS-MML_HTMLorMML-full,Safe'/g\" .jupyter/jupyter_notebook_config.py"

        # LOOP CHECK SHELL
        while True:
            try:
                child.expect('\n')
                out_line = child.before
                # ONLY LOG
                logger(user=self.user, message=str(out_line),
                       level='INFO', verbose=False)

                # CHECK IF USER EXISTS
                if "Last login:" in out_line:
                    # USER EXISTS
                    is_user = True
                    logger(user=self.user, message='USER EXISTS!', level='INFO')
                    # CHANGE PATH
                    child.sendline("cd /storage/scratch2/%s"%(self.user))
                    # DEACTIVATE ANY VIRTUALENV
                    child.sendline('deactivate')
                    child.sendline('conda deactivate')
                    # ADD ENVIROMENT VARIABLES
                    child.sendline('export PATH="$PATH:%s/bin"'%(self.python_path))
                    child.sendline('export PATH="$PATH:%s/lib"'%(self.python_path))
                    child.sendline("export IPYTHONDIR=/storage/scratch2/%s/.ipython"%(self.user))
                    child.sendline("export JUPYTER_CONFIG_DIR=/storage/scratch2/%s/.jupyter"%(self.user))
                    child.sendline("export JUPYTER_DATA_DIR=/storage/scratch2/%s/.jupyter"%(self.user))
                    # JUPYTER GENERATE CONFIG .jupyter
                    child.sendline("%s notebook --generate-config -y"%(self.jupyter_bin_path))
                    # CREATE .ipyton if not created
                    child.sendline('mkdir -p .ipython')
                    # CREATE SYMBOLIK LINK TO HOME DIRECTORY IF NOT EXISTING
                    child.sendline('test ! -d /storage/scratch2/%s_home_dir && ln -s /home/%s/ %s_home_dir'%(self.user, self.user, self.user))
                    logger(user=self.user, message='TRY TO CREATE jupyter_notebook_config.py', level='INFO')
                    # CLEAN BASH HISTORY
                    child.sendline('history -c')
                    # CHECK IF jupyter_notebook_config.json CREATED with
                    child.sendline(shell_jupyter_config_path)

                # CHECK IF JUPYTER CONFIG EXISTS. KNOW THAT USER EXISTS
                if ('Writing default config to:' in out_line) and is_user:
                    # SET PASSWORD
                    # CHANGE jupyter_notebook_config.py
                    is_jupyter_config = True
                    logger(user=self.user, message='CREATED jupyter_notebook_config.py!', level='CRITICAL')
                    # CREATE JUPYTER PASSWORD
                    child.sendline("%s notebook password"%(self.jupyter_bin_path))
                    child.expect("Enter password:")
                    child.sendline(self.credential)
                    child.expect("Verify password:")
                    child.sendline(self.credential)
                    # CLEAN BASH HISTORY
                    child.sendline('history -c')

                # CHANGE jupyter_notebook_config.py. KNOW THAT JUPYTER PASSWORD IS SET.
                if ('Wrote hashed password to' in out_line) and is_jupyter_config:
                    logger(user=self.user, message='JUPYTER PASSWORD IS SET!', level='CRITICAL')
                    # ALLOW REMOTE ACCESS
                    child.sendline(allow_remote_access)
                    # CONFIG FILE jupyter_config.py
                    child.sendline(cull_busy)
                    child.sendline(cull_connected)
                    child.sendline(cull_idle_timeout)
                    child.sendline(cull_interval)
                    child.sendline(kernel_info_timeout)
                    child.sendline(shutdown_no_activity_timeout)
                    child.sendline(notebook_dir)
                    child.sendline(map_root_dir)
                    child.sendline(content_root_dir)
                    child.sendline(iopub_data_rate_limit)
                    child.sendline(mathjax_config)
                    logger(user=self.user, message='CONFIGURED jupyter_notebook_config.py!', level='CRITICAL')
                    # CLEAN BASH HISTORY
                    child.sendline('history -c')
                    # GRAB UNUSED PORT ON HPC
                    child.sendline('history -c && %s'%(shell_free_port))

                # GRAB FREE PORT [IN CASE OF NO RUNNING JUPYTER]
                if ("('free port:'," in out_line):
                    jupyter_port = (out_line.split(",")[1].split(")")[0])
                    logger(user=self.user, message='FOUND FREE PORT: %s'%jupyter_port,level='CRITICAL')
                    # GRAB LAST PORT OF RUNNING NOTEBOOK IF EXISTS
                    child.sendline('history -c && %s notebook list'%(self.jupyter_bin_path))

                # SEE RUNNING JUPYTER PORTS
                if ('Currently running servers:' in out_line):
                    # LOOKING FOR RUNNING JUPYTER INSTANCES
                    logger(user=self.user, message='LOOKING FOR RUNNING SERVERS',level='CRITICAL')
                    checking_running_jupyter = True

                # GRAB LAST RUNNING JUPYTER. KNOW THAT USERS EXISTS, COFIGURE SET AND LOOKING FOR JUPYTER
                if ("http://" in out_line) and checking_running_jupyter:
                    # GRAB LAST RUNNING JUPYTER PORT
                    logger(user=self.user, message='GRAB LAST RUNNING JUPYTER NOTEBOOK PORT', level='INFO')
                    jupyter_port = (out_line.split(":")[2].split("/")[0])
                    # FOUND RUNNING JUPYTER
                    running_jupyter = True
                    # STOP LOOKING FOR RUNNING JUPYTER NOTEBOOKS SERVERS
                    checking_running_jupyter = False
                    logger(user=self.user, message='jupyter_port %s' %
                           jupyter_port, level='CRITICAL')
                    # CLEAN BASH HISTORY
                    child.sendline('history -c')
                    child.close(force=True)
                    return is_user, running_jupyter, jupyter_port

                # CHECK IF JUPYTER CONFIG NOT EXISTS. KNOW THAT USER EXISTS.
                if ("('jupyter_config: ', False)" in out_line) and is_user:
                    # JUPYTER CONFIG NOT EXISTS AND USER EXISTS
                    logger(user=self.user, message='JUPYTER FAILED TO CREATE CONFIG FILE!',level='ERROR')
                    logger(user=self.user, message='FUNCTION ENDED!', level='ERROR')
                    # CLEAN BASH HISTORY
                    child.sendline('history -c')
                    child.close(force=True)
                    break

            except:
                logger(user=self.user, message='jupyter_port %s' %
                       jupyter_port, level='CRITICAL')
                logger(user=self.user, message="FUNCTION 'configure()' ENDED!", level='WARNING')
                # CLEAN BASH HISTORY
                child.sendline('history -c')
                child.close(force=True)
                break

        if jupyter_port is None:
            logger(user=self.user, message='NO RUNNING JUPYTER NOTEBOOK', level='WARNING')
        return is_user, running_jupyter, jupyter_port

    def forward(self, local_port, remote_port, running_instance):
        """Forward port of running instance from HPC to local VM.
        Start new Jupyter instance if nothing running.
        WITH A RUNNING JUPYTER INSTANCE ON A KNOWN PORT:
            - FORWARD KNOWN PORT TO GATEWAY
            - LOGIN IS SUCCESSFULL
            - LOGIN FAILED
        Update users state to 'initialized' before running this function.

        Args:
            remote_port: Port used on HPC.
            running_instance: True / False. If just forward existing running jupyter or start new jupyter.
        """
        try:
            remote_port = int(remote_port)
        except:
            logger(user=self.user, message='INVALID REMOTE PORT!', level='ERROR')
            logger(user=self.user, message="FUNCTION 'forward()' ENDED!", level='ERROR')
            # UPDATE USER IN DATABASE
            add_db(db_name=DATABASE_NAME,
                   user=self.user,
                   last_login=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                   local_port=0,
                   talon_port=0,
                   login_node=self.hostname,
                   pid_session=self.pid,
                   state_session='ended')
            return

        logger(user=self.user, message='FREE LOCAL PORT: %s'%local_port, level='INFO')
        child = pexpect.spawn('ssh -L 0.0.0.0:%s:127.0.0.1:%s  %s@%s -o StrictHostKeyChecking=no' %
                              (local_port, remote_port, self.user, self.hostname), encoding='utf-8', timeout=self.session_length, logfile=None)
        child.expect(['password: '])
        child.sendline(self.credential)
        # login staus variables
        first_line = False
        is_logged = False
        # LOOP CHECK SHELL
        while True:
            try:
                child.expect('\n')
                out_line = child.before
                logger(user=self.user, message=str(out_line), level='INFO', verbose=False)

                # FIRST LINE OF THE LOGIN
                if (" " in out_line) and not first_line:
                    first_line = True

                # LOGIN IS SUCCESSFULL
                if ("Last login:" in out_line) and first_line:
                    is_logged = True
                    logger(user=self.user, message='LOGIN TO HPC SUCCESSFULL! STARTING JUPYTER...', level='CRITICAL')
                    # CHANGE PATH
                    child.sendline("cd /storage/scratch2/%s"%(self.user))
                    # DEACTIVATE ANY VIRTUALENV
                    child.sendline('deactivate')
                    child.sendline('conda deactivate')
                    # ADD ENVIROMENT VARIABLES
                    child.sendline('export PATH="$PATH:%s/bin"'%self.python_path)
                    child.sendline('export PATH="$PATH:%s/lib"'%(self.python_path))
                    child.sendline("export IPYTHONDIR=/storage/scratch2/%s/.ipython"%(self.user))
                    child.sendline("export JUPYTER_CONFIG_DIR=/storage/scratch2/%s/.jupyter"%(self.user))
                    child.sendline("export JUPYTER_DATA_DIR=/storage/scratch2/%s/.jupyter"%(self.user))
                    # CLEAN BASH HISTORY
                    child.sendline('history -c')
                    if running_instance is False:
                        # NEED NEW JUPYTER INSTANCE [MAKE SURE TO KEEP WATCHING OUTPUT]
                        child.sendline("nohup %s lab --no-browser --ip=0.0.0.0 --port=%s &> .jupyter_lab.log & tail -f .jupyter_lab.log"%(self.jupyter_bin_path, remote_port))

                # JUPYTER PORT FORWARDED
                if (running_instance is True) and (is_logged is True):
                    logger(user=self.user, message='JUPYTER FORWARDED. LOCAL PORT: %s REMOTE PORT: %s'%
                            (local_port, remote_port), level='CRITICAL')
                    # UPDATE USER IN DATABASE
                    add_db(db_name=DATABASE_NAME,
                           user=self.user,
                           last_login=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                           local_port=local_port,
                           talon_port=remote_port,
                           login_node=self.hostname,
                           pid_session=self.pid,
                           state_session='running')

                # JUPYTER STARTED SUCCESSFULLY
                if ('The Jupyter Notebook is running at' in out_line) and (is_logged is True):
                    logger(user=self.user, message='JUPYTER STARTED SUCCESSFULLY. LOCAL PORT: %s REMOTE PORT: %s'%
                            (local_port, remote_port), level='CRITICAL')
                    # UPDATE USER IN DATABASE
                    add_db(db_name=DATABASE_NAME,
                           user=self.user,
                           last_login=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                           local_port=local_port,
                           talon_port=remote_port,
                           login_node=self.hostname,
                           pid_session=self.pid,
                           state_session='running')

                # LOGIN FAILED
                if ("Last login:" not in out_line) and first_line and not is_logged:
                    logger(user=self.user, message='LOGIN FAILED!', level='ERROR')
                    # UPDATE USER IN DATABASE
                    add_db(db_name=DATABASE_NAME,
                           user=self.user,
                           last_login=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                           local_port=0,
                           talon_port=0,
                           login_node=0,
                           pid_session=0,
                           state_session='ended')
                    # CLEAN BASH HISTORY
                    child.sendline('history -c')
                    child.close(force=True)
                    break

            except:
                logger(user=self.user, message='forward_port ENDED', level='WARNING')
                # UPDATE USER IN DATABASE
                add_db(db_name=DATABASE_NAME,
                       user=self.user,
                       last_login=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                       local_port=0,
                       talon_port=0,
                       login_node=0,
                       pid_session=0,
                       state_session='ended')
                # CLEAN BASH HISTORY
                child.sendline('history -c')
                child.close(force=True)
                break
        return

def jupyter_run(conn, user, credential, hostname, python_path, jupyter_bin_path, session_length):
    """Functions wrapper for python multiprocessing.

    Args:
        conn: Multi-process connection.
        _euid: User id.
        _pass: Password for login to HPC.
        hostname: Hostname of HPC.
        session_timeout: Seconds until end session.
    """

    try:
        session_length = int(session_length)
    except:
        # UPDATE USER IN DATABASE
        add_db(db_name=DATABASE_NAME,
               user=user,
               last_login=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
               local_port=0,
               talon_port=0,
               login_node=hostname,
               pid_session=0,
               state_session='ended')
        # SEND MESSAGE TO MASTER PROCESS
        conn.send('ended 0')
        conn.close()
        return

    pid = os.getpid()
    # CREATE INSTANCE
    jupyter_instance = JupyterLab(user, credential, hostname, python_path, jupyter_bin_path, pid, session_length)
    # CHECK CONFIGURATION
    user_exists, running_jupyter, jupyter_port = jupyter_instance.configure()

    # GRAB FREE LOCAL PORT
    free_port = find_free_port(min_port=START_OPEN_PORT, max_port=END_OPEN_PORT)

    if user_exists is True:
        if jupyter_port is not None:
            # SEND MESSAGE TO MASTER PROCESS
            conn.send('running %s'%free_port)
            conn.close()
            # FORWARD JUPYTER INSTANCE
            jupyter_instance.forward(local_port=free_port, remote_port=jupyter_port, running_instance=running_jupyter)
        else:
            logger(user=user, message='USER %s WAS NOT ABLE TO GET JUPYTER PORT [MAYBE NOT RUNNING OR FREE PORT]!'%user, level='ERROR')
            # UPDATE USER IN DATABASE
            add_db(db_name=DATABASE_NAME,
                   user=user,
                   last_login=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                   local_port=0,
                   talon_port=0,
                   login_node=hostname,
                   pid_session=0,
                   state_session='ended')
            # SEND MESSAGE TO MASTER PROCESS
            conn.send('ended None')
            conn.close()

    else:
        logger(user=user, message='USER %s DOES NOT EXIST!'%user, level='ERROR')
        # UPDATE USER IN DATABASE
        add_db(db_name=DATABASE_NAME,
               user=user,
               last_login=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
               local_port=0,
               talon_port=0,
               login_node=hostname,
               pid_session=0,
               state_session='ended')
        # SEND MESSAGE TO MASTER PROCESS
        conn.send('ended %s'%free_port)
        conn.close()
    return


# fastapi server starts throw "uvicorn ai:app --port=8088"
from detecto.utils import read_image
from detecto.core import Model
from matplotlib import pyplot as plt
import base64
import io
import matplotlib.patches as patches
import torch
from detecto.utils import reverse_normalize, _is_iterable
from torchvision import transforms
from fastapi import FastAPI
from fastapi import HTTPException

def byte_image(image, boxes, labels=None):
    fig, ax = plt.subplots(1)
    # If the image is already a tensor, convert it back to a PILImage
    # and reverse normalize it
    if isinstance(image, torch.Tensor):
        image = reverse_normalize(image)
        image = transforms.ToPILImage()(image)
    ax.imshow(image)
    # Show a single box or multiple if provided
    if boxes.ndim == 1:
        boxes = boxes.view(1, 4)
    if labels is not None and not _is_iterable(labels):
        labels = [labels]
    # Plot each box
    for i in range(boxes.shape[0]):
        box = boxes[i]
        width, height = (box[2] - box[0]).item(), (box[3] - box[1]).item()
        initial_pos = (box[0].item(), box[1].item())
        rect = patches.Rectangle(initial_pos,  width, height, linewidth=1,
                                 edgecolor='r', facecolor='none')
        if labels:
            ax.text(box[0] + 5, box[1] - 5, '{}'.format(labels[i]), color='red')
        ax.add_patch(rect)
    for item in [fig, ax]:
        item.patch.set_visible(False)
    plt.xticks(())
    plt.yticks(())
    my_stringIObytes = io.BytesIO()
    plt.savefig(my_stringIObytes, format='jpg')
    my_stringIObytes.seek(0)
    my_base64_jpgData = base64.b64encode(my_stringIObytes.read())
    return  my_base64_jpgData

app = FastAPI(title="AI Model")
labels = ['teeth', 'caries']
model = Model.load('trained_model.pth', labels)

@app.get("/get", response_description="Get predict_photo (base64), caries_count, teeth_count "
                                        "by file_path (absolute path to input photo)")
async def get(file_path: str):
    try:
        _image = read_image(file_path)
        _labels, _boxes, scores = model.predict(_image)
        caries_count, teeth_count = _labels.count('caries'), _labels.count('teeth')
        image = byte_image(_image, _boxes, _labels)
        result_of_classification = {"predict_photo": image,
                                    "caries_count": caries_count,
                                    "teeth_count": teeth_count}
    except Exception:
        raise HTTPException(404, str(Exception))
    return result_of_classification


from __future__ import absolute_import, division, print_function, unicode_literals

# THIS MODULE HAS TO WORK WITH PYTHON 2.7!

import argparse
import errno
import logging
import multiprocessing
import os
import resource
import signal
import subprocess
import sys
import threading
import time
import tempfile
sys.dont_write_bytecode = True # prevent creation of .pyc files

from benchexec import baseexecutor
from benchexec.baseexecutor import BaseExecutor
from benchexec.cgroups import *
from benchexec import oomhandler
from benchexec import systeminfo
from benchexec import util

_WALLTIME_LIMIT_DEFAULT_OVERHEAD = 30 # seconds more than cputime limit
_ULIMIT_DEFAULT_OVERHEAD = 30 # seconds after cgroups cputime limit
_BYTE_FACTOR = 1000 # byte in kilobyte
_LOG_SHRINK_MARKER = "\n\n\nWARNING: YOUR LOGFILE WAS TOO LONG, SOME LINES IN THE MIDDLE WERE REMOVED.\n\n\n\n"
_SUDO_ARGS = ['sudo', '--non-interactive', '-u']

try:
    from subprocess import DEVNULL
except ImportError:
    DEVNULL = open(os.devnull, 'rb')

def main(argv=None):
    """
    A simple command-line interface for the runexecutor module of BenchExec.
    """
    if argv is None:
        argv = sys.argv

    # parse options
    parser = argparse.ArgumentParser(
        fromfile_prefix_chars='@',
        description=
        """Execute a command with resource limits and measurements.
           Command-line parameters can additionally be read from a file if file name prefixed with '@' is given as argument.
           Part of BenchExec: https://github.com/sosy-lab/benchexec/""")

    resource_args = parser.add_argument_group("optional arguments for resource limits")
    resource_args.add_argument("--memlimit", type=util.parse_memory_value, metavar="BYTES",
        help="memory limit in bytes")
    resource_args.add_argument("--timelimit", type=util.parse_timespan_value, metavar="SECONDS",
        help="CPU time limit in seconds")
    resource_args.add_argument("--softtimelimit", type=util.parse_timespan_value, metavar="SECONDS",
        help='"soft" CPU time limit in seconds (command will be send the TERM signal at this time)')
    resource_args.add_argument("--walltimelimit", type=util.parse_timespan_value, metavar="SECONDS",
        help='wall time limit in seconds (default is CPU time limit plus a few seconds)')
    resource_args.add_argument("--cores", type=util.parse_int_list, metavar="N,M-K",
        help="list of CPU cores to use")
    resource_args.add_argument("--memoryNodes", type=util.parse_int_list, metavar="N,M-K",
        help="list of memory nodes to use")

    io_args = parser.add_argument_group("optional arguments for run I/O")
    io_args.add_argument("--input", metavar="FILE",
        help="name of file used as stdin for command "
            "(default: /dev/null; use - for stdin passthrough)")
    io_args.add_argument("--output", default="output.log", metavar="FILE",
        help="name of file where command output is written")
    io_args.add_argument("--maxOutputSize", type=util.parse_memory_value, metavar="BYTES",
        help="shrink output file to approximately this size if necessary "
            "(by removing lines from the middle of the output)")
    io_args.add_argument("--skip-cleanup", action="store_false", dest="cleanup",
        help="do not delete files created by the tool in temp directory")

    environment_args = parser.add_argument_group("optional arguments for run environment")
    environment_args.add_argument("--require-cgroup-subsystem", action="append", default=[], metavar="SUBSYSTEM",
        help="additional cgroup system that should be enabled for runs "
            "(may be specified multiple times)")
    environment_args.add_argument("--set-cgroup-value", action="append", dest="cgroup_values", default=[],
        metavar="SUBSYSTEM.OPTION=VALUE",
        help="additional cgroup values that should be set for runs (e.g., 'cpu.shares=1000')")
    environment_args.add_argument("--dir", metavar="DIR",
        help="working directory for executing the command (default is current directory)")
    environment_args.add_argument("--user", metavar="USER",
        help="execute tool under given user account (needs password-less sudo setup)")

    baseexecutor.add_basic_executor_options(parser)

    options = parser.parse_args(argv[1:])
    baseexecutor.handle_basic_executor_options(options)

    # For integrating into some benchmarking frameworks,
    # there is a DEPRECATED special mode
    # where the first and only command-line argument is a serialized dict
    # with additional options
    env = {}
    if len(options.args) == 1 and options.args[0].startswith("{"):
        data = eval(options.args[0])
        options.args = data["args"]
        env = data.get("env", {})
        options.debug = data.get("debug", options.debug)
        if "maxLogfileSize" in data:
            try:
                options.maxOutputSize = int(data["maxLogfileSize"]) * _BYTE_FACTOR * _BYTE_FACTOR # MB to bytes
            except ValueError:
                options.maxOutputSize = util.parse_memory_value(data["maxLogfileSize"])

    if options.input == '-':
        stdin = sys.stdin
    elif options.input is not None:
        if options.input == options.output:
            sys.exit("Input and output files cannot be the same.")
        try:
            stdin = open(options.input, 'rt')
        except IOError as e:
            sys.exit(e)
    else:
        stdin = None

    cgroup_subsystems = set(options.require_cgroup_subsystem)
    cgroup_values = {}
    for arg in options.cgroup_values:
        try:
            key, value = arg.split("=", 1)
            subsystem, option = key.split(".", 1)
            if not subsystem or not option:
                raise ValueError()
        except ValueError:
            sys.exit('Cgroup value "{}" has invalid format, needs to be "subsystem.option=value".'
                     .format(arg))
        cgroup_values[(subsystem, option)] = value
        cgroup_subsystems.add(subsystem)

    executor = RunExecutor(user=options.user, cleanup_temp_dir=options.cleanup,
                           additional_cgroup_subsystems=list(cgroup_subsystems))

    # ensure that process gets killed on interrupt/kill signal
    def signal_handler_kill(signum, frame):
        executor.stop()
    signal.signal(signal.SIGTERM, signal_handler_kill)
    signal.signal(signal.SIGINT,  signal_handler_kill)

    formatted_args = " ".join(map(util.escape_string_shell, options.args))
    logging.info('Starting command %s', formatted_args)
    logging.info('Writing output to %s', options.output)

    # actual run execution
    try:
        result = executor.execute_run(
                            args=options.args,
                            output_filename=options.output,
                            stdin=stdin,
                            hardtimelimit=options.timelimit,
                            softtimelimit=options.softtimelimit,
                            walltimelimit=options.walltimelimit,
                            cores=options.cores,
                            memlimit=options.memlimit,
                            memory_nodes=options.memoryNodes,
                            cgroupValues=cgroup_values,
                            environments=env,
                            workingDir=options.dir,
                            maxLogfileSize=options.maxOutputSize)
    finally:
        if stdin:
            stdin.close()

    executor.check_for_new_files_in_home()

    # exit_code is a special number:
    exit_code = util.ProcessExitCode.from_raw(result['exitcode'])

    def print_optional_result(key):
        if key in result:
            # avoid unicode literals such that the string can be parsed by Python 3.2
            print(key + "=" + str(result[key]).replace("'u", ''))

    # output results
    print_optional_result('terminationreason')
    print("exitcode=" + str(exit_code.raw))
    if exit_code.value is not None:
        print("returnvalue=" + str(exit_code.value))
    if exit_code.signal is not None:
        print("exitsignal=" + str(exit_code.signal))
    print("walltime=" + str(result['walltime']) + "s")
    print("cputime=" + str(result['cputime']) + "s")
    for key in sorted(result.keys()):
        if key.startswith('cputime-'):
            print("{}={:.9f}s".format(key, result[key]))
    print_optional_result('memory')
    if 'energy' in result:
        for key, value in result['energy'].items():
            print("energy-{0}={1}".format(key, value))

class RunExecutor(BaseExecutor):

    # --- object initialization ---

    def __init__(self, user=None, cleanup_temp_dir=True, additional_cgroup_subsystems=[],
                 *args, **kwargs):
        """
        Create an instance of of RunExecutor.
        @param user None or an OS user as which the benchmarked process should be executed (via sudo).
        @param cleanup_temp_dir Whether to remove the temporary directories created for the run.
        @param additional_cgroup_subsystems List of additional cgroup subsystems that should be required and used for runs.
        """
        super(RunExecutor, self).__init__(*args, **kwargs)
        self._termination_reason = None
        self._user = user
        self._should_cleanup_temp_dir = cleanup_temp_dir
        self._cgroup_subsystems = additional_cgroup_subsystems

        if user is not None:
            # Check if we are allowed to execute 'kill' with dummy signal.
            sudo_check = self._build_cmdline(['kill', '-0', '0'])
            logging.debug('Checking for capability to run with sudo as user %s.', user)
            p = subprocess.Popen(sudo_check, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            try:
                if p.wait():
                    logging.error(
                        'Calling "%s" failed with error code %s and the following output: %s\n%s',
                        ' '.join(sudo_check),
                        p.returncode,
                        p.stdout.read().decode().strip(),
                        p.stderr.read().decode().strip())
                    sys.exit('Cannot execute benchmark as user "{0}", please fix your sudo setup.'.format(user))
            finally:
                p.stdout.close()
                p.stderr.close()

            # Check home directory of user
            try:
                self._home_dir = _get_user_account_info(user).pw_dir
            except (KeyError, ValueError) as e:
                sys.exit('Unknown user {}: {}'.format(user, e))
            try:
                self._home_dir_content = set(self._listdir(self._home_dir))
            except (subprocess.CalledProcessError, IOError):
                # Probably directory does not exist
                self._home_dir_content = []
            if self._home_dir_content:
                logging.warning(
                    'Home directory %s of user %s contains files and/or directories, it is '
                    'recommended to do benchmarks with empty home to prevent undesired influences.',
                    self._home_dir, user)

        self._init_cgroups()

    def _init_cgroups(self):
        """
        This function initializes the cgroups for the limitations and measurements.
        """
        self.cgroups = find_my_cgroups()

        for subsystem in self._cgroup_subsystems:
            self.cgroups.require_subsystem(subsystem)
            if subsystem not in self.cgroups:
                sys.exit('Required cgroup subsystem "{}" is missing.'.format(subsystem))

        self.cgroups.require_subsystem(CPUACCT)
        if CPUACCT not in self.cgroups:
            logging.warning('Without cpuacct cgroups, cputime measurement and limit '
                            'might not work correctly if subprocesses are started.')

        self.cgroups.require_subsystem(FREEZER)
        if FREEZER not in self.cgroups:
            if self._user is not None:
                # In sudo mode, we absolutely need at least one cgroup subsystem
                # to be able to find the process where we need to send signals to
                sys.exit('Cannot reliably kill sub-processes without freezer cgroup,'
                         + ' this is necessary if --user is specified.'
                         + ' Please enable this cgroup or do not specify --user.')
            else:
                logging.warning('Cannot reliably kill sub-processes without freezer cgroup.')

        self.cgroups.require_subsystem(MEMORY)
        if MEMORY not in self.cgroups:
            logging.warning('Cannot measure memory consumption without memory cgroup.')
        else:
            if systeminfo.has_swap() and (
                    not self.cgroups.has_value(MEMORY, 'memsw.max_usage_in_bytes')):
                logging.warning(
                    'Kernel misses feature for accounting swap memory, but machine has swap. '
                    'Memory usage may be measured inaccurately. '
                    'Please set swapaccount=1 on your kernel command line or disable swap with '
                    '"sudo swapoff -a".')

        self.cgroups.require_subsystem(CPUSET)
        self.cpus = None # to indicate that we cannot limit cores
        self.memory_nodes = None # to indicate that we cannot limit cores
        if CPUSET in self.cgroups:
            # Read available cpus/memory nodes:
            try:
                self.cpus = util.parse_int_list(self.cgroups.get_value(CPUSET, 'cpus'))
            except ValueError as e:
                logging.warning("Could not read available CPU cores from kernel: %s", e.strerror)
            logging.debug("List of available CPU cores is %s.", self.cpus)

            try:
                self.memory_nodes = util.parse_int_list(self.cgroups.get_value(CPUSET, 'mems'))
            except ValueError as e:
                logging.warning("Could not read available memory nodes from kernel: %s",
                                e.strerror)
            logging.debug("List of available memory nodes is %s.", self.memory_nodes)

    # --- utility functions ---

    def _build_cmdline(self, args, env={}):
        """
        Build the final command line for executing the given command,
        using sudo if necessary.
        """
        if self._user is None:
            return args
        result = _SUDO_ARGS + [self._user]
        for var, value in env.items():
            result.append(var + '=' + value)
        return result + ['--'] + args

    def _kill_process(self, pid, cgroups=None, sig=signal.SIGKILL):
        """
        Try to send signal to given process, either directly of with sudo.
        Because we cannot send signals to the sudo process itself,
        this method checks whether the target is the sudo process
        and redirects the signal to sudo's child in this case.
        """
        if self._user is not None:
            if not cgroups:
                cgroups = find_cgroups_of_process(pid)
            # In case we started a tool with sudo, we cannot kill the started
            # process itself, because sudo always runs as root.
            # So if we are asked to kill the started process itself (the first
            # process in the cgroup), we instead kill the child of sudo
            # (the second process in the cgroup).
            pids = cgroups.get_all_tasks(FREEZER)
            try:
                if pid == next(pids):
                    pid = next(pids)
            except StopIteration:
                # pids seems to not have enough values
                pass
            finally:
                pids.close()
        self._kill_process0(pid, sig)

    def _kill_process0(self, pid, sig=signal.SIGKILL):
        """
        Send signal to given process, either directly or with sudo.
        If the target is the sudo process itself, the signal will be lost,
        because we do not have the rights to send signals to sudo.
        Use _kill_process() because of this.
        """
        if self._user is None:
            super(RunExecutor, self)._kill_process(pid, sig)
        else:
            logging.debug('Sending signal %s to %s with sudo.', sig, pid)
            try:
                # Cast sig to int, under Python 3.5 the signal.SIG* constants are nums, not ints.
                subprocess.check_call(args=self._build_cmdline(['kill', '-'+str(int(sig)), str(pid)]))
            except subprocess.CalledProcessError as e:
                # may happen for example if process no longer exists
                logging.debug(e)

    def _listdir(self, path):
        """Return the list of files in a directory, assuming that our user can read it."""
        if self._user is None:
            return os.listdir(path)
        else:
            args = self._build_cmdline(['/bin/ls', '-1', path])
            return subprocess.check_output(args, stderr=DEVNULL).decode('utf-8', errors='ignore').split('\n')

    def _set_termination_reason(self, reason):
        if not self._termination_reason:
            self._termination_reason = reason

    # --- setup and cleanup for a single run ---

    def _setup_cgroups(self, my_cpus, memlimit, memory_nodes, cgroup_values):
        """
        This method creates the CGroups for the following execution.
        @param my_cpus: None or a list of the CPU cores to use
        @param memlimit: None or memory limit in bytes
        @param memory_nodes: None or a list of memory nodes of a NUMA system to use
        @param cgroup_values: dict of additional values to set
        @return cgroups: a map of all the necessary cgroups for the following execution.
                         Please add the process of the following execution to all those cgroups!
        """
        logging.debug("Setting up cgroups for run.")

        # Setup cgroups, need a single call to create_cgroup() for all subsystems
        subsystems = [CPUACCT, FREEZER, MEMORY] + self._cgroup_subsystems
        if my_cpus is not None:
            subsystems.append(CPUSET)
        subsystems = [s for s in subsystems if s in self.cgroups]

        cgroups = self.cgroups.create_fresh_child_cgroup(*subsystems)

        logging.debug("Created cgroups %s.", cgroups)

        # First, set user-specified values such that they get overridden by our settings if necessary.
        for ((subsystem, option), value) in cgroup_values.items():
            try:
                cgroups.set_value(subsystem, option, value)
            except EnvironmentError as e:
                cgroups.remove()
                sys.exit('{} for setting cgroup option {}.{} to "{}" (error code {}).'
                         .format(e.strerror, subsystem, option, value, e.errno))
            logging.debug('Cgroup value %s.%s was set to "%s", new value is now "%s".',
                          subsystem, option, value, cgroups.get_value(subsystem, option))

        # Setup cpuset cgroup if necessary to limit the CPU cores/memory nodes to be used.
        if my_cpus is not None:
            my_cpus_str = ','.join(map(str, my_cpus))
            cgroups.set_value(CPUSET, 'cpus', my_cpus_str)
            my_cpus_str = cgroups.get_value(CPUSET, 'cpus')
            logging.debug('Using cpu cores [%s].', my_cpus_str)

        if memory_nodes is not None:
            cgroups.set_value(CPUSET, 'mems', ','.join(map(str, memory_nodes)))
            memory_nodesStr = cgroups.get_value(CPUSET, 'mems')
            logging.debug('Using memory nodes [%s].', memory_nodesStr)

        # Setup memory limit
        if memlimit is not None:
            limit = 'limit_in_bytes'
            cgroups.set_value(MEMORY, limit, memlimit)

            swap_limit = 'memsw.limit_in_bytes'
            # We need swap limit because otherwise the kernel just starts swapping
            # out our process if the limit is reached.
            # Some kernels might not have this feature,
            # which is ok if there is actually no swap.
            if not cgroups.has_value(MEMORY, swap_limit):
                if systeminfo.has_swap():
                    sys.exit('Kernel misses feature for accounting swap memory, but machine has swap. Please set swapaccount=1 on your kernel command line or disable swap with "sudo swapoff -a".')
            else:
                try:
                    cgroups.set_value(MEMORY, swap_limit, memlimit)
                except IOError as e:
                    if e.errno == errno.ENOTSUP: # kernel responds with operation unsupported if this is disabled
                        sys.exit('Memory limit specified, but kernel does not allow limiting swap memory. Please set swapaccount=1 on your kernel command line or disable swap with "sudo swapoff -a".')
                    raise e

            memlimit = cgroups.get_value(MEMORY, limit)
            logging.debug('Effective memory limit is %s bytes.', memlimit)

        if MEMORY in cgroups:
            try:
                # Note that this disables swapping completely according to
                # https://www.kernel.org/doc/Documentation/cgroups/memory.txt
                # (unlike setting the global swappiness to 0).
                # Our process might get killed because of this.
                cgroups.set_value(MEMORY, 'swappiness', '0')
            except IOError as e:
                logging.warning('Could not disable swapping for benchmarked process: %s', e)

        return cgroups

    def _setup_temp_dir(self):
        """Setup the temporary directories for the run.
        @return a triple of base directory, home directory, and temp directory
            (all can be removed after the run)
        """
        if self._user is None:
            base_dir = tempfile.mkdtemp(prefix="BenchExec_run_")
        else:
            create_temp_dir = self._build_cmdline([
                'python', '-c',
                'import tempfile;'
                'print(tempfile.mkdtemp(prefix="BenchExec_run_"))'
                ])
            base_dir = subprocess.check_output(create_temp_dir).decode().strip()
        temp_dir = os.path.join(base_dir, "tmp")
        home_dir = os.path.join(base_dir, "home")
        if self._user is None:
            os.mkdir(temp_dir)
            os.mkdir(home_dir)
        else:
            subprocess.check_call(self._build_cmdline(["mkdir", "--mode=700", temp_dir, home_dir]))

        logging.debug("Executing run with $HOME and $TMPDIR below %s.", base_dir)
        return base_dir, home_dir, temp_dir

    def _cleanup_temp_dir(self, base_dir):
        """Delete given temporary directory and all its contents."""
        if self._should_cleanup_temp_dir:
            logging.debug('Cleaning up temporary directory.')
            if self._user is None:
                util.rmtree(base_dir, onerror=util.log_rmtree_error)
            else:
                rm = subprocess.Popen(self._build_cmdline(['rm', '-rf', '--', base_dir]),
                                      stderr=subprocess.PIPE)
                rm_output = rm.stderr.read().decode()
                rm.stderr.close()
                if rm.wait() != 0 or rm_output:
                    logging.warning("Failed to clean up temp directory %s: %s.",
                                    base_dir, rm_output)
        else:
            logging.info("Skipping cleanup of temporary directory %s.", base_dir)

    def _setup_environment(self, environments, home_dir, temp_dir):
        """Return map with desired environment variables for run."""
        # If keepEnv is set or sudo is used, start from a fresh environment,
        # otherwise with the current one.
        # Set HOME and TMPDIR to the temporary directories create above.
        # keepEnv specifies variables to copy from the current environment,
        # newEnv specifies variables to set to a new value,
        # additionalEnv specifies variables where some value should be appended, and
        # clearEnv specifies variables to delete.
        if self._user is not None or environments.get("keepEnv", None) is not None:
            run_environment = {}
        else :
            run_environment = os.environ.copy()
        run_environment["HOME"] = home_dir
        run_environment["TMPDIR"] = temp_dir
        run_environment["TMP"] = temp_dir
        run_environment["TEMPDIR"] = temp_dir
        run_environment["TEMP"] = temp_dir
        for key, value in environments.get("keepEnv", {}).items():
            if key in os.environ:
                run_environment[key] = os.environ[key]
        for key, value in environments.get("newEnv", {}).items():
            run_environment[key] = value
        for key, value in environments.get("additionalEnv", {}).items():
            run_environment[key] = os.environ.get(key, "") + value
        for key in environments.get("clearEnv", {}).items():
            run_environment.pop(key, None)

        logging.debug("Using additional environment %s.", environments)
        return run_environment

    def _setup_output_file(self, output_filename, args):
        """Open and prepare output file."""
        # write command line into outputFile
        # (without environment variables, they are documented by benchexec)
        try:
            output_file = open(output_filename, 'w') # override existing file
        except IOError as e:
            sys.exit(e)
        output_file.write(' '.join(map(util.escape_string_shell, self._build_cmdline(args)))
                          + '\n\n\n' + '-' * 80 + '\n\n\n')
        output_file.flush()
        return output_file

    def _setup_cgroup_time_limit(self, hardtimelimit, softtimelimit, walltimelimit,
                                 cgroups, cores, pid_to_kill):
        """Start time-limit handler.
        @return None or the time-limit handler for calling cancel()
        """
        # hard time limit with cgroups is optional (additionally enforce by ulimit)
        cgroup_hardtimelimit = hardtimelimit if CPUACCT in cgroups else None

        if any([cgroup_hardtimelimit, softtimelimit, walltimelimit]):
            # Start a timer to periodically check timelimit
            timelimitThread = _TimelimitThread(cgroups=cgroups,
                                               hardtimelimit=cgroup_hardtimelimit,
                                               softtimelimit=softtimelimit,
                                               walltimelimit=walltimelimit,
                                               pid_to_kill=pid_to_kill,
                                               cores=cores,
                                               callbackFn=self._set_termination_reason,
                                               kill_process_fn=self._kill_process)
            timelimitThread.start()
            return timelimitThread
        return None

    def _setup_cgroup_memory_limit(self, memlimit, cgroups, pid_to_kill):
        """Start memory-limit handler.
        @return None or the memory-limit handler for calling cancel()
        """
        if memlimit is not None:
            try:
                oomThread = oomhandler.KillProcessOnOomThread(
                    cgroups=cgroups, pid_to_kill=pid_to_kill,
                    callbackFn=self._set_termination_reason,
                    kill_process_fn=self._kill_process)
                oomThread.start()
                return oomThread
            except OSError as e:
                logging.critical("OSError %s during setup of OomEventListenerThread: %s.",
                                 e.errno, e.strerror)
        return None

    def _setup_ulimit_time_limit(self, hardtimelimit, cgroups):
        """Setup time limit with ulimit for the current process."""
        if hardtimelimit is not None:
            # Also use ulimit for CPU time limit as a fallback if cgroups don't work.
            if CPUACCT in cgroups:
                # Use a slightly higher limit to ensure cgroups get used
                # (otherwise we cannot detect the timeout properly).
                ulimit = hardtimelimit + _ULIMIT_DEFAULT_OVERHEAD
            else:
                ulimit = hardtimelimit
            resource.setrlimit(resource.RLIMIT_CPU, (ulimit, ulimit))

    # --- run execution ---

    def execute_run(self, args, output_filename, stdin=None,
                   hardtimelimit=None, softtimelimit=None, walltimelimit=None,
                   cores=None, memlimit=None, memory_nodes=None,
                   environments={}, workingDir=None, maxLogfileSize=None,
                   cgroupValues={}):
        """
        This function executes a given command with resource limits,
        and writes the output to a file.
        @param args: the command line to run
        @param output_filename: the file where the output should be written to
        @param stdin: What to uses as stdin for the process (None: /dev/null, a file descriptor, or a file object)
        @param hardtimelimit: None or the CPU time in seconds after which the tool is forcefully killed.
        @param softtimelimit: None or the CPU time in seconds after which the tool is sent a kill signal.
        @param walltimelimit: None or the wall time in seconds after which the tool is forcefully killed (default: hardtimelimit + a few seconds)
        @param cores: None or a list of the CPU cores to use
        @param memlimit: None or memory limit in bytes
        @param memory_nodes: None or a list of memory nodes in a NUMA system to use
        @param environments: special environments for running the command
        @param workingDir: None or a directory which the execution should use as working directory
        @param maxLogfileSize: None or a number of bytes to which the output of the tool should be truncated approximately if there is too much output.
        @param cgroupValues: dict of additional cgroup values to set (key is tuple of subsystem and option, respective subsystem needs to be enabled in RunExecutor; cannot be used to override values set by BenchExec)
        @return: dict with result of run (measurement results and process exitcode)
        """
        # Check argument values and call the actual method _execute()

        if stdin == subprocess.PIPE:
            sys.exit('Illegal value subprocess.PIPE for stdin')
        elif stdin is None:
            stdin = DEVNULL

        if hardtimelimit is not None:
            if hardtimelimit <= 0:
                sys.exit("Invalid time limit {0}.".format(hardtimelimit))
        if softtimelimit is not None:
            if softtimelimit <= 0:
                sys.exit("Invalid soft time limit {0}.".format(softtimelimit))
            if hardtimelimit and (softtimelimit > hardtimelimit):
                sys.exit("Soft time limit cannot be larger than the hard time limit.")
            if not CPUACCT in self.cgroups:
                sys.exit("Soft time limit cannot be specified without cpuacct cgroup.")

        if walltimelimit is None:
            if hardtimelimit is not None:
                walltimelimit = hardtimelimit + _WALLTIME_LIMIT_DEFAULT_OVERHEAD
            elif softtimelimit is not None:
                walltimelimit = softtimelimit + _WALLTIME_LIMIT_DEFAULT_OVERHEAD
        else:
            if walltimelimit <= 0:
                sys.exit("Invalid wall time limit {0}.".format(walltimelimit))

        if cores is not None:
            if self.cpus is None:
                sys.exit("Cannot limit CPU cores without cpuset cgroup.")
            if not cores:
                sys.exit("Cannot execute run without any CPU core.")
            if not set(cores).issubset(self.cpus):
                sys.exit("Cores {0} are not allowed to be used".format(list(set(cores).difference(self.cpus))))

        if memlimit is not None:
            if memlimit <= 0:
                sys.exit("Invalid memory limit {0}.".format(memlimit))
            if not MEMORY in self.cgroups:
                sys.exit("Memory limit specified, but cannot be implemented without cgroup support.")

        if memory_nodes is not None:
            if self.memory_nodes is None:
                sys.exit("Cannot restrict memory nodes without cpuset cgroup.")
            if len(memory_nodes) == 0:
                sys.exit("Cannot execute run without any memory node.")
            if not set(memory_nodes).issubset(self.memory_nodes):
                sys.exit("Memory nodes {0} are not allowed to be used".format(list(set(memory_nodes).difference(self.memory_nodes))))

        if workingDir:
            if not os.path.exists(workingDir):
                sys.exit("Working directory {0} does not exist.".format(workingDir))
            if not os.path.isdir(workingDir):
                sys.exit("Working directory {0} is not a directory.".format(workingDir))
            if not os.access(workingDir, os.X_OK):
                sys.exit("Permission denied for working directory {0}.".format(workingDir))

        for ((subsystem, option), _) in cgroupValues.items():
            if not subsystem in self._cgroup_subsystems:
                sys.exit('Cannot set option "{option}" for subsystem "{subsystem}" that is not enabled. '
                         'Please specify "--require-cgroup-subsystem {subsystem}".'.format(
                            option=option, subsystem=subsystem))
            if not self.cgroups.has_value(subsystem, option):
                sys.exit('Cannot set option "{option}" for subsystem "{subsystem}", it does not exist.'
                         .format(option=option, subsystem=subsystem))

        try:
            return self._execute(args, output_filename, stdin,
                                 hardtimelimit, softtimelimit, walltimelimit, memlimit,
                                 cores, memory_nodes,
                                 cgroupValues,
                                 environments, workingDir, maxLogfileSize)

        except OSError as e:
            logging.critical("OSError %s while starting '%s' in '%s': %s.",
                             e.errno, args[0], workingDir or '.', e.strerror)
            return {'terminationreason': 'failed', 'exitcode': 0,
                    'cputime': 0, 'walltime': 0}

    def _execute(self, args, output_filename, stdin,
                 hardtimelimit, softtimelimit, walltimelimit, memlimit,
                 cores, memory_nodes,
                 cgroup_values,
                 environments, workingDir, max_output_size):
        """
        This method executes the command line and waits for the termination of it,
        handling all setup and cleanup, but does not check whether arguments are valid.
        """

        def preParent():
            """Setup that is executed in the parent process immediately before the actual tool is started."""
            # start measurements
            energy_before = util.measure_energy()
            walltime_before = util.read_monotonic_time()
            return (walltime_before, energy_before)

        def postParent(preParent_result):
            """Cleanup that is executed in the parent process immediately after the actual tool terminated."""
            # finish measurements
            walltime_before, energy_before = preParent_result
            walltime = util.read_monotonic_time() - walltime_before
            energy = util.measure_energy(energy_before)
            return (walltime, energy)

        def preSubprocess():
            """Setup that is executed in the forked process before the actual tool is started."""
            os.setpgrp() # make subprocess to group-leader
            os.nice(5) # increase niceness of subprocess
            self._setup_ulimit_time_limit(hardtimelimit, cgroups)

        # preparations that are not time critical
        cgroups = self._setup_cgroups(cores, memlimit, memory_nodes, cgroup_values)
        base_dir, home_dir, temp_dir = self._setup_temp_dir()
        run_environment = self._setup_environment(environments, home_dir, temp_dir)
        outputFile = self._setup_output_file(output_filename, args)
        args = self._build_cmdline(args, env=run_environment)

        timelimitThread = None
        oomThread = None
        pid = None
        returnvalue = 0
        ru_child = None
        self._termination_reason = None

        throttle_check = systeminfo.CPUThrottleCheck(cores)
        swap_check = systeminfo.SwapCheck()

        logging.debug('Starting process.')

        try:
            pid, result_fn = self._start_execution(args=args,
                stdin=stdin, stdout=outputFile, stderr=outputFile,
                env=run_environment, cwd=workingDir,
                cgroups=cgroups,
                parent_setup_fn=preParent, child_setup_fn=preSubprocess,
                parent_cleanup_fn=postParent)

            with self.SUB_PROCESS_PIDS_LOCK:
                self.SUB_PROCESS_PIDS.add(pid)

            timelimitThread = self._setup_cgroup_time_limit(
                hardtimelimit, softtimelimit, walltimelimit, cgroups, cores, pid)
            oomThread = self._setup_cgroup_memory_limit(memlimit, cgroups, pid)

            returnvalue, ru_child, (walltime, energy) = result_fn() # blocks until process has terminated

            result = {}
            result['walltime'] = walltime
            if energy:
                result['energy'] = energy

            # needs to come before cgroups.remove()
            self._get_cgroup_measurements(cgroups, ru_child, result)

        finally:
            # cleanup steps that need to get executed even in case of failure
            logging.debug('Process terminated, exit code %s.', returnvalue)

            with self.SUB_PROCESS_PIDS_LOCK:
                self.SUB_PROCESS_PIDS.discard(pid)

            if timelimitThread:
                timelimitThread.cancel()

            if oomThread:
                oomThread.cancel()

            # Kill all remaining processes if some managed to survive.
            # Because we send signals to all processes anyway we use the
            # internal function.
            cgroups.kill_all_tasks(self._kill_process0)

            # normally subprocess closes file, we do this again after all tasks terminated
            outputFile.close()

            logging.debug("Cleaning up cgroups.")
            cgroups.remove()

            self._cleanup_temp_dir(base_dir)

        # cleanup steps that are only relevant in case of success
        if throttle_check.has_throttled():
            logging.warning('CPU throttled itself during benchmarking due to overheating. '
                            'Benchmark results are unreliable!')
        if swap_check.has_swapped():
            logging.warning('System has swapped during benchmarking. '
                            'Benchmark results are unreliable!')

        _reduce_file_size_if_necessary(output_filename, max_output_size)

        if returnvalue not in [0,1]:
            _get_debug_output_after_crash(output_filename)

        result['exitcode'] = returnvalue
        if self._termination_reason:
            result['terminationreason'] = self._termination_reason
        elif memlimit and 'memory' in result and result['memory'] >= memlimit:
            # The kernel does not always issue OOM notifications and thus the OOMHandler
            # does not always run even in case of OOM. We detect this there and report OOM.
            result['terminationreason'] = 'memory'

        return result

    def _get_cgroup_measurements(self, cgroups, ru_child, result):
        """
        This method calculates the exact results for time and memory measurements.
        It is not important to call this method as soon as possible after the run.
        """
        logging.debug("Getting cgroup measurements.")

        cputime_wait = ru_child.ru_utime + ru_child.ru_stime if ru_child else 0
        cputime_cgroups = None
        if CPUACCT in cgroups:
            # We want to read the value from the cgroup.
            # The documentation warns about outdated values.
            # So we read twice with 0.1s time difference,
            # and continue reading as long as the values differ.
            # This has never happened except when interrupting the script with Ctrl+C,
            # but just try to be on the safe side here.
            tmp = cgroups.read_cputime()
            tmp2 = None
            while tmp != tmp2:
                time.sleep(0.1)
                tmp2 = tmp
                tmp = cgroups.read_cputime()
            cputime_cgroups = tmp

            # Usually cputime_cgroups seems to be 0.01s greater than cputime_wait.
            # Furthermore, cputime_wait might miss some subprocesses,
            # therefore we expect cputime_cgroups to be always greater (and more correct).
            # However, sometimes cputime_wait is a little bit bigger than cputime2.
            # For small values, this is probably because cputime_wait counts since fork,
            # whereas cputime_cgroups counts only after cgroups.add_task()
            # (so overhead from runexecutor is correctly excluded in cputime_cgroups).
            # For large values, a difference may also indicate a problem with cgroups,
            # for example another process moving our benchmarked process between cgroups,
            # thus we warn if the difference is substantial and take the larger cputime_wait value.
            if cputime_wait > 0.5 and (cputime_wait * 0.95) > cputime_cgroups:
                logging.warning(
                    'Cputime measured by wait was %s, cputime measured by cgroup was only %s, '
                    'perhaps measurement is flawed.',
                    cputime_wait, cputime_cgroups)
                result['cputime'] = cputime_wait
            else:
                result['cputime'] = cputime_cgroups

            for (core, coretime) in enumerate(cgroups.get_value(CPUACCT, 'usage_percpu').split(" ")):
                try:
                    coretime = int(coretime)
                    if coretime != 0:
                        result['cputime-cpu'+str(core)] = coretime/1000000000 # nano-seconds to seconds
                except (OSError, ValueError) as e:
                    logging.debug("Could not read CPU time for core %s from kernel: %s", core, e)

        if MEMORY in cgroups:
            # This measurement reads the maximum number of bytes of RAM+Swap the process used.
            # For more details, c.f. the kernel documentation:
            # https://www.kernel.org/doc/Documentation/cgroups/memory.txt
            memUsageFile = 'memsw.max_usage_in_bytes'
            if not cgroups.has_value(MEMORY, memUsageFile):
                memUsageFile = 'max_usage_in_bytes'
            if not cgroups.has_value(MEMORY, memUsageFile):
                logging.warning('Memory-usage is not available due to missing files.')
            else:
                try:
                    result['memory'] = int(cgroups.get_value(MEMORY, memUsageFile))
                except IOError as e:
                    if e.errno == errno.ENOTSUP: # kernel responds with operation unsupported if this is disabled
                        logging.critical(
                            "Kernel does not track swap memory usage, cannot measure memory usage."
                            " Please set swapaccount=1 on your kernel command line.")
                    else:
                        raise e

        logging.debug(
            'Resource usage of run: walltime=%s, cputime=%s, cgroup-cputime=%s, memory=%s',
            result['walltime'], cputime_wait, cputime_cgroups, result.get('memory', None))

    # --- other public functions ---

    def stop(self):
        self._set_termination_reason('killed')
        super(RunExecutor, self).stop()

    def check_for_new_files_in_home(self):
        """Check that the user account's home directory now does not contain more files than
        when this instance was created, and warn otherwise.
        Does nothing if no user account was given to RunExecutor.
        @return set of newly created files
        """
        if not self._user:
            return None
        try:
            created_files = set(self._listdir(self._home_dir)).difference(self._home_dir_content)
        except (subprocess.CalledProcessError, IOError):
            # Probably home directory does not exist
            created_files = []
        if created_files:
            logging.warning('The tool created the following files in %s, '
                            'this may influence later runs:\n\t%s',
                            self._home_dir, '\n\t'.join(created_files))
        return created_files

def _get_user_account_info(user):
    """Get the user account info from the passwd database. Only works on Linux.
    @param user The name of a user account or a numeric uid prefixed with '#'
    @return a tuple that corresponds to the members of the passwd structure
    @raise KeyError: If user account is unknown
    @raise ValueError: If uid is not a valid number
    """
    import pwd # Import here to avoid problems on other platforms
    if user[0] == '#':
        return pwd.getpwuid(int(user[1:]))
    else:
        return pwd.getpwnam(user)

def _reduce_file_size_if_necessary(fileName, maxSize):
    """
    This function shrinks a file.
    We remove only the middle part of a file,
    the file-start and the file-end remain unchanged.
    """
    fileSize = os.path.getsize(fileName)

    if maxSize is None:
        logging.debug("Size of logfile '%s' is %s bytes, size limit disabled.", fileName, fileSize)
        return # disabled, nothing to do

    if fileSize < (maxSize + 500):
        logging.debug("Size of logfile '%s' is %s bytes, nothing to do.", fileName, fileSize)
        return

    logging.warning("Logfile '%s' is too big (size %s bytes). Removing lines.", fileName, fileSize)
    util.shrink_text_file(fileName, maxSize, _LOG_SHRINK_MARKER)

def _get_debug_output_after_crash(output_filename):
    """
    Segmentation faults and some memory failures reference a file
    with more information (hs_err_pid_*). We append this file to the log.
    The format that we expect is a line
    "# An error report file with more information is saved as:"
    and the file name of the dump file on the next line.
    """
    logging.debug("Analysing output for crash info.")
    foundDumpFile = False
    with open(output_filename, 'r+') as outputFile:
        for line in outputFile:
            if foundDumpFile:
                try:
                    dumpFileName = line.strip(' #\n')
                    outputFile.seek(0, os.SEEK_END) # jump to end of log file
                    with open(dumpFileName, 'r') as dumpFile:
                        util.copy_all_lines_from_to(dumpFile, outputFile)
                    os.remove(dumpFileName)
                except IOError as e:
                    logging.warning('Could not append additional segmentation fault information '
                                    'from %s (%s)',
                                    dumpFile, e.strerror)
                break
            try:
                if util.decode_to_string(line).startswith('# An error report file with more information is saved as:'):
                    logging.debug('Going to append error report file')
                    foundDumpFile = True
            except UnicodeDecodeError:
                pass
                # ignore invalid chars from logfile

class _TimelimitThread(threading.Thread):
    """
    Thread that periodically checks whether the given process has already
    reached its timelimit. After this happens, the process is terminated.
    """
    def __init__(self, cgroups, kill_process_fn, hardtimelimit, softtimelimit, walltimelimit, pid_to_kill, cores,
                 callbackFn=lambda reason: None):
        super(_TimelimitThread, self).__init__()

        if hardtimelimit or softtimelimit:
            assert CPUACCT in cgroups
        assert walltimelimit is not None

        if cores:
            self.cpuCount = len(cores)
        else:
            try:
                self.cpuCount = multiprocessing.cpu_count()
            except NotImplementedError:
                self.cpuCount = 1

        self.daemon = True
        self.cgroups = cgroups
        self.timelimit = hardtimelimit or (60*60*24*365*100) # large dummy value
        self.softtimelimit = softtimelimit or (60*60*24*365*100) # large dummy value
        self.latestKillTime = util.read_monotonic_time() + walltimelimit
        self.pid_to_kill = pid_to_kill
        self.callback = callbackFn
        self.kill_process = kill_process_fn
        self.finished = threading.Event()

    def read_cputime(self):
        while True:
            try:
                return self.cgroups.read_cputime()
            except ValueError:
                # Sometimes the kernel produces strange values with linebreaks in them
                time.sleep(1)

    def run(self):
        while not self.finished.is_set():
            usedCpuTime = self.read_cputime() if CPUACCT in self.cgroups else 0
            remainingCpuTime = self.timelimit - usedCpuTime
            remainingSoftCpuTime = self.softtimelimit - usedCpuTime
            remainingWallTime = self.latestKillTime - util.read_monotonic_time()
            logging.debug(
                "TimelimitThread for process %s: used CPU time: %s, remaining CPU time: %s, "
                "remaining soft CPU time: %s, remaining wall time: %s.",
                self.pid_to_kill, usedCpuTime, remainingCpuTime,
                remainingSoftCpuTime, remainingWallTime)
            if remainingCpuTime <= 0:
                self.callback('cputime')
                logging.debug('Killing process %s due to CPU time timeout.', self.pid_to_kill)
                self.kill_process(self.pid_to_kill, self.cgroups)
                self.finished.set()
                return
            if remainingWallTime <= 0:
                self.callback('walltime')
                logging.warning('Killing process %s due to wall time timeout.', self.pid_to_kill)
                self.kill_process(self.pid_to_kill, self.cgroups)
                self.finished.set()
                return

            if remainingSoftCpuTime <= 0:
                self.callback('cputime-soft')
                # soft time limit violated, ask process to terminate
                self.kill_process(self.pid_to_kill, self.cgroups, signal.SIGTERM)
                self.softtimelimit = self.timelimit

            remainingTime = min(remainingCpuTime/self.cpuCount,
                                remainingSoftCpuTime/self.cpuCount,
                                remainingWallTime)
            self.finished.wait(remainingTime + 1)

    def cancel(self):
        self.finished.set()

if __name__ == '__main__':
    main()


# Import the pyeapi library
import pyeapi

# Open a session with leaf1-DC1. The script will find .eapi.conf and reference the credentials automatically

connect = pyeapi.connect_to('leaf1-DC1')

# "Create" sets the port as a Layer 3 port (no switchport)

connect.api("ipinterfaces").create('Ethernet4')

# Set Ethernet4 for the IP address of 4.4.4.4 and put the result into the variable (boolean) result

result = connect.api("ipinterfaces").set_address('Ethernet4','4.4.4.4/24')

# This is just very basic error handling here, it gives a yes or no answer depending on whether a "200 OK" response was given, or another error was occurred.

if result == True:
    print("Completed!")
if result == False:
    print("Error!")


# import packages
import sys
import pandas as pd
import torch
import torchvision
from torch.utils.data import DataLoader
from torchvision import transforms

# load the functions to do spatial CV
from analysis_utils.spatial_CV import split_lsms_spatial

# import the torch_framework package
from analysis_utils.torch_framework.CrossValidator import CrossValidator
from analysis_utils.torch_framework.ResNet18 import ResNet18
from analysis_utils.torch_framework.SatDataset import SatDataset
from analysis_utils.torch_framework.torch_helpers import get_agg_img_stats, get_feat_stats, get_target_stats
from analysis_utils.torch_framework.torch_helpers import standardise as torch_standardise
from analysis_utils.torch_framework.CvEvaluator import CvEvaluator

# import the helpers to demean the data and get the deltas etc.
from analysis_utils.analysis_helpers import demean_df, make_delta_df

# check if the number of command line arguments is correct
if len(sys.argv) != 4:
    print("Please provide the values for the following variables:")
    print("model_name, cv_object_name, target_var")
    print("Usage: python my_script.py model_name cv_object_name target_var")
    raise(ValueError("Incorrect number of command line arguments"))

print("Hello!")
print("Initialising Training for the Within model using Remote sensing images")
print(f"Target variable: {sys.argv[3]}")
print("This is the delta version of the target variable. Just keeps the same name.")
print("\n")

####################################################################################################
# Accept the command line arguments
####################################################################################################

# get the command line arguments
model_name = sys.argv[1]
cv_object_name = sys.argv[2]
delta_target_var = sys.argv[3]

####################################################################################################
# Set the hyper-parameters
####################################################################################################
data_type = 'RS_v2'
id_var = 'delta_id'
cv_evaluator_object_name = f"evaluator_demeaned_{cv_object_name}"

# set the random seed
spatial_cv_random_seed = 348
random_seed = 230897

# set the number of folds for k-fold CV
n_folds = 5

# share of data used
max_obs = 10000000

# set hyper-parameters
hyper_params = {
    'lr': [1e-2, 1e-3],
    'batch_size': [256],
    'alpha': [1e-1, 1e-2, 1e-3],
    'step_size': [1],
    'gamma': [0.96],
    'n_epochs': [200],
    'patience': [40]
}

# training device
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
print(f"Training device: {device} \n")

####################################################################################################
# Set the data paths
####################################################################################################

# set the global file paths
root_data_dir = "/scratch/users/nschmid5/data_analysis/Data"
#root_data_dir = "../../Data"
# the lsms data
lsms_pth = f"{root_data_dir}/lsms/processed/labels_cluster_v1.csv"

# set the image directories
sat_img_dir = f"{root_data_dir}/satellite_imgs"

# median LS images at the cluster level
delta_img_dir = f"{sat_img_dir}/{data_type}/RS_v2_delta"
delta_stats_pth = f"{sat_img_dir}/{data_type}/RS_v2_delta_img_stats.pkl"

####################################################################################################
# Load the data
####################################################################################################
# load the LSMS data
lsms_df = pd.read_csv(lsms_pth).iloc[:max_obs, :]

cl_df = lsms_df[['cluster_id', 'country', 'lat', 'lon']].copy().drop_duplicates().reset_index(drop=True)
within_df = lsms_df[['cluster_id', 'unique_id', 'log_mean_pc_cons_usd_2017']]

# demean the df
demeaned_df = demean_df(within_df)
demeaned_df = demeaned_df.rename(columns={'unique_id': 'delta_id'})
demeaned_df = pd.merge(demeaned_df, cl_df, on='cluster_id', how='left')

# create delta df
print("Make Delta DF")
delta_df = make_delta_df(within_df)
delta_df = pd.merge(delta_df, cl_df, on='cluster_id', how='left')

# add the demeaned df to the delta df
delta_df = pd.concat([delta_df, demeaned_df]).reset_index(drop=True)

# get the fold ids from spatial CV
print("\nDividing the data into k different folds using spatial CV")
fold_ids = split_lsms_spatial(lsms_df, n_folds=n_folds, random_seed=spatial_cv_random_seed)

# extract the image statistics for the demeaned images
delta_img_stats = get_agg_img_stats(delta_stats_pth, delta_df, id_var='delta_id')
delta_feat_stats = get_feat_stats(delta_img_stats)

# get the target stats
delta_target_stats = get_target_stats(delta_df, delta_target_var)

# define the transforms
# get the data transforms for the target --> is used in the DataLoader object
delta_target_transform = transforms.Compose([
    torchvision.transforms.Lambda(
        lambda t: torch_standardise(t, delta_target_stats['mean'], delta_target_stats['std'])),
])

# get the data transform for the Landsat image (normalisation and random horizontal + vertical flips)
delta_feat_transform = torchvision.transforms.Compose(
    [torchvision.transforms.RandomVerticalFlip(.5),
     torchvision.transforms.RandomHorizontalFlip(.5),
     transforms.Normalize(delta_feat_stats['mean'], delta_feat_stats['std'])]
)

####################################################################################################
# Train the model
####################################################################################################
print("\n\n")
print("=====================================================================")
print("Initialise CROSS VALIDATION")
print("=====================================================================")

# load the data into RAM first
# this increases training times
_dat = SatDataset(delta_df,
                  delta_img_dir,
                  data_type,
                  delta_target_var,
                  id_var,
                  delta_feat_transform,
                  delta_target_transform,
                  random_seed)
_loader = DataLoader(_dat, batch_size=hyper_params['batch_size'][0], shuffle=False)
_, _ = next(iter(_loader))

# run model training
# initialise the model and the CrossValidator object
resnet18 = ResNet18(
    input_channels=4,
    use_pretrained_weights=False,
    scaled_weight_init=False,
    random_seed=random_seed
)

cv = CrossValidator(
    model_class=resnet18,
    lsms_df=delta_df,
    fold_ids=fold_ids,
    img_dir=delta_img_dir,
    data_type=data_type,
    target_var=delta_target_var,
    id_var=id_var,
    feat_transform=delta_feat_transform,
    target_transform=delta_target_transform,
    device=device,
    model_name=model_name,
    random_seed=random_seed
)

# run k-fold-cv
cv.run_cv(hyper_params=hyper_params, tune_hyper_params=True)

# save the cv object
cv.save_object(name=cv_object_name)

# output the overall performance of the model
print("\n")
print('=' * 100)
print('Cross-Validated performance:')
print('=' * 100)
print(cv.compute_overall_performance(use_fold_weights=True))
print('\n\n')

####################################################################################################
# Use the trained model to predict only the demeaned images
####################################################################################################
print("Evaluate the model on the demeaned images")
print("\n")
# get the statistics for the demeaned dataset

# extract the image statistics for the demeaned images
demeaned_img_stats = get_agg_img_stats(delta_stats_pth, demeaned_df, id_var='delta_id')
demeaned_feat_stats = get_feat_stats(demeaned_img_stats)

# get the target stats
demeaned_target_stats = get_target_stats(demeaned_df, delta_target_var)

# define the transforms
# get the data transforms for the target --> is used in the DataLoader object
demeaned_target_transform = transforms.Compose([
    torchvision.transforms.Lambda(
        lambda t: torch_standardise(t, demeaned_target_stats['mean'], demeaned_target_stats['std'])),
])

# get the data transform for the Landsat image (normalisation and random horizontal + vertical flips)
demeaned_feat_transform = torchvision.transforms.Compose([
     transforms.Normalize(demeaned_feat_stats['mean'], demeaned_feat_stats['std'])
])

# load the CvEvaluator object
cv_evaluator = CvEvaluator(
    cv_object=cv,
    lsms_df=demeaned_df,
    fold_ids=fold_ids,
    device=device
)

cv_evaluator.evaluate(
    feat_transform=demeaned_feat_transform,
    target_transform=demeaned_target_transform
)

cv_evaluator.save_object(name=cv_evaluator_object_name)

# output the overall performance of the model
print("\n")
print('=' * 100)
print('Cross-Validated performance on delta images:')
print('=' * 100)
print(cv.compute_overall_performance(use_fold_weights=True))
print('\n\n')

# output the overall performance of the model
print("\n")
print('=' * 100)
print('Cross-Validated performance on demeaned images:')
print('=' * 100)
print(cv_evaluator.compute_overall_performance(use_fold_weights=True))
print('\n\n')
print('BYE BYE')


import json

from common.exceptions import BaseAPIError
from common.log import logger
from datahub.databus.exceptions import (
    TaskAccessResourceNotFound,
    TaskConfigErr,
    TaskCreateErr,
    TaskGenConfigErr,
)
from datahub.databus.settings import MODULE_PULLER
from datahub.databus.task.puller_task import start_puller_task
from datahub.databus.task.pulsar import topic as pulsar_topic
from datahub.databus.task.pulsar.topic import get_pulsar_channel_token
from datahub.databus.task.status import set_databus_task_status
from datahub.databus.task.task_utils import (
    _stop_task_without_exception,
    config_factory,
    generate_puller_connector_name,
    get_available_puller_cluster_name,
    is_databus_task_exists,
)

from datahub.databus import model_manager, models, rawdata, settings

class BasePuller(object):
    component = None
    module = MODULE_PULLER

    def __init__(self, data_id, raw_data, storage_channel, bk_username):
        """
        :param rt_id: 该任务的result_table_id
        :param bk_username: 该任务的操作者
        :param storage_params: 此次puller任务创建时的原始参数字符串，用做日志记录
        """
        self.data_id = data_id
        self.raw_data = raw_data
        self.bk_username = bk_username
        self.storage_channel = storage_channel
        self.connector_name = generate_puller_connector_name(self.component, data_id)
        self.cluster_name = self._get_connector_cluster_name()
        self.sink_topic = rawdata.get_topic_name(raw_data)
        self.resource_info, self.resource_json = self._get_resource_info()
        self.config_factory = config_factory[storage_channel.cluster_type]

    def _get_connector_cluster_name(self):
        connectors = model_manager.get_connector_route(self.connector_name)
        if connectors:
            # 若存在记录则使用上次的集群
            return connectors.cluster_name
        else:
            return get_available_puller_cluster_name(
                self.storage_channel.cluster_type,
                self.storage_channel.cluster_name,
                self.component,
            )

    def _get_resource_info(self):
        try:
            resource_info = models.AccessResourceInfo.objects.get(raw_data_id=self.data_id)
            resource_json = json.loads(resource_info.resource)
        except models.AccessResourceInfo.DoesNotExist:
            raise TaskAccessResourceNotFound(message_kv={"data_id": self.data_id})
        except Exception:
            raise TaskConfigErr()
        return resource_info, resource_json

    def start(self):
        try:
            self._before_add_task()
            self._add_databus_puller_task()
            conf = self._get_puller_task_conf()
            if self._need_start_task():
                logger.info(
                    "going to start {} puller {} with conf {}".format(self.component, self.connector_name, conf)
                )
                self._start_databus_puller_task(conf)
        except BaseAPIError as task_error:
            logger.error("start databus puller task failed, exception %s" % task_error.message)
            raise TaskCreateErr(
                message_kv={
                    "data_id": self.data_id,
                    "message": "{}({})".format(task_error.message, task_error.CODE),
                }
            )

    def stop(self):
        obj = model_manager.get_connector_route(self.connector_name)
        if not obj:
            return
        _stop_task_without_exception(self.connector_name, self.cluster_name)

    def _before_add_task(self):
        if self.storage_channel.cluster_type == settings.TYPE_PULSAR:
            topic = rawdata.get_topic_name(self.raw_data)
            pulsar_topic.create_partition_topic_if_not_exist(
                self.storage_channel,
                topic,
                settings.PULSAR_DEFAULT_PARTITION,
                get_pulsar_channel_token(self.storage_channel),
            )

    def _add_databus_puller_task(self):
        """
        添加拉取任务
        """
        logger.info(
            "add %s task: component=%s topic=%s channel_name=%s"
            % (
                self.connector_name,
                self.component,
                self.sink_topic,
                self.storage_channel.cluster_name,
            )
        )

        if not is_databus_task_exists(
            self.connector_name,
            self.storage_channel.cluster_type,
            self.module,
            self.component,
            self.storage_channel.cluster_name,
        ):
            # 注意，puller目前没有rt_id信息，后面可能要确认寻找
            model_manager.add_databus_task(
                "raw_data_id-%s" % self.data_id,
                self.connector_name,
                self.storage_channel.cluster_name,
                "{}#{}".format(self.component, self.data_id),
                self.component,
                "{}#{}".format(self.storage_channel.cluster_name, self.sink_topic),
                self.storage_channel.cluster_type,
                settings.MODULE_PULLER,
            )

    def _get_puller_task_conf(self):
        """
        获取分发任务配置
        """
        raise TaskGenConfigErr(message_kv={"task": self.connector_name, "cluster": self.cluster_name})

    def _need_start_task(self):
        return True

    def _start_databus_puller_task(self, conf):
        # 启动任务
        start_puller_task(
            self.storage_channel.cluster_type,
            self.cluster_name,
            self.connector_name,
            conf,
        )

    def set_databus_task_status(self):
        set_databus_task_status(self.connector_name, self.cluster_name, models.DataBusTaskStatus.RUNNING)

    @classmethod
    def _compare_connector_conf(cls, cluster_name, connector_name, conf, running_conf):
        return False