import os
import re
import sys
import time
import fcntl
import shelve
import datetime
import traceback
import requests

from urllib3.exceptions import InsecureRequestWarning

import pandaserver.userinterface.Client as Client
from pandaserver.taskbuffer import EventServiceUtils
from pandacommon.pandalogger.PandaLogger import PandaLogger
from pandaserver.jobdispatcher.Watcher import Watcher
from pandaserver.brokerage.SiteMapper import SiteMapper
# from pandaserver.dataservice.MailUtils import MailUtils
from pandaserver.srvcore.CoreUtils import commands_get_status_output
from pandaserver.config import panda_config

# logger
_logger = PandaLogger().getLogger('copyArchive')

# main
def main(argv=tuple(), tbuf=None, **kwargs):
    # password
    requests.packages.urllib3.disable_warnings(category=InsecureRequestWarning)

    _logger.debug("===================== start =====================")

    # memory checker
    def _memoryCheck(str):
        try:
            proc_status = '/proc/%d/status' % os.getpid()
            procfile = open(proc_status)
            name   = ""
            vmSize = ""
            vmRSS  = ""
            # extract Name,VmSize,VmRSS
            for line in procfile:
                if line.startswith("Name:"):
                    name = line.split()[-1]
                    continue
                if line.startswith("VmSize:"):
                    vmSize = ""
                    for item in line.split()[1:]:
                        vmSize += item
                    continue
                if line.startswith("VmRSS:"):
                    vmRSS = ""
                    for item in line.split()[1:]:
                        vmRSS += item
                    continue
            procfile.close()
            _logger.debug('MemCheck - %s Name=%s VSZ=%s RSS=%s : %s' % (os.getpid(),name,vmSize,vmRSS,str))
        except Exception:
            type, value, traceBack = sys.exc_info()
            _logger.error("memoryCheck() : %s %s" % (type,value))
            _logger.debug('MemCheck - %s unknown : %s' % (os.getpid(),str))
        return

    _memoryCheck("start")

    # # kill old dq2 process
    # try:
    #     # time limit
    #     timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=2)
    #     # get process list
    #     scriptName = sys.argv[0]
    #     out = commands_get_status_output(
    #         'ps axo user,pid,lstart,args | grep dq2.clientapi | grep -v PYTHONPATH | grep -v grep')[-1]
    #     for line in out.split('\n'):
    #         if line == '':
    #             continue
    #         items = line.split()
    #         # owned process
    #         if items[0] not in ['sm','atlpan','pansrv','root']: # ['os.getlogin()']: doesn't work in cron
    #             continue
    #         # look for python
    #         if re.search('python',line) is None:
    #             continue
    #         # PID
    #         pid = items[1]
    #         # start time
    #         timeM = re.search('(\S+\s+\d+ \d+:\d+:\d+ \d+)',line)
    #         startTime = datetime.datetime(*time.strptime(timeM.group(1),'%b %d %H:%M:%S %Y')[:6])
    #         # kill old process
    #         if startTime < timeLimit:
    #             _logger.debug("old dq2 process : %s %s" % (pid,startTime))
    #             _logger.debug(line)
    #             commands_get_status_output('kill -9 %s' % pid)
    # except Exception:
    #     type, value, traceBack = sys.exc_info()
    #     _logger.error("kill dq2 process : %s %s" % (type,value))
    #
    #
    # # kill old process
    # try:
    #     # time limit
    #     timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=7)
    #     # get process list
    #     scriptName = sys.argv[0]
    #     out = commands_get_status_output('ps axo user,pid,lstart,args | grep %s' % scriptName)[-1]
    #     for line in out.split('\n'):
    #         items = line.split()
    #         # owned process
    #         if items[0] not in ['sm','atlpan','pansrv','root']: # ['os.getlogin()']: doesn't work in cron
    #             continue
    #         # look for python
    #         if re.search('python',line) is None:
    #             continue
    #         # PID
    #         pid = items[1]
    #         # start time
    #         timeM = re.search('(\S+\s+\d+ \d+:\d+:\d+ \d+)',line)
    #         startTime = datetime.datetime(*time.strptime(timeM.group(1),'%b %d %H:%M:%S %Y')[:6])
    #         # kill old process
    #         if startTime < timeLimit:
    #             _logger.debug("old process : %s %s" % (pid,startTime))
    #             _logger.debug(line)
    #             commands_get_status_output('kill -9 %s' % pid)
    # except Exception:
    #     type, value, traceBack = sys.exc_info()
    #     _logger.error("kill process : %s %s" % (type,value))

    # instantiate TB
    # if tbuf is None:
    from pandaserver.taskbuffer.TaskBuffer import taskBuffer
    taskBuffer.init(panda_config.dbhost,panda_config.dbpasswd,nDBConnection=1)
    # else:
    #     taskBuffer = tbuf

    # instantiate sitemapper
    siteMapper = SiteMapper(taskBuffer)

    # send email for access requests
    _logger.debug("Site Access")
    try:
        # get contact
        contactAddr = {}
        siteContactAddr = {}
        sql = "SELECT name,email FROM ATLAS_PANDAMETA.cloudconfig"
        status,res = taskBuffer.querySQLS(sql,{})
        for cloudName,cloudEmail in res:
            contactAddr[cloudName] = cloudEmail
        # get requests
        sql  = "SELECT pandaSite,status,dn FROM ATLAS_PANDAMETA.siteaccess WHERE status IN (:status1,:status2,:status3) "
        sql += "ORDER BY pandaSite,status "
        varMap = {}
        varMap[':status1'] = 'requested'
        varMap[':status2'] = 'tobeapproved'
        varMap[':status3'] = 'toberejected'
        status,res = taskBuffer.querySQLS(sql,varMap)
        requestsInCloud = {}
        # mailUtils = MailUtils()
        # loop over all requests
        for pandaSite,reqStatus,userName in res:
            cloud = siteMapper.getSite(pandaSite).cloud
            _logger.debug("request : '%s' site=%s status=%s cloud=%s" % (userName,pandaSite,reqStatus,cloud))
            # send emails to user
            if reqStatus in ['tobeapproved','toberejected']:
                # set status
                if reqStatus == 'tobeapproved':
                    newStatus = 'approved'
                else:
                    newStatus = 'rejected'
                # get mail address for user
                userMailAddr = ''
                sqlUM = "SELECT email FROM ATLAS_PANDAMETA.users WHERE name=:userName"
                varMap = {}
                varMap[':userName'] = userName
                stUM,resUM = taskBuffer.querySQLS(sqlUM,varMap)
                if resUM is None or len(resUM) == 0:
                    _logger.error("email address is unavailable for '%s'" % userName)
                else:
                    userMailAddr = resUM[0][0]
                # send
                # if userMailAddr not in ['',None,'None','notsend']:
                #     _logger.debug("send update to %s" % userMailAddr)
                #     retMail = mailUtils.sendSiteAccessUpdate(userMailAddr,newStatus,pandaSite)
                #     _logger.debug(retMail)
                # update database
                sqlUp  = "UPDATE ATLAS_PANDAMETA.siteaccess SET status=:newStatus "
                sqlUp += "WHERE pandaSite=:pandaSite AND dn=:userName"
                varMap = {}
                varMap[':userName']  = userName
                varMap[':newStatus'] = newStatus
                varMap[':pandaSite'] = pandaSite
                stUp,resUp = taskBuffer.querySQLS(sqlUp,varMap)
            else:
                # append cloud
                requestsInCloud.setdefault(cloud, {})
                # append site
                requestsInCloud[cloud].setdefault(pandaSite, [])
                # append user
                requestsInCloud[cloud][pandaSite].append(userName)
        # send requests to the cloud responsible
        for cloud in requestsInCloud:
            requestsMap = requestsInCloud[cloud]
            _logger.debug("requests for approval : cloud=%s" % cloud)
            # send
            if cloud in contactAddr and contactAddr[cloud] not in ['',None,'None']:
                # get site contact
                for pandaSite in requestsMap:
                    userNames = requestsMap[pandaSite]
                    if pandaSite not in siteContactAddr:
                        varMap = {}
                        varMap[':siteid'] = pandaSite
                        sqlSite = "SELECT email FROM ATLAS_PANDAMETA.schedconfig WHERE siteid=:siteid AND rownum<=1"
                        status,res = taskBuffer.querySQLS(sqlSite,varMap)
                        siteContactAddr[pandaSite] = res[0][0]
                        # append
                        if siteContactAddr[pandaSite] not in ['',None,'None']:
                            contactAddr[cloud] += ',%s' % siteContactAddr[pandaSite]
            else:
                _logger.error("contact email address is unavailable for %s" % cloud)
    except Exception:
        type, value, traceBack = sys.exc_info()
        _logger.error("Failed with %s %s" % (type,value))
    _logger.debug("Site Access : done")

    # finalize failed jobs
    _logger.debug("AnalFinalizer session")
    try:
        # get min PandaID for failed jobs in Active table
        sql  = "SELECT MIN(PandaID),prodUserName,jobDefinitionID,jediTaskID,computingSite FROM ATLAS_PANDA.jobsActive4 "
        sql += "WHERE prodSourceLabel=:prodSourceLabel AND jobStatus=:jobStatus "
        sql += "GROUP BY prodUserName,jobDefinitionID,jediTaskID,computingSite "
        varMap = {}
        varMap[':jobStatus']       = 'failed'
        varMap[':prodSourceLabel'] = 'user'
        status,res = taskBuffer.querySQLS(sql,varMap)
        if res is not None:
            # loop over all user/jobdefID
            for pandaID,prodUserName,jobDefinitionID,jediTaskID,computingSite in res:
                # check
                _logger.debug("check finalization for %s task=%s jobdefID=%s site=%s" % (prodUserName,jediTaskID,
                                                                                jobDefinitionID,
                                                                                computingSite))
                sqlC  = "SELECT COUNT(*) FROM ("
                sqlC += "SELECT PandaID FROM ATLAS_PANDA.jobsActive4 "
                sqlC += "WHERE prodSourceLabel=:prodSourceLabel AND prodUserName=:prodUserName "
                sqlC += "AND jediTaskID=:jediTaskID "
                sqlC += "AND computingSite=:computingSite "
                sqlC += "AND jobDefinitionID=:jobDefinitionID "
                sqlC += "AND NOT jobStatus IN (:jobStatus1,:jobStatus2) "
                sqlC += "UNION "
                sqlC += "SELECT PandaID FROM ATLAS_PANDA.jobsDefined4 "
                sqlC += "WHERE prodSourceLabel=:prodSourceLabel AND prodUserName=:prodUserName "
                sqlC += "AND jediTaskID=:jediTaskID "
                sqlC += "AND computingSite=:computingSite "
                sqlC += "AND jobDefinitionID=:jobDefinitionID "
                sqlC += "AND NOT jobStatus IN (:jobStatus1,:jobStatus2) "
                sqlC += ") "
                varMap = {}
                varMap[':jobStatus1']      = 'failed'
                varMap[':jobStatus2']      = 'merging'
                varMap[':prodSourceLabel'] = 'user'
                varMap[':jediTaskID']      = jediTaskID
                varMap[':computingSite']   = computingSite
                varMap[':prodUserName']    = prodUserName
                varMap[':jobDefinitionID'] = jobDefinitionID
                statC,resC = taskBuffer.querySQLS(sqlC,varMap)
                # finalize if there is no non-failed jobs
                if resC is not None:
                    _logger.debug("n of non-failed jobs : %s" % resC[0][0])
                    if resC[0][0] == 0:
                        jobSpecs = taskBuffer.peekJobs([pandaID],fromDefined=False,fromArchived=False,fromWaiting=False)
                        jobSpec = jobSpecs[0]
                        if jobSpec is None:
                            _logger.debug("skip PandaID={0} not found in jobsActive".format(pandaID))
                            continue
                        _logger.debug("finalize %s %s" % (prodUserName,jobDefinitionID))
                        finalizedFlag = taskBuffer.finalizePendingJobs(prodUserName,jobDefinitionID)
                        _logger.debug("finalized with %s" % finalizedFlag)
                        if finalizedFlag and jobSpec.produceUnMerge():
                            # collect sub datasets
                            subDsNames = set()
                            subDsList = []
                            for tmpFileSpec in jobSpec.Files:
                                if tmpFileSpec.type in ['log','output'] and \
                                        re.search('_sub\d+$',tmpFileSpec.destinationDBlock) is not None:
                                    if tmpFileSpec.destinationDBlock in subDsNames:
                                        continue
                                    subDsNames.add(tmpFileSpec.destinationDBlock)
                                    datasetSpec = taskBuffer.queryDatasetWithMap({'name':tmpFileSpec.destinationDBlock})
                                    subDsList.append(datasetSpec)
                            _logger.debug("update unmerged datasets")
                            taskBuffer.updateUnmergedDatasets(jobSpec,subDsList)
                else:
                    _logger.debug("n of non-failed jobs : None")
    except Exception:
        errType,errValue = sys.exc_info()[:2]
        _logger.error("AnalFinalizer failed with %s %s" % (errType,errValue))

    # finalize failed jobs
    _logger.debug("check stuck mergeing jobs")
    try:
        timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=2)
        # get PandaIDs
        varMap = {}
        varMap[':prodSourceLabel'] = 'managed'
        varMap[':jobStatus'] = 'merging'
        varMap[':timeLimit'] = timeLimit
        sql  = "SELECT distinct jediTaskID FROM ATLAS_PANDA.jobsActive4 "
        sql += "WHERE prodSourceLabel=:prodSourceLabel AND jobStatus=:jobStatus and modificationTime<:timeLimit "
        tmp,res = taskBuffer.querySQLS(sql,varMap)
        checkedDS = set()
        for jediTaskID, in res:
            varMap = {}
            varMap[':jediTaskID'] = jediTaskID
            varMap[':dsType'] = 'trn_log'
            sql = "SELECT datasetID FROM ATLAS_PANDA.JEDI_Datasets WHERE jediTaskID=:jediTaskID AND type=:dsType AND nFilesUsed=nFilesTobeUsed "
            tmpP,resD = taskBuffer.querySQLS(sql,varMap)
            for datasetID, in resD:
                varMap = {}
                varMap[':jediTaskID'] = jediTaskID
                varMap[':fileStatus'] = 'ready'
                varMap[':datasetID'] = datasetID
                sql  = "SELECT PandaID FROM ATLAS_PANDA.JEDI_Dataset_Contents "
                sql += "WHERE jediTaskID=:jediTaskID AND datasetid=:datasetID AND status=:fileStatus AND PandaID=OutPandaID AND rownum<=1 "
                tmpP,resP = taskBuffer.querySQLS(sql,varMap)
                if resP == []:
                    continue
                PandaID = resP[0][0]
                varMap = {}
                varMap[':PandaID'] = PandaID
                varMap[':fileType'] = 'log'
                sql = "SELECT d.status FROM ATLAS_PANDA.filesTable4 f,ATLAS_PANDA.datasets d WHERE PandaID=:PandaID AND f.type=:fileType AND d.name=f.destinationDBlock "
                tmpS,resS = taskBuffer.querySQLS(sql,varMap)
                if resS is not None:
                    subStatus, = resS[0]
                    if subStatus in ['completed']:
                        jobSpecs = taskBuffer.peekJobs([PandaID],fromDefined=False,fromArchived=False,fromWaiting=False)
                        jobSpec = jobSpecs[0]
                        subDsNames = set()
                        subDsList = []
                        for tmpFileSpec in jobSpec.Files:
                            if tmpFileSpec.type in ['log','output'] and \
                                    re.search('_sub\d+$',tmpFileSpec.destinationDBlock) is not None:
                                if tmpFileSpec.destinationDBlock in subDsNames:
                                    continue
                                subDsNames.add(tmpFileSpec.destinationDBlock)
                                datasetSpec = taskBuffer.queryDatasetWithMap({'name':tmpFileSpec.destinationDBlock})
                                subDsList.append(datasetSpec)
                        _logger.debug("update unmerged datasets for jediTaskID={0} PandaID={1}".format(jediTaskID,PandaID))
                        taskBuffer.updateUnmergedDatasets(jobSpec,subDsList,updateCompleted=True)
    except Exception:
        errType,errValue = sys.exc_info()[:2]
        _logger.error("check for stuck merging jobs failed with %s %s" % (errType,errValue))

    # get sites to skip various timeout
    varMap = {}
    varMap[':status'] = 'paused'
    sql = "SELECT siteid FROM ATLAS_PANDAMETA.schedconfig WHERE status=:status "
    sitesToSkipTO = set()
    status,res = taskBuffer.querySQLS(sql,varMap)
    for siteid, in res:
        sitesToSkipTO.add(siteid)
    _logger.debug("PQs to skip timeout : {0}".format(','.join(sitesToSkipTO)))

    sitesToDisableReassign = set()
    # get sites to disable reassign
    for siteName in siteMapper.siteSpecList:
        siteSpec = siteMapper.siteSpecList[siteName]
        if siteSpec.capability == 'ucore' and not siteSpec.is_unified:
            continue
        if siteSpec.disable_reassign():
            sitesToDisableReassign.add(siteName)
    _logger.debug("PQs to disable reassign : {0}".format(','.join(sitesToDisableReassign)))

    _memoryCheck("watcher")

    _logger.debug("Watcher session")

    # get the list of workflows
    sql = "SELECT DISTINCT workflow FROM ATLAS_PANDAMETA.schedconfig WHERE status='online' "
    status, res = taskBuffer.querySQLS(sql, {})
    workflow_timeout_map = {}
    for workflow, in res + [('production',), ('analysis',)]:
        timeout = taskBuffer.getConfigValue('watcher', 'HEARTBEAT_TIMEOUT_{0}'.format(workflow), 'pandaserver', 'atlas')
        if timeout is not None:
            workflow_timeout_map[workflow] = timeout
        elif workflow in ['production', 'analysis']:
            workflow_timeout_map[workflow] = 2

    workflows = list(workflow_timeout_map)

    _logger.debug("timeout : {0}".format(str(workflow_timeout_map)))

    # check heartbeat for analysis jobs
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=workflow_timeout_map['analysis'])
    varMap = {}
    varMap[':modificationTime'] = timeLimit
    varMap[':prodSourceLabel1'] = 'panda'
    varMap[':prodSourceLabel2'] = 'user'
    varMap[':jobStatus1'] = 'running'
    varMap[':jobStatus2'] = 'starting'
    varMap[':jobStatus3'] = 'stagein'
    varMap[':jobStatus4'] = 'stageout'
    sql  = "SELECT PandaID FROM ATLAS_PANDA.jobsActive4 WHERE (prodSourceLabel=:prodSourceLabel1 OR prodSourceLabel=:prodSourceLabel2) "
    sql += "AND jobStatus IN (:jobStatus1,:jobStatus2,:jobStatus3,:jobStatus4) AND modificationTime<:modificationTime"
    status,res = taskBuffer.querySQLS(sql,varMap)
    if res is None:
        _logger.debug("# of Anal Watcher : %s" % res)
    else:
        _logger.debug("# of Anal Watcher : %s" % len(res))
        for (id,) in res:
            _logger.debug("Anal Watcher %s" % id)
            thr = Watcher(taskBuffer,id,single=True,sleepTime=60,sitemapper=siteMapper)
            thr.start()
            thr.join()

    # check heartbeat for analysis jobs in transferring
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=workflow_timeout_map['analysis'])
    varMap = {}
    varMap[':modificationTime'] = timeLimit
    varMap[':prodSourceLabel1'] = 'panda'
    varMap[':prodSourceLabel2'] = 'user'
    varMap[':jobStatus1'] = 'transferring'
    sql  = "SELECT PandaID FROM ATLAS_PANDA.jobsActive4 WHERE prodSourceLabel IN (:prodSourceLabel1,:prodSourceLabel2) "
    sql += "AND jobStatus=:jobStatus1 AND modificationTime<:modificationTime"
    status,res = taskBuffer.querySQLS(sql,varMap)
    if res is None:
        _logger.debug("# of Transferring Anal Watcher : %s" % res)
    else:
        _logger.debug("# of Transferring Anal Watcher : %s" % len(res))
        for (id,) in res:
            _logger.debug("Trans Anal Watcher %s" % id)
            thr = Watcher(taskBuffer,id,single=True,sleepTime=60,sitemapper=siteMapper)
            thr.start()
            thr.join()

    # check heartbeat for sent jobs
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(minutes=30)
    varMap = {}
    varMap[':jobStatus'] = 'sent'
    varMap[':modificationTime'] = timeLimit
    status,res = taskBuffer.querySQLS("SELECT PandaID FROM ATLAS_PANDA.jobsActive4 WHERE jobStatus=:jobStatus AND modificationTime<:modificationTime",
                                  varMap)
    if res is None:
        _logger.debug("# of Sent Watcher : %s" % res)
    else:
        _logger.debug("# of Sent Watcher : %s" % len(res))
        for (id,) in res:
            _logger.debug("Sent Watcher %s" % id)
            thr = Watcher(taskBuffer,id,single=True,sleepTime=30,sitemapper=siteMapper)
            thr.start()
            thr.join()

    # check heartbeat for 'holding' analysis/ddm jobs
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=3)
    # get XMLs
    xmlIDs = set()
    # xmlFiles = os.listdir(panda_config.logdir)
    # for file in xmlFiles:
    #     match = re.search('^(\d+)_([^_]+)_.{36}$',file)
    #     if match is not None:
    #         id = match.group(1)
    #         xmlIDs.append(int(id))
    job_output_report_list = taskBuffer.listJobOutputReport()
    if job_output_report_list is not None:
        for panda_id, job_status, attempt_nr, time_stamp in job_output_report_list:
            xmlIDs.add(int(panda_id))
    sql = "SELECT PandaID FROM ATLAS_PANDA.jobsActive4 WHERE jobStatus=:jobStatus AND (modificationTime<:modificationTime OR (endTime IS NOT NULL AND endTime<:endTime)) AND (prodSourceLabel=:prodSourceLabel1 OR prodSourceLabel=:prodSourceLabel2 OR prodSourceLabel=:prodSourceLabel3) AND stateChangeTime != modificationTime"
    varMap = {}
    varMap[':modificationTime'] = timeLimit
    varMap[':endTime'] = timeLimit
    varMap[':jobStatus'] = 'holding'
    varMap[':prodSourceLabel1'] = 'panda'
    varMap[':prodSourceLabel2'] = 'user'
    varMap[':prodSourceLabel3'] = 'ddm'
    status,res = taskBuffer.querySQLS(sql,varMap)
    if res is None:
        _logger.debug("# of Holding Anal/DDM Watcher : %s" % res)
    else:
        _logger.debug("# of Holding Anal/DDM Watcher : %s - XMLs : %s" % (len(res),len(xmlIDs)))
        for (id,) in res:
            _logger.debug("Holding Anal/DDM Watcher %s" % id)
            if int(id) in xmlIDs:
                _logger.debug("   found XML -> skip %s" % id)
                continue
            thr = Watcher(taskBuffer,id,single=True,sleepTime=180,sitemapper=siteMapper)
            thr.start()
            thr.join()

    # check heartbeat for high prio production jobs
    timeOutVal = 3
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=timeOutVal)
    sql  = "SELECT PandaID FROM ATLAS_PANDA.jobsActive4 WHERE jobStatus=:jobStatus AND currentPriority>:pLimit "
    sql += "AND (modificationTime<:modificationTime OR (endTime IS NOT NULL AND endTime<:endTime))"
    varMap = {}
    varMap[':modificationTime'] = timeLimit
    varMap[':endTime'] = timeLimit
    varMap[':jobStatus'] = 'holding'
    varMap[':pLimit'] = 800
    status,res = taskBuffer.querySQLS(sql,varMap)
    if res is None:
        _logger.debug("# of High prio Holding Watcher : %s" % res)
    else:
        _logger.debug("# of High prio Holding Watcher : %s" % len(res))
        for (id,) in res:
            _logger.debug("High prio Holding Watcher %s" % id)
            thr = Watcher(taskBuffer,id,single=True,sleepTime=60*timeOutVal,sitemapper=siteMapper)
            thr.start()
            thr.join()

    # check heartbeat for production jobs
    timeOutVal = 48
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=timeOutVal)
    sql = "SELECT PandaID FROM ATLAS_PANDA.jobsActive4 WHERE jobStatus=:jobStatus AND (modificationTime<:modificationTime OR (endTime IS NOT NULL AND endTime<:endTime))"
    varMap = {}
    varMap[':modificationTime'] = timeLimit
    varMap[':endTime'] = timeLimit
    varMap[':jobStatus'] = 'holding'
    status,res = taskBuffer.querySQLS(sql,varMap)
    if res is None:
        _logger.debug("# of Holding Watcher : %s" % res)
    else:
        _logger.debug("# of Holding Watcher : %s" % len(res))
        for (id,) in res:
            _logger.debug("Holding Watcher %s" % id)
            thr = Watcher(taskBuffer,id,single=True,sleepTime=60*timeOutVal,sitemapper=siteMapper)
            thr.start()
            thr.join()

    # check heartbeat for production jobs with internal stage-out
    sql = "SELECT PandaID,jobStatus,jobSubStatus FROM ATLAS_PANDA.jobsActive4 j,ATLAS_PANDAMETA.schedconfig s "
    sql += "WHERE j.computingSite=s.siteid AND jobStatus=:jobStatus1 AND jobSubStatus IS NOT NULL AND modificationTime<:modificationTime "
    for workflow in workflows:
        if workflow == 'analysis':
            continue
        varMap = {}
        varMap[':modificationTime'] = timeLimit
        varMap[':jobStatus1'] = 'transferring'
        sqlX = sql
        if workflow == 'production':
            if len(workflows) > 2:
                sqlX += "AND (s.workflow IS NULL OR s.workflow NOT IN ("
                for ng_workflow in workflows:
                    if ng_workflow in ['production', 'analysis']:
                        continue
                    tmp_key = ':w_{0}'.format(ng_workflow)
                    varMap[tmp_key] = ng_workflow
                    sqlX += '{0},'.format(tmp_key)
                sqlX = sqlX[:-1]
                sqlX += ")) "
        else:
            tmp_key = ':w_{0}'.format(workflow)
            sqlX += "AND s.workflow={0} ".format(tmp_key)
            varMap[tmp_key] = workflow
        timeOutVal = workflow_timeout_map[workflow]
        timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=timeOutVal)
        varMap[':modificationTime'] = timeLimit
        status,res = taskBuffer.querySQLS(sqlX, varMap)
        if res is None:
            _logger.debug("# of Internal Staging Watcher with workflow={0}: {1}".format(workflow, res))
        else:
            _logger.debug("# of Internal Staging Watcher with workflow={0}: {1}".format(workflow, len(res)))
            for pandaID, jobStatus, jobSubStatus in res:
                _logger.debug("Internal Staging  Watcher %s %s:%s" % (pandaID, jobStatus, jobSubStatus))
                thr = Watcher(taskBuffer,pandaID,single=True,sleepTime=60*timeOutVal,sitemapper=siteMapper)
                thr.start()
                thr.join()

    # check heartbeat for production jobs
    sql = "SELECT PandaID,jobStatus,j.computingSite FROM ATLAS_PANDA.jobsActive4 j, ATLAS_PANDAMETA.schedconfig s "
    sql += "WHERE j.computingSite=s.siteid AND jobStatus IN (:jobStatus1,:jobStatus2,:jobStatus3,:jobStatus4) AND modificationTime<:modificationTime "
    for workflow in workflows:
        if workflow == 'analysis':
            continue
        varMap = {}
        varMap[':jobStatus1'] = 'running'
        varMap[':jobStatus2'] = 'starting'
        varMap[':jobStatus3'] = 'stagein'
        varMap[':jobStatus4'] = 'stageout'
        sqlX = sql
        if workflow == 'production':
            if len(workflows) > 2:
                sqlX += "AND (s.workflow IS NULL OR s.workflow NOT IN ("
                for ng_workflow in workflows:
                    if ng_workflow in ['production', 'analysis']:
                        continue
                    tmp_key = ':w_{0}'.format(ng_workflow)
                    varMap[tmp_key] = ng_workflow
                    sqlX += '{0},'.format(tmp_key)
                sqlX = sqlX[:-1]
                sqlX += ")) "
        else:
            tmp_key = ':w_{0}'.format(workflow)
            sqlX += "AND s.workflow={0} ".format(tmp_key)
            varMap[tmp_key] = workflow
        timeOutVal = workflow_timeout_map[workflow]
        timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=timeOutVal)
        varMap[':modificationTime'] = timeLimit
        status,res = taskBuffer.querySQLS(sqlX, varMap)
        if res is None:
            _logger.debug("# of General Watcher with workflow={0}: {1}".format(workflow, res))
        else:
            _logger.debug("# of General Watcher with workflow={0}: {1}".format(workflow, len(res)))
            for pandaID,jobStatus,computingSite in res:
                if computingSite in sitesToSkipTO:
                    _logger.debug("skip General Watcher for PandaID={0} at {1} since timeout is disabled for {2}".format(pandaID,computingSite,jobStatus))
                    continue
                _logger.debug("General Watcher %s" % pandaID)
                thr = Watcher(taskBuffer,pandaID,single=True,sleepTime=60*timeOutVal,sitemapper=siteMapper)
                thr.start()
                thr.join()

    _memoryCheck("reassign")

    # kill long-waiting jobs in defined table
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(days=7)
    status,res = taskBuffer.querySQLS("SELECT PandaID,cloud,prodSourceLabel FROM ATLAS_PANDA.jobsDefined4 WHERE creationTime<:creationTime",
                                  {':creationTime':timeLimit})
    jobs=[]
    dashFileMap = {}
    if res is not None:
        for pandaID,cloud,prodSourceLabel in res:
            # collect PandaIDs
            jobs.append(pandaID)
    if len(jobs):
        _logger.debug("killJobs for Defined (%s)" % str(jobs))
        Client.killJobs(jobs,2)

    # kill long-waiting jobs in active table
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(days=7)
    varMap = {}
    varMap[':jobStatus'] = 'activated'
    varMap[':creationTime'] = timeLimit
    status,res = taskBuffer.querySQLS("SELECT PandaID from ATLAS_PANDA.jobsActive4 WHERE jobStatus=:jobStatus AND creationTime<:creationTime",
                                  varMap)
    jobs=[]
    if res is not None:
        for (id,) in res:
            jobs.append(id)
    if len(jobs):
        _logger.debug("killJobs for Active (%s)" % str(jobs))
        Client.killJobs(jobs,2)

    # kill long-waiting ddm jobs for dispatch
    _logger.debug("kill PandaMovers")
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=12)
    sql = "SELECT PandaID from ATLAS_PANDA.jobsActive4 WHERE prodSourceLabel=:prodSourceLabel AND transferType=:transferType AND creationTime<:creationTime"
    varMap = {}
    varMap[':creationTime']    = timeLimit
    varMap[':prodSourceLabel'] = 'ddm'
    varMap[':transferType']    = 'dis'
    _logger.debug(sql+str(varMap))
    status,res = taskBuffer.querySQLS(sql,varMap)
    _logger.debug(res)
    jobs=[]
    if res is not None:
        for (id,) in res:
            jobs.append(id)
    if len(jobs):
        _logger.debug("kill DDM Jobs (%s)" % str(jobs))
        Client.killJobs(jobs,2)

    # reassign activated jobs in inactive sites
    inactiveTimeLimitSite = 2
    inactiveTimeLimitJob  = 4
    inactivePrioLimit = 800
    timeLimitSite = datetime.datetime.utcnow() - datetime.timedelta(hours=inactiveTimeLimitSite)
    timeLimitJob  = datetime.datetime.utcnow() - datetime.timedelta(hours=inactiveTimeLimitJob)
    # get PandaIDs
    sql  = 'SELECT distinct computingSite FROM ATLAS_PANDA.jobsActive4 '
    sql += 'WHERE prodSourceLabel=:prodSourceLabel '
    sql += 'AND ((modificationTime<:timeLimit AND jobStatus=:jobStatus1) '
    sql += 'OR (stateChangeTime<:timeLimit AND jobStatus=:jobStatus2)) '
    sql += 'AND lockedby=:lockedby AND currentPriority>=:prioLimit '
    sql += 'AND NOT processingType IN (:pType1) AND relocationFlag<>:rFlag1 '
    varMap = {}
    varMap[':prodSourceLabel'] = 'managed'
    varMap[':jobStatus1'] = 'activated'
    varMap[':jobStatus2'] = 'starting'
    varMap[':lockedby']  = 'jedi'
    varMap[':timeLimit'] = timeLimitJob
    varMap[':prioLimit'] = inactivePrioLimit
    varMap[':pType1']    = 'pmerge'
    varMap[':rFlag1']    = 2
    stDS,resDS = taskBuffer.querySQLS(sql,varMap)
    sqlSS  = 'SELECT laststart FROM ATLAS_PANDAMETA.siteData '
    sqlSS += 'WHERE site=:site AND flag=:flag AND hours=:hours AND laststart<:laststart '
    sqlPI  = 'SELECT PandaID,eventService,attemptNr FROM ATLAS_PANDA.jobsActive4 '
    sqlPI += 'WHERE prodSourceLabel=:prodSourceLabel AND jobStatus IN (:jobStatus1,:jobStatus2) '
    sqlPI += 'AND (modificationTime<:timeLimit OR stateChangeTime<:timeLimit) '
    sqlPI += 'AND lockedby=:lockedby AND currentPriority>=:prioLimit '
    sqlPI += 'AND computingSite=:site AND NOT processingType IN (:pType1) AND relocationFlag<>:rFlag1 '
    for tmpSite, in resDS:
        if tmpSite in sitesToDisableReassign:
            _logger.debug('skip reassignJobs at inactive site %s since reassign is disabled' % (tmpSite))
            continue
        # check if the site is inactive
        varMap = {}
        varMap[':site']  = tmpSite
        varMap[':flag']  = 'production'
        varMap[':hours'] = 3
        varMap[':laststart'] = timeLimitSite
        stSS,resSS = taskBuffer.querySQLS(sqlSS,varMap)
        if stSS is not None and len(resSS) > 0:
            # get jobs
            varMap = {}
            varMap[':prodSourceLabel'] = 'managed'
            varMap[':jobStatus1'] = 'activated'
            varMap[':jobStatus2'] = 'starting'
            varMap[':lockedby']  = 'jedi'
            varMap[':timeLimit'] = timeLimitJob
            varMap[':prioLimit'] = inactivePrioLimit
            varMap[':site']      = tmpSite
            varMap[':pType1']    = 'pmerge'
            varMap[':rFlag1']    = 2
            stPI,resPI = taskBuffer.querySQLS(sqlPI,varMap)
            jediJobs = []
            # reassign
            _logger.debug('reassignJobs for JEDI at inactive site %s laststart=%s' % (tmpSite,resSS[0][0]))
            if resPI is not None:
                for pandaID, eventService, attemptNr in resPI:
                    if eventService in [EventServiceUtils.esMergeJobFlagNumber]:
                        _logger.debug('retrying es merge %s at inactive site %s' % (pandaID,tmpSite))
                        taskBuffer.retryJob(pandaID,{},getNewPandaID=True,attemptNr=attemptNr,
                                                     recoverableEsMerge=True)
                    jediJobs.append(pandaID)
            if len(jediJobs) != 0:
                nJob = 100
                iJob = 0
                while iJob < len(jediJobs):
                    _logger.debug('reassignJobs for JEDI at inactive site %s (%s)' % (tmpSite,jediJobs[iJob:iJob+nJob]))
                    Client.killJobs(jediJobs[iJob:iJob+nJob],51,keepUnmerged=True)
                    iJob += nJob

    # reassign defined jobs in defined table
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=4)
    # get PandaIDs
    status,res = taskBuffer.lockJobsForReassign("ATLAS_PANDA.jobsDefined4",timeLimit,['defined'],['managed'],[],[],[],
                                                True)
    jobs = []
    jediJobs = []
    if res is not None:
        for (id,lockedby) in res:
            if lockedby == 'jedi':
                jediJobs.append(id)
            else:
                jobs.append(id)
    # reassign
    _logger.debug('reassignJobs for defined jobs -> #%s' % len(jobs))
    if len(jobs) > 0:
        nJob = 100
        iJob = 0
        while iJob < len(jobs):
            _logger.debug('reassignJobs for defined jobs (%s)' % jobs[iJob:iJob+nJob])
            taskBuffer.reassignJobs(jobs[iJob:iJob+nJob],joinThr=True)
            iJob += nJob
    _logger.debug('reassignJobs for JEDI defined jobs -> #%s' % len(jediJobs))
    if len(jediJobs) != 0:
        nJob = 100
        iJob = 0
        while iJob < len(jediJobs):
            _logger.debug('reassignJobs for JEDI defined jobs (%s)' % jediJobs[iJob:iJob+nJob])
            Client.killJobs(jediJobs[iJob:iJob+nJob],51,keepUnmerged=True)
            iJob += nJob

    # reassign long-waiting jobs in defined table
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=12)
    status,res = taskBuffer.lockJobsForReassign("ATLAS_PANDA.jobsDefined4",timeLimit,[],['managed'],[],[],[],
                                                True)
    jobs = []
    jediJobs = []
    if res is not None:
        for (id,lockedby) in res:
            if lockedby == 'jedi':
                jediJobs.append(id)
            else:
                jobs.append(id)
    # reassign
    _logger.debug('reassignJobs for long in defined table -> #%s' % len(jobs))
    if len(jobs) > 0:
        nJob = 100
        iJob = 0
        while iJob < len(jobs):
            _logger.debug('reassignJobs for long in defined table (%s)' % jobs[iJob:iJob+nJob])
            taskBuffer.reassignJobs(jobs[iJob:iJob+nJob],joinThr=True)
            iJob += nJob
    _logger.debug('reassignJobs for long JEDI in defined table -> #%s' % len(jediJobs))
    if len(jediJobs) != 0:
        nJob = 100
        iJob = 0
        while iJob < len(jediJobs):
            _logger.debug('reassignJobs for long JEDI in defined table (%s)' % jediJobs[iJob:iJob+nJob])
            Client.killJobs(jediJobs[iJob:iJob+nJob],51,keepUnmerged=True)
            iJob += nJob

    # reassign too long-standing evgen/simul jobs with active state at T1
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=6)
    for tmpCloud in siteMapper.getCloudList():
        # ignore special clouds
        if tmpCloud in ['CERN','OSG']:
            continue
        status,res = taskBuffer.lockJobsForReassign("ATLAS_PANDA.jobsActive4",timeLimit,['activated'],['managed'],
                                                    ['evgen','simul'],[siteMapper.getCloud(tmpCloud)['tier1']],[],
                                                    True,onlyReassignable=True)
        jobs = []
        jediJobs = []
        if res is not None:
            for (id,lockedby) in res:
                if lockedby == 'jedi':
                    jediJobs.append(id)
                else:
                    jobs.append(id)
        _logger.debug('reassignJobs for Active T1 evgensimul in %s -> #%s' % (tmpCloud,len(jobs)))
        if len(jobs) != 0:
            nJob = 100
            iJob = 0
            while iJob < len(jobs):
                _logger.debug('reassignJobs for Active T1 evgensimul (%s)' % jobs[iJob:iJob+nJob])
                taskBuffer.reassignJobs(jobs[iJob:iJob+nJob],joinThr=True)
                iJob += nJob
        _logger.debug('reassignJobs for Active T1 JEDI evgensimul in %s -> #%s' % (tmpCloud,len(jediJobs)))
        if len(jediJobs) != 0:
            nJob = 100
            iJob = 0
            while iJob < len(jediJobs):
                _logger.debug('reassignJobs for Active T1 JEDI evgensimul (%s)' % jediJobs[iJob:iJob+nJob])
                Client.killJobs(jediJobs[iJob:iJob+nJob],51,keepUnmerged=True)
                iJob += nJob

    # reassign too long-standing evgen/simul jobs with active state at T2
    try:
        _logger.debug('looking for stuck T2s to reassign evgensimul')
        timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=6)
        varMap = {}
        varMap[':jobStatus1']       = 'activated'
        varMap[':jobStatus2']       = 'running'
        varMap[':prodSourceLabel']  = 'managed'
        varMap[':processingType1']  = 'evgen'
        varMap[':processingType2']  = 'simul'
        sql = "SELECT cloud,computingSite,jobStatus,COUNT(*) FROM ATLAS_PANDA.jobsActive4 "\
              "WHERE jobStatus IN (:jobStatus1,:jobStatus2) AND prodSourceLabel=:prodSourceLabel "\
              "AND processingType IN (:processingType1,:processingType2) GROUP BY cloud,computingSite,jobStatus "
        status,res = taskBuffer.querySQLS(sql, varMap)
        if res is not None:
            # get ratio of activated/running
            siteStatData = {}
            for tmpCloud,tmpComputingSite,tmpJobStatus,tmpCount in res:
                # skip T1
                if tmpComputingSite == siteMapper.getCloud(tmpCloud)['tier1']:
                    continue
                # skip if reassign is disabled
                if tmpComputingSite in sitesToDisableReassign:
                    continue
                # add cloud/site
                tmpKey = (tmpCloud,tmpComputingSite)
                if tmpKey not in siteStatData:
                    siteStatData[tmpKey] = {'activated':0,'running':0}
                # add the number of jobs
                if tmpJobStatus in siteStatData[tmpKey]:
                    siteStatData[tmpKey][tmpJobStatus] += tmpCount
            # look for stuck site
            stuckThr = 10
            stuckSites = []
            for tmpKey in siteStatData:
                tmpStatData = siteStatData[tmpKey]
                if tmpStatData['running'] == 0 or \
                       float(tmpStatData['activated'])/float(tmpStatData['running']) > stuckThr:
                    tmpCloud,tmpComputingSite = tmpKey
                    _logger.debug(' %s:%s %s/%s > %s' % (tmpCloud,tmpComputingSite,tmpStatData['activated'],tmpStatData['running'],stuckThr))
                    # get stuck jobs
                    status,res = taskBuffer.lockJobsForReassign("ATLAS_PANDA.jobsActive4",timeLimit,['activated'],['managed'],
                                                                ['evgen','simul'],[tmpComputingSite],[tmpCloud],True,
                                                                onlyReassignable=True)
                    jobs = []
                    jediJobs = []
                    if res is not None:
                        for (id,lockedby) in res:
                            if lockedby == 'jedi':
                                jediJobs.append(id)
                            else:
                                jobs.append(id)
                    _logger.debug('reassignJobs for Active T2 evgensimul %s:%s -> #%s' % (tmpCloud,tmpComputingSite,len(jobs)))
                    if len(jobs) > 0:
                        nJob = 100
                        iJob = 0
                        while iJob < len(jobs):
                            _logger.debug('reassignJobs for Active T2 evgensimul (%s)' % jobs[iJob:iJob+nJob])
                            taskBuffer.reassignJobs(jobs[iJob:iJob+nJob],joinThr=True)
                            iJob += nJob
                    _logger.debug('reassignJobs for Active T2 JEDI evgensimul %s:%s -> #%s' % (tmpCloud,tmpComputingSite,len(jediJobs)))
                    if len(jediJobs) > 0:
                        nJob = 100
                        iJob = 0
                        while iJob < len(jediJobs):
                            _logger.debug('reassignJobs for Active T2 JEDI evgensimul (%s)' % jediJobs[iJob:iJob+nJob])
                            Client.killJobs(jediJobs[iJob:iJob+nJob],51,keepUnmerged=True)
                            iJob += nJob
    except Exception:
        errType,errValue = sys.exc_info()[:2]
        _logger.error("failed to reassign T2 evgensimul with %s:%s" % (errType,errValue))

    # reassign too long activated jobs in active table
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(days=2)
    status,res = taskBuffer.lockJobsForReassign("ATLAS_PANDA.jobsActive4",timeLimit,['activated'],['managed'],[],[],[],True,
                                                onlyReassignable=True,getEventService=True)
    jobs = []
    jediJobs = []
    if res is not None:
        for pandaID, lockedby, eventService, attemptNr, computingSite in res:
            if computingSite in sitesToDisableReassign:
                _logger.debug('skip reassignJobs for long activated PandaID={0} since disabled at {1}'.format(pandaID,computingSite))
                continue
            if lockedby == 'jedi':
                if eventService in [EventServiceUtils.esMergeJobFlagNumber]:
                    _logger.debug('retrying {0} in long activated' % pandaID)
                    taskBuffer.retryJob(pandaID,{},getNewPandaID=True,attemptNr=attemptNr,
                                        recoverableEsMerge=True)
                jediJobs.append(pandaID)
            else:
                jobs.append(pandaID)
    _logger.debug('reassignJobs for long activated in active table -> #%s' % len(jobs))
    if len(jobs) != 0:
        nJob = 100
        iJob = 0
        while iJob < len(jobs):
            _logger.debug('reassignJobs for long activated in active table (%s)' % jobs[iJob:iJob+nJob])
            taskBuffer.reassignJobs(jobs[iJob:iJob+nJob],joinThr=True)
            iJob += nJob
    _logger.debug('reassignJobs for long activated JEDI in active table -> #%s' % len(jediJobs))
    if len(jediJobs) != 0:
        nJob = 100
        iJob = 0
        while iJob < len(jediJobs):
            _logger.debug('reassignJobs for long activated JEDI in active table (%s)' % jediJobs[iJob:iJob+nJob])
            Client.killJobs(jediJobs[iJob:iJob+nJob],51,keepUnmerged=True)
            iJob += nJob

    # reassign too long starting jobs in active table
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=48)
    status,res = taskBuffer.lockJobsForReassign("ATLAS_PANDA.jobsActive4",timeLimit,['starting'],['managed'],[],[],[],True,
                                                onlyReassignable=True,useStateChangeTime=True,getEventService=True)
    jobs = []
    jediJobs = []
    if res is not None:
        for pandaID, lockedby, eventService, attemptNr, computingSite in res:
            if computingSite in sitesToDisableReassign:
                _logger.debug('skip reassignJobs for long starting PandaID={0} since disabled at {1}'.format(pandaID,computingSite))
                continue
            if lockedby == 'jedi':
                jediJobs.append(pandaID)
            else:
                jobs.append(pandaID)
    _logger.debug('reassignJobs for long starting in active table -> #%s' % len(jobs))
    if len(jobs) != 0:
        nJob = 100
        iJob = 0
        while iJob < len(jobs):
            _logger.debug('reassignJobs for long starting in active table (%s)' % jobs[iJob:iJob+nJob])
            taskBuffer.reassignJobs(jobs[iJob:iJob+nJob],joinThr=True)
            iJob += nJob
    _logger.debug('reassignJobs for long starting JEDI in active table -> #%s' % len(jediJobs))
    if len(jediJobs) != 0:
        nJob = 100
        iJob = 0
        while iJob < len(jediJobs):
            _logger.debug('reassignJobs for long stating JEDI in active table (%s)' % jediJobs[iJob:iJob+nJob])
            Client.killJobs(jediJobs[iJob:iJob+nJob],51,keepUnmerged=True)
            iJob += nJob

    # kill too long-standing analysis jobs in active table
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(days=7)
    varMap = {}
    varMap[':prodSourceLabel1'] = 'test'
    varMap[':prodSourceLabel2'] = 'panda'
    varMap[':prodSourceLabel3'] = 'user'
    varMap[':modificationTime'] = timeLimit
    status,res = taskBuffer.querySQLS("SELECT PandaID FROM ATLAS_PANDA.jobsActive4 WHERE (prodSourceLabel=:prodSourceLabel1 OR prodSourceLabel=:prodSourceLabel2 OR prodSourceLabel=:prodSourceLabel3) AND modificationTime<:modificationTime ORDER BY PandaID",
                                  varMap)
    jobs = []
    if res is not None:
        for (id,) in res:
            jobs.append(id)
    # kill
    if len(jobs):
        Client.killJobs(jobs,2)
        _logger.debug("killJobs for Anal Active (%s)" % str(jobs))

    # kill too long pending jobs
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=1)
    varMap = {}
    varMap[':jobStatus']    = 'pending'
    varMap[':creationTime'] = timeLimit
    status,res = taskBuffer.querySQLS("SELECT PandaID FROM ATLAS_PANDA.jobsWaiting4 WHERE jobStatus=:jobStatus AND creationTime<:creationTime",
                                      varMap)
    jobs = []
    if res is not None:
        for (id,) in res:
            jobs.append(id)
    # kill
    if len(jobs):
        if len(jobs):
            nJob = 100
            iJob = 0
            while iJob < len(jobs):
                _logger.debug("killJobs for Pending (%s)" % str(jobs[iJob:iJob+nJob]))
                Client.killJobs(jobs[iJob:iJob+nJob],4)
                iJob += nJob

    # kick waiting ES merge jobs which were generated from fake co-jumbo
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(minutes=10)
    varMap = {}
    varMap[':jobStatus']    = 'waiting'
    varMap[':creationTime'] = timeLimit
    varMap[':esMerge'] = EventServiceUtils.esMergeJobFlagNumber
    sql  = "SELECT PandaID,computingSite FROM ATLAS_PANDA.jobsWaiting4 WHERE jobStatus=:jobStatus AND creationTime<:creationTime "
    sql += "AND eventService=:esMerge ORDER BY jediTaskID "
    status,res = taskBuffer.querySQLS(sql, varMap)
    jobsMap = {}
    if res is not None:
        for id,site in res:
            if site not in jobsMap:
                jobsMap[site] = []
            jobsMap[site].append(id)
    # kick
    if len(jobsMap):
        for site in jobsMap:
            jobs = jobsMap[site]
            nJob = 100
            iJob = 0
            while iJob < len(jobs):
                _logger.debug("kick waiting ES merge (%s)" % str(jobs[iJob:iJob+nJob]))
                Client.reassignJobs(jobs[iJob:iJob+nJob], )
                iJob += nJob

    # kill too long waiting jobs
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=1)
    varMap = {}
    varMap[':jobStatus']    = 'waiting'
    varMap[':creationTime'] = timeLimit
    varMap[':coJumbo'] = EventServiceUtils.coJumboJobFlagNumber
    sql  = "SELECT PandaID FROM ATLAS_PANDA.jobsWaiting4 WHERE jobStatus=:jobStatus AND creationTime<:creationTime "
    sql += "AND (eventService IS NULL OR eventService<>:coJumbo) "
    status,res = taskBuffer.querySQLS(sql, varMap)
    jobs = []
    if res is not None:
        for (id,) in res:
            jobs.append(id)
    # kill
    if len(jobs):
        if len(jobs):
            nJob = 100
            iJob = 0
            while iJob < len(jobs):
                _logger.debug("killJobs for Waiting (%s)" % str(jobs[iJob:iJob+nJob]))
                Client.killJobs(jobs[iJob:iJob+nJob],4)
                iJob += nJob

    # kill too long running ES jobs
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=24)
    varMap = {}
    varMap[':jobStatus1']  = 'running'
    varMap[':jobStatus2']  = 'starting'
    varMap[':timeLimit'] = timeLimit
    varMap[':esJob'] = EventServiceUtils.esJobFlagNumber
    varMap[':coJumbo'] = EventServiceUtils.coJumboJobFlagNumber
    sql  = "SELECT PandaID FROM ATLAS_PANDA.jobsActive4 WHERE jobStatus IN (:jobStatus1,:jobStatus2) AND stateChangeTime<:timeLimit "
    sql += "AND eventService IN (:esJob,:coJumbo) AND currentPriority>=900 "
    status,res = taskBuffer.querySQLS(sql, varMap)
    jobs = []
    if res is not None:
        for (id,) in res:
            jobs.append(id)
    # kill
    if len(jobs):
        nJob = 100
        iJob = 0
        while iJob < len(jobs):
            _logger.debug("killJobs for long running ES jobs (%s)" % str(jobs[iJob:iJob+nJob]))
            Client.killJobs(jobs[iJob:iJob+nJob], 2, keepUnmerged=True, jobSubStatus='es_toolong')
            iJob += nJob

    # kill too long running ES merge jobs
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=24)
    varMap = {}
    varMap[':jobStatus1']  = 'running'
    varMap[':jobStatus2']  = 'starting'
    varMap[':timeLimit'] = timeLimit
    varMap[':esMergeJob'] = EventServiceUtils.esMergeJobFlagNumber
    sql  = "SELECT PandaID FROM ATLAS_PANDA.jobsActive4 WHERE jobStatus IN (:jobStatus1,:jobStatus2) AND stateChangeTime<:timeLimit "
    sql += "AND eventService=:esMergeJob "
    status,res = taskBuffer.querySQLS(sql, varMap)
    jobs = []
    if res is not None:
        for (id,) in res:
            jobs.append(id)
    # kill
    if len(jobs):
        nJob = 100
        iJob = 0
        while iJob < len(jobs):
            _logger.debug("killJobs for long running ES merge jobs (%s)" % str(jobs[iJob:iJob+nJob]))
            Client.killJobs(jobs[iJob:iJob+nJob], 2)
            iJob += nJob

    # kill too long waiting jobs
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(days=7)
    sql = "SELECT PandaID FROM ATLAS_PANDA.jobsWaiting4 WHERE ((creationTime<:timeLimit AND (eventService IS NULL OR eventService<>:coJumbo)) "
    sql += "OR modificationTime<:timeLimit) "
    varMap = {}
    varMap[':timeLimit'] = timeLimit
    varMap[':coJumbo'] = EventServiceUtils.coJumboJobFlagNumber
    status,res = taskBuffer.querySQLS(sql, varMap)
    jobs = []
    if res is not None:
        for (id,) in res:
            jobs.append(id)
    # kill
    if len(jobs):
        Client.killJobs(jobs,4)
        _logger.debug("killJobs in jobsWaiting (%s)" % str(jobs))

    # rebrokerage
    _logger.debug("Rebrokerage start")

    # get timeout value
    timeoutVal = taskBuffer.getConfigValue('rebroker','ANALY_TIMEOUT')
    if timeoutVal is None:
        timeoutVal = 12
    _logger.debug("timeout value : {0}h".format(timeoutVal))
    try:
        normalTimeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=timeoutVal)
        sortTimeLimit   = datetime.datetime.utcnow() - datetime.timedelta(hours=3)
        sql = "WITH p AS ("\
              "SELECT MIN(PandaID) PandaID,jobDefinitionID,prodUserName,prodUserID,computingSite,jediTaskID,processingType "\
              "FROM ATLAS_PANDA.jobsActive4 "\
              "WHERE prodSourceLabel IN (:prodSourceLabel1,:prodSourceLabel2) "\
              "AND jobStatus IN (:jobStatus1,:jobStatus2,:jobStatus3) "\
              "AND jobsetID IS NOT NULL AND lockedBy=:lockedBy "\
              "GROUP BY jobDefinitionID,prodUserName,prodUserID,computingSite,jediTaskID,processingType "\
              ") "\
              "SELECT /*+ INDEX (s JOBS_STATUSLOG_PANDAID_IDX) */ "\
              "p.jobDefinitionID,p.prodUserName,p.prodUserID,p.computingSite,s.modificationTime,p.jediTaskID,p.processingType " \
              "FROM p, ATLAS_PANDA.jobs_statuslog s "\
              "WHERE s.PandaID=p.PandaID AND s.jobStatus=:s_jobStatus AND s.modificationTime<:modificationTime "
        varMap = {}
        varMap[':prodSourceLabel1'] = 'user'
        varMap[':prodSourceLabel2'] = 'panda'
        varMap[':modificationTime'] = sortTimeLimit
        varMap[':lockedBy']         = 'jedi'
        varMap[':jobStatus1']       = 'activated'
        varMap[':jobStatus2']       = 'dummy'
        varMap[':jobStatus3']       = 'starting'
        varMap[':s_jobStatus'] = 'activated'
        # get jobs older than threshold
        ret,res = taskBuffer.querySQLS(sql, varMap)
        resList = []
        keyList = set()
        if res is not None:
            for tmpItem in res:
                jobDefinitionID,prodUserName,prodUserID,computingSite,maxTime,jediTaskID,processingType = tmpItem
                tmpKey = (jediTaskID,jobDefinitionID)
                keyList.add(tmpKey)
                resList.append(tmpItem)
        # get stalled assigned job
        sqlA  = "SELECT jobDefinitionID,prodUserName,prodUserID,computingSite,MAX(creationTime),jediTaskID,processingType "
        sqlA += "FROM ATLAS_PANDA.jobsDefined4 "
        sqlA += "WHERE prodSourceLabel IN (:prodSourceLabel1,:prodSourceLabel2) AND jobStatus IN (:jobStatus1,:jobStatus2) "
        sqlA += "AND creationTime<:modificationTime AND lockedBy=:lockedBy "
        sqlA += "GROUP BY jobDefinitionID,prodUserName,prodUserID,computingSite,jediTaskID,processingType "
        varMap = {}
        varMap[':prodSourceLabel1'] = 'user'
        varMap[':prodSourceLabel2'] = 'panda'
        varMap[':modificationTime'] = sortTimeLimit
        varMap[':lockedBy']         = 'jedi'
        varMap[':jobStatus1']       = 'assigned'
        varMap[':jobStatus2']       = 'defined'
        retA,resA = taskBuffer.querySQLS(sqlA, varMap)
        if resA is not None:
            for tmpItem in resA:
                jobDefinitionID,prodUserName,prodUserID,computingSite,maxTime,jediTaskID,processingType = tmpItem
                tmpKey = (jediTaskID,jobDefinitionID)
                if tmpKey not in keyList:
                    keyList.add(tmpKey)
                    resList.append(tmpItem)
        # sql to check recent activity
        sql  = "SELECT PandaID,stateChangeTime,jobStatus FROM %s "
        sql += "WHERE prodUserName=:prodUserName AND jobDefinitionID=:jobDefinitionID "
        sql += "AND computingSite=:computingSite AND jediTaskID=:jediTaskID "
        sql += "AND jobStatus NOT IN (:jobStatus1,:jobStatus2,:jobStatus3) "
        sql += "AND stateChangeTime>:modificationTime "
        sql += "AND rownum <= 1"
        # sql to get associated jobs with jediTaskID
        sqlJJ  = "SELECT PandaID FROM %s "
        sqlJJ += "WHERE jediTaskID=:jediTaskID AND jobStatus IN (:jobS1,:jobS2,:jobS3,:jobS4,:jobS5) "
        sqlJJ += "AND jobDefinitionID=:jobDefID AND computingSite=:computingSite "
        if resList != []:
            recentRuntimeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=3)
            # loop over all user/jobID combinations
            iComb = 0
            nComb = len(resList)
            _logger.debug("total combinations = %s" % nComb)
            for jobDefinitionID,prodUserName,prodUserID,computingSite,maxModificationTime,jediTaskID,processingType in resList:
                # check if jobs with the jobID have run recently
                varMap = {}
                varMap[':jediTaskID']       = jediTaskID
                varMap[':computingSite']    = computingSite
                varMap[':prodUserName']     = prodUserName
                varMap[':jobDefinitionID']  = jobDefinitionID
                varMap[':modificationTime'] = recentRuntimeLimit
                varMap[':jobStatus1'] = 'closed'
                varMap[':jobStatus2'] = 'failed'
                varMap[':jobStatus3'] = 'starting'
                _logger.debug(" rebro:%s/%s:ID=%s:%s jediTaskID=%s site=%s" % (iComb,nComb,jobDefinitionID,
                                                                               prodUserName,jediTaskID,
                                                                               computingSite))
                iComb += 1
                hasRecentJobs = False
                # check site
                if not siteMapper.checkSite(computingSite):
                    _logger.debug("    -> skip unknown site=%s" % computingSite)
                    continue
                # check site status
                tmpSiteStatus = siteMapper.getSite(computingSite).status
                if tmpSiteStatus not in ['offline','test']:
                    # use normal time limit for normal site status
                    if maxModificationTime > normalTimeLimit:
                        _logger.debug("    -> skip wait for normal timelimit=%s<maxModTime=%s" % (normalTimeLimit,maxModificationTime))
                        continue
                    for tableName in ['ATLAS_PANDA.jobsActive4','ATLAS_PANDA.jobsArchived4']:
                        retU,resU = taskBuffer.querySQLS(sql % tableName, varMap)
                        if resU is None:
                            # database error
                            raise RuntimeError("failed to check modTime")
                        if resU != []:
                            # found recent jobs
                            hasRecentJobs = True
                            _logger.debug("    -> skip due to recent activity %s to %s at %s" % (resU[0][0],
                                                                                resU[0][2],
                                                                                resU[0][1]))
                            break
                else:
                    _logger.debug("    -> immediate rebro due to site status=%s" % tmpSiteStatus)
                if hasRecentJobs:
                    # skip since some jobs have run recently
                    continue
                else:
                    if jediTaskID is None:
                        _logger.debug("    -> rebro for normal task : no action")
                    else:
                        _logger.debug("    -> rebro for JEDI task")
                        killJobs = []
                        varMap = {}
                        varMap[':jediTaskID'] = jediTaskID
                        varMap[':jobDefID'] = jobDefinitionID
                        varMap[':computingSite'] = computingSite
                        varMap[':jobS1'] = 'defined'
                        varMap[':jobS2'] = 'assigned'
                        varMap[':jobS3'] = 'activated'
                        varMap[':jobS4'] = 'dummy'
                        varMap[':jobS5'] = 'starting'
                        for tableName in ['ATLAS_PANDA.jobsDefined4','ATLAS_PANDA.jobsActive4']:
                            retJJ,resJJ = taskBuffer.querySQLS(sqlJJ % tableName, varMap)
                            for tmpPandaID, in resJJ:
                                killJobs.append(tmpPandaID)
                        # reverse sort to kill buildJob in the end
                        killJobs.sort()
                        killJobs.reverse()
                        # kill to reassign
                        taskBuffer.killJobs(killJobs,'JEDI','51',True)
    except Exception as e:
        _logger.error("rebrokerage failed with {0} : {1}".format(str(e), traceback.format_exc()))

    # kill too long running jobs
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(days=21)
    status,res = taskBuffer.querySQLS("SELECT PandaID FROM ATLAS_PANDA.jobsActive4 WHERE creationTime<:creationTime",
                                  {':creationTime':timeLimit})
    jobs = []
    if res is not None:
        for (id,) in res:
            jobs.append(id)
    # kill
    if len(jobs):
        nJob = 100
        iJob = 0
        while iJob < len(jobs):
            # set tobekill
            _logger.debug('killJobs for Running (%s)' % jobs[iJob:iJob+nJob])
            Client.killJobs(jobs[iJob:iJob+nJob],2)
            # run watcher
            for id in jobs[iJob:iJob+nJob]:
                thr = Watcher(taskBuffer,id,single=True,sitemapper=siteMapper,sleepTime=60*24*21)
                thr.start()
                thr.join()
                time.sleep(1)
            iJob += nJob
            time.sleep(10)

    # kill too long waiting ddm jobs
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(days=5)
    varMap = {}
    varMap[':prodSourceLabel'] = 'ddm'
    varMap[':creationTime'] = timeLimit
    status,res = taskBuffer.querySQLS("SELECT PandaID FROM ATLAS_PANDA.jobsActive4 WHERE prodSourceLabel=:prodSourceLabel AND creationTime<:creationTime",
                                  varMap)
    jobs = []
    if res is not None:
        for (id,) in res:
            jobs.append(id)
    # kill
    if len(jobs):
        Client.killJobs(jobs,2)
        _logger.debug("killJobs for DDM (%s)" % str(jobs))

    # kill too long throttled jobs
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(days=7)
    varMap = {}
    varMap[':jobStatus'] = 'throttled'
    varMap[':creationTime'] = timeLimit
    status,res = taskBuffer.querySQLS("SELECT PandaID FROM ATLAS_PANDA.jobsActive4 WHERE jobStatus=:jobStatus AND creationTime<:creationTime ",
                                      varMap)
    jobs = []
    if res is not None:
        for (id,) in res:
            jobs.append(id)
    # kill
    if len(jobs):
        Client.killJobs(jobs,2)
        _logger.debug("killJobs for throttled (%s)" % str(jobs))

    # check if merge job is valid
    _logger.debug('kill invalid pmerge')
    varMap = {}
    varMap[':processingType'] = 'pmerge'
    varMap[':timeLimit'] = datetime.datetime.utcnow() - datetime.timedelta(minutes=30)
    sql  = "SELECT PandaID,jediTaskID FROM ATLAS_PANDA.jobsDefined4 WHERE processingType=:processingType AND modificationTime<:timeLimit "
    sql += "UNION "
    sql += "SELECT PandaID,jediTaskID FROM ATLAS_PANDA.jobsActive4 WHERE processingType=:processingType AND modificationTime<:timeLimit "
    status,res = taskBuffer.querySQLS(sql,varMap)
    nPmerge = 0
    badPmerge = 0
    _logger.debug('check {0} pmerge'.format(len(res)))
    for pandaID,jediTaskID in res:
        nPmerge += 1
        isValid,tmpMsg = taskBuffer.isValidMergeJob(pandaID,jediTaskID)
        if isValid is False:
            _logger.debug("kill pmerge {0} since {1} gone".format(pandaID,tmpMsg))
            taskBuffer.killJobs([pandaID],'killed since pre-merge job {0} gone'.format(tmpMsg),
                                '52',True)
            badPmerge += 1
    _logger.debug('killed invalid pmerge {0}/{1}'.format(badPmerge,nPmerge))

    # cleanup of jumbo jobs
    _logger.debug('jumbo job cleanup')
    res = taskBuffer.cleanupJumboJobs()
    _logger.debug(res)

    _memoryCheck("delete XML")

    # delete old files in DA cache
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(days=7)
    files = os.listdir(panda_config.cache_dir)
    for file in files:
        # skip special test file
        if file == 'sources.72c48dc5-f055-43e5-a86e-4ae9f8ea3497.tar.gz':
            continue
        if file == 'sources.090f3f51-fc81-4e80-9749-a5e4b2bd58de.tar.gz':
            continue
        try:
            # get timestamp
            timestamp = datetime.datetime.fromtimestamp(os.stat('%s/%s' % (panda_config.cache_dir,file)).st_mtime)
            # delete
            if timestamp < timeLimit:
                _logger.debug("delete %s " % file)
                os.remove('%s/%s' % (panda_config.cache_dir,file))
        except Exception:
            pass

    _memoryCheck("delete core")

    # delete core
    dirName = '%s/..' % panda_config.logdir
    for file in os.listdir(dirName):
        if file.startswith('core.'):
            _logger.debug("delete %s " % file)
            try:
                os.remove('%s/%s' % (dirName,file))
            except Exception:
                pass

    # update email DB
    _memoryCheck("email")
    _logger.debug("Update emails")

    # lock file
    _lockGetMail = open(panda_config.lockfile_getMail, 'w')
    # lock email DB
    fcntl.flock(_lockGetMail.fileno(), fcntl.LOCK_EX)
    # open email DB
    pDB = shelve.open(panda_config.emailDB)
    # read
    mailMap = {}
    for name in pDB:
        addr = pDB[name]
        mailMap[name] = addr
    # close DB
    pDB.close()
    # release file lock
    fcntl.flock(_lockGetMail.fileno(), fcntl.LOCK_UN)
    # set email address
    for name in mailMap:
        addr = mailMap[name]
        # remove _
        name = re.sub('_$','',name)
        status,res = taskBuffer.querySQLS("SELECT email FROM ATLAS_PANDAMETA.users WHERE name=:name",{':name':name})
        # failed or not found
        if status == -1 or len(res) == 0:
            _logger.error("%s not found in user DB" % name)
            continue
        # already set
        if res[0][0] not in ['','None',None]:
            continue
        # update email
        _logger.debug("set '%s' to %s" % (name,addr))
        status,res = taskBuffer.querySQLS("UPDATE ATLAS_PANDAMETA.users SET email=:addr WHERE name=:name",{':addr':addr,':name':name})

    # sandbox
    _logger.debug("Touch sandbox")
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(days=1)
    sqlC = "SELECT hostName,fileName,creationTime,userName FROM ATLAS_PANDAMETA.userCacheUsage "\
           "WHERE creationTime>:timeLimit AND creationTime>modificationTime "\
           "AND (fileName like 'sources%' OR fileName like 'jobO%') "
    sqlU = "UPDATE ATLAS_PANDAMETA.userCacheUsage SET modificationTime=CURRENT_DATE "\
           "WHERE userName=:userName AND fileName=:fileName "
    status, res = taskBuffer.querySQLS(sqlC, {':timeLimit': timeLimit})
    if res is None:
        _logger.error("failed to get files")
    elif len(res) > 0:
        _logger.debug("{0} files to touch".format(len(res)))
        for hostName, fileName, creationTime, userName in res:
            base_url = 'https://{0}:{1}'.format(hostName, panda_config.pserverport)
            _logger.debug("touch {0} on {1} created at {2}".format(fileName, hostName, creationTime))
            s,o = Client.touchFile(base_url, fileName)
            _logger.debug(o)
            if o == 'True':
                varMap = dict()
                varMap[':userName'] = userName
                varMap[':fileName'] = fileName
                taskBuffer.querySQLS(sqlU, varMap)

    _logger.debug("Check sandbox")
    timeLimit = datetime.datetime.utcnow() - datetime.timedelta(days=1)
    expireLimit = datetime.datetime.utcnow() - datetime.timedelta(days=30)
    sqlD = "DELETE FROM ATLAS_PANDAMETA.userCacheUsage WHERE userName=:userName AND fileName=:fileName "
    nRange = 100
    for i in range(nRange):
        _logger.debug("{0}/{1} {2} files to check".format(nRange, i, len(res)))
        res = taskBuffer.getLockSandboxFiles(timeLimit, 1000)
        if res is None:
            _logger.error("failed to get files")
            break
        elif len(res) == 0:
            break
        for userName, hostName, fileName, creationTime, modificationTime in res:
            url = 'https://{0}:{1}/cache/{2}'.format(hostName, panda_config.pserverport, fileName)
            _logger.debug("checking {0} created at {1}".format(url, creationTime))
            toDelete = False
            try:
                x = requests.head(url, verify=False)
                _logger.debug("code {0}".format(x.status_code))
                if x.status_code == 404:
                    _logger.debug("delete")
                    toDelete = True
            except Exception as e:
                _logger.debug("failed with {0}".format(str(e)))
                if creationTime < expireLimit:
                    toDelete = True
                    _logger.debug("delete due to creationTime={0}".format(creationTime))
            # update or delete
            varMap = dict()
            varMap[':userName'] = userName
            varMap[':fileName'] = fileName
            if toDelete:
                taskBuffer.querySQLS(sqlD, varMap)
            else:
                _logger.debug("keep")

    _memoryCheck("end")

    _logger.debug("===================== end =====================")

# run
if __name__ == '__main__':
    main(argv=sys.argv)


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

import unittest
import nose

import datetime
from pandas.msgpack import packb, unpackb

class DummyException(Exception):
    pass

class TestExceptions(unittest.TestCase):

    def test_raise_on_find_unsupported_value(self):
        import datetime
        self.assertRaises(TypeError, packb, datetime.datetime.now())

    def test_raise_from_object_hook(self):
        def hook(obj):
            raise DummyException
        self.assertRaises(DummyException, unpackb, packb({}), object_hook=hook)
        self.assertRaises(DummyException, unpackb, packb({'fizz': 'buzz'}), object_hook=hook)
        self.assertRaises(DummyException, unpackb, packb({'fizz': 'buzz'}), object_pairs_hook=hook)
        self.assertRaises(DummyException, unpackb, packb({'fizz': {'buzz': 'spam'}}), object_hook=hook)
        self.assertRaises(DummyException, unpackb, packb({'fizz': {'buzz': 'spam'}}), object_pairs_hook=hook)

    def test_invalidvalue(self):
        self.assertRaises(ValueError, unpackb, b'\xd9\x97#DL_')


#
# Author:  Travis Oliphant  2002-2011 with contributions from
#          SciPy Developers 2004-2011
#
from __future__ import division, print_function, absolute_import

from scipy import special
from scipy.special import entr, gammaln as gamln

from numpy import floor, ceil, log, exp, sqrt, log1p, expm1, tanh, cosh, sinh

import numpy as np

from ._distn_infrastructure import (
        rv_discrete, _lazywhere, _ncx2_pdf, _ncx2_cdf, get_distribution_names)

class binom_gen(rv_discrete):
    """A binomial discrete random variable.

    %(before_notes)s

    Notes
    -----
    The probability mass function for `binom` is::

       binom.pmf(k) = choose(n, k) * p**k * (1-p)**(n-k)

    for ``k`` in ``{0, 1,..., n}``.

    `binom` takes ``n`` and ``p`` as shape parameters.

    %(after_notes)s

    %(example)s

    """
    def _rvs(self, n, p):
        return self._random_state.binomial(n, p, self._size)

    def _argcheck(self, n, p):
        self.b = n
        return (n >= 0) & (p >= 0) & (p <= 1)

    def _logpmf(self, x, n, p):
        k = floor(x)
        combiln = (gamln(n+1) - (gamln(k+1) + gamln(n-k+1)))
        return combiln + special.xlogy(k, p) + special.xlog1py(n-k, -p)

    def _pmf(self, x, n, p):
        return exp(self._logpmf(x, n, p))

    def _cdf(self, x, n, p):
        k = floor(x)
        vals = special.bdtr(k, n, p)
        return vals

    def _sf(self, x, n, p):
        k = floor(x)
        return special.bdtrc(k, n, p)

    def _ppf(self, q, n, p):
        vals = ceil(special.bdtrik(q, n, p))
        vals1 = np.maximum(vals - 1, 0)
        temp = special.bdtr(vals1, n, p)
        return np.where(temp >= q, vals1, vals)

    def _stats(self, n, p, moments='mv'):
        q = 1.0 - p
        mu = n * p
        var = n * p * q
        g1, g2 = None, None
        if 's' in moments:
            g1 = (q - p) / sqrt(var)
        if 'k' in moments:
            g2 = (1.0 - 6*p*q) / var
        return mu, var, g1, g2

    def _entropy(self, n, p):
        k = np.r_[0:n + 1]
        vals = self._pmf(k, n, p)
        return np.sum(entr(vals), axis=0)
binom = binom_gen(name='binom')

class bernoulli_gen(binom_gen):
    """A Bernoulli discrete random variable.

    %(before_notes)s

    Notes
    -----
    The probability mass function for `bernoulli` is::

       bernoulli.pmf(k) = 1-p  if k = 0
                        = p    if k = 1

    for ``k`` in ``{0, 1}``.

    `bernoulli` takes ``p`` as shape parameter.

    %(after_notes)s

    %(example)s

    """
    def _rvs(self, p):
        return binom_gen._rvs(self, 1, p)

    def _argcheck(self, p):
        return (p >= 0) & (p <= 1)

    def _logpmf(self, x, p):
        return binom._logpmf(x, 1, p)

    def _pmf(self, x, p):
        return binom._pmf(x, 1, p)

    def _cdf(self, x, p):
        return binom._cdf(x, 1, p)

    def _sf(self, x, p):
        return binom._sf(x, 1, p)

    def _ppf(self, q, p):
        return binom._ppf(q, 1, p)

    def _stats(self, p):
        return binom._stats(1, p)

    def _entropy(self, p):
        return entr(p) + entr(1-p)
bernoulli = bernoulli_gen(b=1, name='bernoulli')

class nbinom_gen(rv_discrete):
    """A negative binomial discrete random variable.

    %(before_notes)s

    Notes
    -----
    The probability mass function for `nbinom` is::

         nbinom.pmf(k) = choose(k+n-1, n-1) * p**n * (1-p)**k

    for ``k >= 0``.

    `nbinom` takes ``n`` and ``p`` as shape parameters.

    %(after_notes)s

    %(example)s

    """
    def _rvs(self, n, p):
        return self._random_state.negative_binomial(n, p, self._size)

    def _argcheck(self, n, p):
        return (n > 0) & (p >= 0) & (p <= 1)

    def _pmf(self, x, n, p):
        return exp(self._logpmf(x, n, p))

    def _logpmf(self, x, n, p):
        coeff = gamln(n+x) - gamln(x+1) - gamln(n)
        return coeff + n*log(p) + special.xlog1py(x, -p)

    def _cdf(self, x, n, p):
        k = floor(x)
        return special.betainc(n, k+1, p)

    def _sf_skip(self, x, n, p):
        # skip because special.nbdtrc doesn't work for 0<n<1
        k = floor(x)
        return special.nbdtrc(k, n, p)

    def _ppf(self, q, n, p):
        vals = ceil(special.nbdtrik(q, n, p))
        vals1 = (vals-1).clip(0.0, np.inf)
        temp = self._cdf(vals1, n, p)
        return np.where(temp >= q, vals1, vals)

    def _stats(self, n, p):
        Q = 1.0 / p
        P = Q - 1.0
        mu = n*P
        var = n*P*Q
        g1 = (Q+P)/sqrt(n*P*Q)
        g2 = (1.0 + 6*P*Q) / (n*P*Q)
        return mu, var, g1, g2
nbinom = nbinom_gen(name='nbinom')

class geom_gen(rv_discrete):
    """A geometric discrete random variable.

    %(before_notes)s

    Notes
    -----
    The probability mass function for `geom` is::

        geom.pmf(k) = (1-p)**(k-1)*p

    for ``k >= 1``.

    `geom` takes ``p`` as shape parameter.

    %(after_notes)s

    %(example)s

    """
    def _rvs(self, p):
        return self._random_state.geometric(p, size=self._size)

    def _argcheck(self, p):
        return (p <= 1) & (p >= 0)

    def _pmf(self, k, p):
        return np.power(1-p, k-1) * p

    def _logpmf(self, k, p):
        return special.xlog1py(k - 1, -p) + log(p)

    def _cdf(self, x, p):
        k = floor(x)
        return -expm1(log1p(-p)*k)

    def _sf(self, x, p):
        return np.exp(self._logsf(x, p))

    def _logsf(self, x, p):
        k = floor(x)
        return k*log1p(-p)

    def _ppf(self, q, p):
        vals = ceil(log(1.0-q)/log(1-p))
        temp = self._cdf(vals-1, p)
        return np.where((temp >= q) & (vals > 0), vals-1, vals)

    def _stats(self, p):
        mu = 1.0/p
        qr = 1.0-p
        var = qr / p / p
        g1 = (2.0-p) / sqrt(qr)
        g2 = np.polyval([1, -6, 6], p)/(1.0-p)
        return mu, var, g1, g2
geom = geom_gen(a=1, name='geom', longname="A geometric")

class hypergeom_gen(rv_discrete):
    """A hypergeometric discrete random variable.

    The hypergeometric distribution models drawing objects from a bin.
    M is the total number of objects, n is total number of Type I objects.
    The random variate represents the number of Type I objects in N drawn
    without replacement from the total population.

    %(before_notes)s

    Notes
    -----
    The probability mass function is defined as::

        pmf(k, M, n, N) = choose(n, k) * choose(M - n, N - k) / choose(M, N),
                                       for max(0, N - (M-n)) <= k <= min(n, N)

    %(after_notes)s

    Examples
    --------
    >>> from scipy.stats import hypergeom
    >>> import matplotlib.pyplot as plt

    Suppose we have a collection of 20 animals, of which 7 are dogs.  Then if
    we want to know the probability of finding a given number of dogs if we
    choose at random 12 of the 20 animals, we can initialize a frozen
    distribution and plot the probability mass function:

    >>> [M, n, N] = [20, 7, 12]
    >>> rv = hypergeom(M, n, N)
    >>> x = np.arange(0, n+1)
    >>> pmf_dogs = rv.pmf(x)

    >>> fig = plt.figure()
    >>> ax = fig.add_subplot(111)
    >>> ax.plot(x, pmf_dogs, 'bo')
    >>> ax.vlines(x, 0, pmf_dogs, lw=2)
    >>> ax.set_xlabel('# of dogs in our group of chosen animals')
    >>> ax.set_ylabel('hypergeom PMF')
    >>> plt.show()

    Instead of using a frozen distribution we can also use `hypergeom`
    methods directly.  To for example obtain the cumulative distribution
    function, use:

    >>> prb = hypergeom.cdf(x, M, n, N)

    And to generate random numbers:

    >>> R = hypergeom.rvs(M, n, N, size=10)

    """
    def _rvs(self, M, n, N):
        return self._random_state.hypergeometric(n, M-n, N, size=self._size)

    def _argcheck(self, M, n, N):
        cond = (M > 0) & (n >= 0) & (N >= 0)
        cond &= (n <= M) & (N <= M)
        self.a = max(N-(M-n), 0)
        self.b = min(n, N)
        return cond

    def _logpmf(self, k, M, n, N):
        tot, good = M, n
        bad = tot - good
        return gamln(good+1) - gamln(good-k+1) - gamln(k+1) + gamln(bad+1) \
            - gamln(bad-N+k+1) - gamln(N-k+1) - gamln(tot+1) + gamln(tot-N+1) \
            + gamln(N+1)

    def _pmf(self, k, M, n, N):
        # same as the following but numerically more precise
        # return comb(good, k) * comb(bad, N-k) / comb(tot, N)
        return exp(self._logpmf(k, M, n, N))

    def _stats(self, M, n, N):
        # tot, good, sample_size = M, n, N
        # "wikipedia".replace('N', 'M').replace('n', 'N').replace('K', 'n')
        M, n, N = 1.*M, 1.*n, 1.*N
        m = M - n
        p = n/M
        mu = N*p

        var = m*n*N*(M - N)*1.0/(M*M*(M-1))
        g1 = (m - n)*(M-2*N) / (M-2.0) * sqrt((M-1.0) / (m*n*N*(M-N)))

        g2 = M*(M+1) - 6.*N*(M-N) - 6.*n*m
        g2 *= (M-1)*M*M
        g2 += 6.*n*N*(M-N)*m*(5.*M-6)
        g2 /= n * N * (M-N) * m * (M-2.) * (M-3.)
        return mu, var, g1, g2

    def _entropy(self, M, n, N):
        k = np.r_[N - (M - n):min(n, N) + 1]
        vals = self.pmf(k, M, n, N)
        return np.sum(entr(vals), axis=0)

    def _sf(self, k, M, n, N):
        """More precise calculation, 1 - cdf doesn't cut it."""
        # This for loop is needed because `k` can be an array. If that's the
        # case, the sf() method makes M, n and N arrays of the same shape. We
        # therefore unpack all inputs args, so we can do the manual
        # integration.
        res = []
        for quant, tot, good, draw in zip(k, M, n, N):
            # Manual integration over probability mass function. More accurate
            # than integrate.quad.
            k2 = np.arange(quant + 1, draw + 1)
            res.append(np.sum(self._pmf(k2, tot, good, draw)))
        return np.asarray(res)
hypergeom = hypergeom_gen(name='hypergeom')

# FIXME: Fails _cdfvec
class logser_gen(rv_discrete):
    """A Logarithmic (Log-Series, Series) discrete random variable.

    %(before_notes)s

    Notes
    -----
    The probability mass function for `logser` is::

        logser.pmf(k) = - p**k / (k*log(1-p))

    for ``k >= 1``.

    `logser` takes ``p`` as shape parameter.

    %(after_notes)s

    %(example)s

    """
    def _rvs(self, p):
        # looks wrong for p>0.5, too few k=1
        # trying to use generic is worse, no k=1 at all
        return self._random_state.logseries(p, size=self._size)

    def _argcheck(self, p):
        return (p > 0) & (p < 1)

    def _pmf(self, k, p):
        return -np.power(p, k) * 1.0 / k / log(1 - p)

    def _stats(self, p):
        r = log(1 - p)
        mu = p / (p - 1.0) / r
        mu2p = -p / r / (p - 1.0)**2
        var = mu2p - mu*mu
        mu3p = -p / r * (1.0+p) / (1.0 - p)**3
        mu3 = mu3p - 3*mu*mu2p + 2*mu**3
        g1 = mu3 / np.power(var, 1.5)

        mu4p = -p / r * (
            1.0 / (p-1)**2 - 6*p / (p - 1)**3 + 6*p*p / (p-1)**4)
        mu4 = mu4p - 4*mu3p*mu + 6*mu2p*mu*mu - 3*mu**4
        g2 = mu4 / var**2 - 3.0
        return mu, var, g1, g2
logser = logser_gen(a=1, name='logser', longname='A logarithmic')

class poisson_gen(rv_discrete):
    """A Poisson discrete random variable.

    %(before_notes)s

    Notes
    -----
    The probability mass function for `poisson` is::

        poisson.pmf(k) = exp(-mu) * mu**k / k!

    for ``k >= 0``.

    `poisson` takes ``mu`` as shape parameter.

    %(after_notes)s

    %(example)s

    """
    def _rvs(self, mu):
        return self._random_state.poisson(mu, self._size)

    def _logpmf(self, k, mu):
        Pk = k*log(mu)-gamln(k+1) - mu
        return Pk

    def _pmf(self, k, mu):
        return exp(self._logpmf(k, mu))

    def _cdf(self, x, mu):
        k = floor(x)
        return special.pdtr(k, mu)

    def _sf(self, x, mu):
        k = floor(x)
        return special.pdtrc(k, mu)

    def _ppf(self, q, mu):
        vals = ceil(special.pdtrik(q, mu))
        vals1 = np.maximum(vals - 1, 0)
        temp = special.pdtr(vals1, mu)
        return np.where(temp >= q, vals1, vals)

    def _stats(self, mu):
        var = mu
        tmp = np.asarray(mu)
        g1 = sqrt(1.0 / tmp)
        g2 = 1.0 / tmp
        return mu, var, g1, g2
poisson = poisson_gen(name="poisson", longname='A Poisson')

class planck_gen(rv_discrete):
    """A Planck discrete exponential random variable.

    %(before_notes)s

    Notes
    -----
    The probability mass function for `planck` is::

        planck.pmf(k) = (1-exp(-lambda_))*exp(-lambda_*k)

    for ``k*lambda_ >= 0``.

    `planck` takes ``lambda_`` as shape parameter.

    %(after_notes)s

    %(example)s

    """
    def _argcheck(self, lambda_):
        if (lambda_ > 0):
            self.a = 0
            self.b = np.inf
            return 1
        elif (lambda_ < 0):
            self.a = -np.inf
            self.b = 0
            return 1
        else:
            return 0

    def _pmf(self, k, lambda_):
        fact = (1-exp(-lambda_))
        return fact*exp(-lambda_*k)

    def _cdf(self, x, lambda_):
        k = floor(x)
        return 1-exp(-lambda_*(k+1))

    def _ppf(self, q, lambda_):
        vals = ceil(-1.0/lambda_ * log1p(-q)-1)
        vals1 = (vals-1).clip(self.a, np.inf)
        temp = self._cdf(vals1, lambda_)
        return np.where(temp >= q, vals1, vals)

    def _stats(self, lambda_):
        mu = 1/(exp(lambda_)-1)
        var = exp(-lambda_)/(expm1(-lambda_))**2
        g1 = 2*cosh(lambda_/2.0)
        g2 = 4+2*cosh(lambda_)
        return mu, var, g1, g2

    def _entropy(self, lambda_):
        l = lambda_
        C = (1-exp(-l))
        return l*exp(-l)/C - log(C)
planck = planck_gen(name='planck', longname='A discrete exponential ')

class boltzmann_gen(rv_discrete):
    """A Boltzmann (Truncated Discrete Exponential) random variable.

    %(before_notes)s

    Notes
    -----
    The probability mass function for `boltzmann` is::

        boltzmann.pmf(k) = (1-exp(-lambda_)*exp(-lambda_*k)/(1-exp(-lambda_*N))

    for ``k = 0,..., N-1``.

    `boltzmann` takes ``lambda_`` and ``N`` as shape parameters.

    %(after_notes)s

    %(example)s

    """
    def _pmf(self, k, lambda_, N):
        fact = (1-exp(-lambda_))/(1-exp(-lambda_*N))
        return fact*exp(-lambda_*k)

    def _cdf(self, x, lambda_, N):
        k = floor(x)
        return (1-exp(-lambda_*(k+1)))/(1-exp(-lambda_*N))

    def _ppf(self, q, lambda_, N):
        qnew = q*(1-exp(-lambda_*N))
        vals = ceil(-1.0/lambda_ * log(1-qnew)-1)
        vals1 = (vals-1).clip(0.0, np.inf)
        temp = self._cdf(vals1, lambda_, N)
        return np.where(temp >= q, vals1, vals)

    def _stats(self, lambda_, N):
        z = exp(-lambda_)
        zN = exp(-lambda_*N)
        mu = z/(1.0-z)-N*zN/(1-zN)
        var = z/(1.0-z)**2 - N*N*zN/(1-zN)**2
        trm = (1-zN)/(1-z)
        trm2 = (z*trm**2 - N*N*zN)
        g1 = z*(1+z)*trm**3 - N**3*zN*(1+zN)
        g1 = g1 / trm2**(1.5)
        g2 = z*(1+4*z+z*z)*trm**4 - N**4 * zN*(1+4*zN+zN*zN)
        g2 = g2 / trm2 / trm2
        return mu, var, g1, g2
boltzmann = boltzmann_gen(name='boltzmann',
        longname='A truncated discrete exponential ')

class randint_gen(rv_discrete):
    """A uniform discrete random variable.

    %(before_notes)s

    Notes
    -----
    The probability mass function for `randint` is::

        randint.pmf(k) = 1./(high - low)

    for ``k = low, ..., high - 1``.

    `randint` takes ``low`` and ``high`` as shape parameters.

    Note the difference to the numpy ``random_integers`` which
    returns integers on a *closed* interval ``[low, high]``.

    %(after_notes)s

    %(example)s

    """
    def _argcheck(self, low, high):
        self.a = low
        self.b = high - 1
        return (high > low)

    def _pmf(self, k, low, high):
        p = np.ones_like(k) / (high - low)
        return np.where((k >= low) & (k < high), p, 0.)

    def _cdf(self, x, low, high):
        k = floor(x)
        return (k - low + 1.) / (high - low)

    def _ppf(self, q, low, high):
        vals = ceil(q * (high - low) + low) - 1
        vals1 = (vals - 1).clip(low, high)
        temp = self._cdf(vals1, low, high)
        return np.where(temp >= q, vals1, vals)

    def _stats(self, low, high):
        m2, m1 = np.asarray(high), np.asarray(low)
        mu = (m2 + m1 - 1.0) / 2
        d = m2 - m1
        var = (d*d - 1) / 12.0
        g1 = 0.0
        g2 = -6.0/5.0 * (d*d + 1.0) / (d*d - 1.0)
        return mu, var, g1, g2

    def _rvs(self, low, high=None):
        """An array of *size* random integers >= ``low`` and < ``high``.

        If ``high`` is ``None``, then range is >=0  and < low
        """
        return self._random_state.randint(low, high, self._size)

    def _entropy(self, low, high):
        return log(high - low)
randint = randint_gen(name='randint', longname='A discrete uniform '
                      '(random integer)')

# FIXME: problems sampling.
class zipf_gen(rv_discrete):
    """A Zipf discrete random variable.

    %(before_notes)s

    Notes
    -----
    The probability mass function for `zipf` is::

        zipf.pmf(k, a) = 1/(zeta(a) * k**a)

    for ``k >= 1``.

    `zipf` takes ``a`` as shape parameter.

    %(after_notes)s

    %(example)s

    """
    def _rvs(self, a):
        return self._random_state.zipf(a, size=self._size)

    def _argcheck(self, a):
        return a > 1

    def _pmf(self, k, a):
        Pk = 1.0 / special.zeta(a, 1) / k**a
        return Pk

    def _munp(self, n, a):
        return _lazywhere(
            a > n + 1, (a, n),
            lambda a, n: special.zeta(a - n, 1) / special.zeta(a, 1),
            np.inf)
zipf = zipf_gen(a=1, name='zipf', longname='A Zipf')

class dlaplace_gen(rv_discrete):
    """A  Laplacian discrete random variable.

    %(before_notes)s

    Notes
    -----
    The probability mass function for `dlaplace` is::

        dlaplace.pmf(k) = tanh(a/2) * exp(-a*abs(k))

    for ``a > 0``.

    `dlaplace` takes ``a`` as shape parameter.

    %(after_notes)s

    %(example)s

    """
    def _pmf(self, k, a):
        return tanh(a/2.0) * exp(-a * abs(k))

    def _cdf(self, x, a):
        k = floor(x)
        f = lambda k, a: 1.0 - exp(-a * k) / (exp(a) + 1)
        f2 = lambda k, a: exp(a * (k+1)) / (exp(a) + 1)
        return _lazywhere(k >= 0, (k, a), f=f, f2=f2)

    def _ppf(self, q, a):
        const = 1 + exp(a)
        vals = ceil(np.where(q < 1.0 / (1 + exp(-a)), log(q*const) / a - 1,
                                                      -log((1-q) * const) / a))
        vals1 = vals - 1
        return np.where(self._cdf(vals1, a) >= q, vals1, vals)

    def _stats(self, a):
        ea = exp(a)
        mu2 = 2.*ea/(ea-1.)**2
        mu4 = 2.*ea*(ea**2+10.*ea+1.) / (ea-1.)**4
        return 0., mu2, 0., mu4/mu2**2 - 3.

    def _entropy(self, a):
        return a / sinh(a) - log(tanh(a/2.0))
dlaplace = dlaplace_gen(a=-np.inf,
                        name='dlaplace', longname='A discrete Laplacian')

class skellam_gen(rv_discrete):
    """A  Skellam discrete random variable.

    %(before_notes)s

    Notes
    -----
    Probability distribution of the difference of two correlated or
    uncorrelated Poisson random variables.

    Let k1 and k2 be two Poisson-distributed r.v. with expected values
    lam1 and lam2. Then, ``k1 - k2`` follows a Skellam distribution with
    parameters ``mu1 = lam1 - rho*sqrt(lam1*lam2)`` and
    ``mu2 = lam2 - rho*sqrt(lam1*lam2)``, where rho is the correlation
    coefficient between k1 and k2. If the two Poisson-distributed r.v.
    are independent then ``rho = 0``.

    Parameters mu1 and mu2 must be strictly positive.

    For details see: http://en.wikipedia.org/wiki/Skellam_distribution

    `skellam` takes ``mu1`` and ``mu2`` as shape parameters.

    %(after_notes)s

    %(example)s

    """
    def _rvs(self, mu1, mu2):
        n = self._size
        return (self._random_state.poisson(mu1, n) -
                self._random_state.poisson(mu2, n))

    def _pmf(self, x, mu1, mu2):
        px = np.where(x < 0,
                _ncx2_pdf(2*mu2, 2*(1-x), 2*mu1)*2,
                _ncx2_pdf(2*mu1, 2*(1+x), 2*mu2)*2)
        # ncx2.pdf() returns nan's for extremely low probabilities
        return px

    def _cdf(self, x, mu1, mu2):
        x = floor(x)
        px = np.where(x < 0,
                _ncx2_cdf(2*mu2, -2*x, 2*mu1),
                1-_ncx2_cdf(2*mu1, 2*(x+1), 2*mu2))
        return px

    def _stats(self, mu1, mu2):
        mean = mu1 - mu2
        var = mu1 + mu2
        g1 = mean / sqrt((var)**3)
        g2 = 1 / var
        return mean, var, g1, g2
skellam = skellam_gen(a=-np.inf, name="skellam", longname='A Skellam')

# Collect names of classes and objects in this module.
pairs = list(globals().items())
_distn_names, _distn_gen_names = get_distribution_names(pairs, rv_discrete)

__all__ = _distn_names + _distn_gen_names