import sys
import traceback
import zipfile
from array import *
import re as regEx
import random
import socket
import struct
import ipaddress
import subprocess
import os
from datetime import datetime
import time
import logging

# Regex strings for all it should search for in the files
ipv4Pattern = regEx.compile(r'(25[0-5]|2[0-4][0-9]|1[0-9]{2}|[1-9][0-9]|[0-9])\.(25[0-5]|2[0-4][0-9]|1[0-9]{2}|[1-9]'
                            r'[0-9]|[1-9]|0)\.(25[0-5]|2[0-4][0-9]|1[0-9]{2}|[1-9][0-9]|[1-9]|0)\.(25[0-5]|2[0-4][0-9]|'
                            r'1[0-9]{2}|[1-9][0-9]|[0-9])')

ipv6Pattern = regEx.compile(r'(([0-9a-fA-F]{1,4}:){7,7}[0-9a-fA-F]{1,4}|([0-9a-fA-F]{1,4}:){1,7}:|([0-9a-fA-F]{1,4}:)'
                            r'{1,6}:[0-9a-fA-F]{1,4}|([0-9a-fA-F]{1,4}:){1,5}(:[0-9a-fA-F]{1,4}){1,2}|([0-9a-fA-F]'
                            r'{1,4}:){1,4}(:[0-9a-fA-F]{1,4}){1,3}|([0-9a-fA-F]{1,4}:){1,3}(:[0-9a-fA-F]{1,4}){1,4}|'
                            r'([0-9a-fA-F]{1,4}:){1,2}(:[0-9a-fA-F]{1,4}){1,5}|[0-9a-fA-F]{1,4}:((:[0-9a-fA-F]{1,4})'
                            r'{1,6})|:((:[0-9a-fA-F]{1,4}){1,7}|:)|fe80:(:[0-9a-fA-F]{0,4}){0,4}%[0-9a-zA-Z]{1,}|::'
                            r'(ffff(:0{1,4}){0,1}:){0,1}((25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9])\.){3,3}(25[0-5]|'
                            r'(2[0-4]|1{0,1}[0-9]){0,1}[0-9])|([0-9a-fA-F]{1,4}:){1,4}:((25[0-5]|(2[0-4]|1{0,1}[0-9])'
                            r'{0,1}[0-9])\.){3,3}(25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9]))')

imsiPattern = regEx.compile(r'[0-9]{13,17}')

imsiHexPattern = regEx.compile(r'(IMSI :   { |IMSI : )(\d{2}\ |\d[A-Z]\ |[A-Z]\d\ )(\d{2}\ |\d[A-Z]\ |[A-Z]\d\ )'
                               r'(\d{2}\ |\d[A-Z]\ |[A-Z]\d\ )(\d{2}\ |\d[A-Z]\ |[A-Z]\d\ )(\d{2}\ |\d[A-Z]\ |'
                               r'[A-Z]\d\ )(\d{2}\ |\d[A-Z]\ |[A-Z]\d\ )(\d{2}\ |\d[A-Z]\ |[A-Z]\d\ )(\d{2}\ |'
                               r'\d[A-Z]\ |[A-Z]\d\ )(\d{2}|\d[A-Z]|[A-Z]\d)')

macPattern = regEx.compile(r'(?:[0-9A-Fa-f]{2}[:-]){5}(?:[0-9A-Fa-f]{2})')

hostnamePattern = regEx.compile(r'(?i)epcgw')

usernamePattern = regEx.compile(r'(?i)serviceuser')

urlPattern = regEx.compile(r'(?i)splunk')

hostnameMatches = {}

usernameMatches = {}

urlMatches = {}

# Dictionaries and lists for all matches found during the washing
ipv4AddressMatches = {}
ipv6AddressMatches = {}
imsiMatches = {}
imsiHexMatches = {}
macMatches = {}
urlFound = []
hostnameFound = []
usernameFound = []

# Different variables needed throughout the script
inboxLocation = '/nfs/data/inbox/'
folderLocation = '/nfs/data/'
tmpLocation = '/nfs/data/tmp/'
folder_dt = datetime.now()
dt_string = folder_dt.strftime('%y-%m-%d-%H.%M')
outboxLocation = folderLocation + '/outbox/'
outboxDirname = 'washed-' + dt_string
outboxDir = folderLocation + '/outbox/' + outboxDirname
tmpDir = tmpLocation + outboxDirname
script_log = '/nfs/data/cron-script.log'
washing_log = '/local/scramble/washing-script/log/washing-script.log'
zipEnd = '.zip'
zipLogEnd = '.log.zip'
gzEnd = '.gz'
gzLogEnd = '.log.gz'
tarEnd = '.tar'
tarLogEnd = '.log.tar'
targzEnd = '.tar.gz'
targzLogEnd = '.log.tar.gz'

def replaceCharsInTuple(tuple):
    try:
        # Clean up IP-addresses in tuple and return as string
        tupleToReturn = tuple.replace(',', '.').replace('(', '').replace(')', '').replace(' ', '').replace('\'', '')
        return ''.join(tupleToReturn)
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def replaceIpv6Tuple(tuple):
    try:
        ipv6TupleToReturn = tuple[0]
        return ipv6TupleToReturn
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def InsertSpaceInTupleImsiHex(tuple):
    try:
        # Clean up IMSI hex in tuple and return as string
        tupleToReturn = tuple.replace(',', '').replace('(', '').replace(')', '').replace('\'', '') \
            .replace('IMSI : ', '').replace('IMSI :   { ', '').replace(':', '').replace('{', '').replace('   ', '') \
            .replace('  ', ' ')
        return ''.join(tupleToReturn)
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def checkIfIpv4ExistsAndReplace(match):
    try:
        # If match not found in dictionary, generate a new IPv4 address with the first two octates as x
        if ipv4AddressMatches.get(match) == None:
            ipv4List = ['x', 'x']
            ipv4List.append(str(random.randint(1, 255)))
            ipv4List.append(str(random.randint(1, 255)))
            ipv4AddressMatches[match] = '.'.join(ipv4List)
        return ipv4AddressMatches.get(match)
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def checkIfIpv6ExistsAndReplace(match):
    try:
        # If match not found in dictionary, generate a new IPv6 address
        false_ipv6 = regEx.compile(r'[a-fA-F]{1,3}::|::[a-fA-F]{1,3}| ::|:: | :: ')
        false_match = regEx.findall(false_ipv6, match)
        socket_match = socket.inet_pton(socket.AF_INET6, match)
        if True:
            if match != '::':
                if len(false_match) == 0:
                    if ipv6AddressMatches.get(match) == None:
                        ipv6AddressMatches[match] = ipaddress.IPv6Address(
                            random.randint(0, 2 ** 128 - 1))  # Add random IPv6
                    return ipv6AddressMatches.get(match)
                if len(false_match) <= 1:
                    ipv6AddressMatches[match] = false_match[0]
                    return ipv6AddressMatches.get(match)
            if match == '::':
                ipv6AddressMatches[match] = '::'
                return ipv6AddressMatches.get(match)
        if False:
            ipv6AddressMatches[match] = false_match[0]
            return ipv6AddressMatches.get(match)

    except socket.error:
        ipv6AddressMatches[match] = false_match[0]
        return ipv6AddressMatches.get(match)

def checkIfMacExistsAndReplace(match):
    try:
        # If match not found in dictionary, generate a new mac address
        if macMatches.get(match) == None:
            macMatches[match] = '%02x:%02x:%02x:%02x:%02x:%02x' % (
                random.randint(0, 255),
                random.randint(0, 255),
                random.randint(0, 255),
                random.randint(0, 255),
                random.randint(0, 255),
                random.randint(0, 255))
        return macMatches.get(match)
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def checkIfImsiExistsAndReplace(match):
    try:
        # If match not found in dictionary, generate a new random number with xxxxxx in front
        if imsiMatches.get(match) == None:
            imsiList = ['xxxxx']
            imsiList.append(str(random.randint(1000000000, 9999999999)))
            imsiMatches[match] = ''.join(imsiList)
        return imsiMatches.get(match)
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def checkIfImsiHexExistsAndReplace(match):
    try:
        # If match not found in dictionary, generate a new random number with xx on part of the hex
        if imsiHexMatches.get(match) == None:
            imsiHexList = ['xx ', 'xx ', 'xx ', 'xx ']
            imsiHexList.append(str(random.randint(10, 99)) + ' ' + str(random.randint(10, 99)) + ' ' +
                               str(random.randint(10, 99)) + ' ' + str(random.randint(10, 99)) + ' ' +
                               str(random.randint(0, 9)) + 'F')
            imsiHexMatches[match] = str(''.join(imsiHexList))
        return imsiHexMatches.get(match)
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def checkIfHostnameExistsAndReplace(match):
    try:
        # If match not found in dictionary, replace it. If found replace item listed in the dictionary
        if hostnameMatches.get(match) == None:
            logger.warning(current_time() + ' - !!!!!!!!!!' + match + ' not found in hostname list!!!!!!!!!!\n'
                                                                       'Match removed from file')
            hostnameMatches[match] = str('xxxxxxx')
        return hostnameMatches.get(match)
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def checkIfUsernameExistsAndReplace(match):
    try:
        # If match not found in dictionary, replace it. If found replace item listed in the dictionary
        if usernameMatches.get(match) == None:
            usernameMatches[match] = str('xxxxxxxxx')
        return usernameMatches.get(match)
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def checkIfUrlExistsAndReplace(match):
    try:
        # If match not found in dictionary, replace it. If found replace item listed in the dictionary
        if urlMatches.get(match) == None:
            logger.warning(current_time() + ' - !!!!!!!!!!' + match + ' not found in URL list!!!!!!!!!!\n'
                                                                       'Match removed from file')
            urlMatches[match] = str('xxxxx')
        return urlMatches.get(match)
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def current_time():
    return time.strftime('%d-%m-%y %H:%M:%S', time.localtime())

logging.basicConfig(filename=washing_log,
                    format='%(message)s')
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)

def wash_filename(item):
    try:
        global washed_filename
        current_filename = str(item)
        filename_ipv4Match = regEx.findall(ipv4Pattern, item)
        filename_ipv6Match = regEx.findall(ipv6Pattern, item)
        filename_imsiMatch = regEx.findall(imsiPattern, item)
        filename_imsiHexMatch = regEx.findall(imsiHexPattern, item)
        filename_hostnameMatch = regEx.findall(hostnamePattern, item)
        filename_macMatch = regEx.findall(macPattern, item)
        filename_usernameMatch = regEx.findall(usernamePattern, item)
        filename_urlMatch = regEx.findall(urlPattern, item)
        filename_ipv4Array = []
        filename_ipv6Array = []
        filename_imsiArray = []
        filename_imsiHexArray = []
        filename_macArray = []
        filename_hostnameArray = []
        filename_usernameArray = []
        filename_urlArray = []
        for i in range(len(filename_ipv4Match)):
            filename_ipv4Array.append(replaceCharsInTuple(str(filename_ipv4Match[i])))
        for i in range(len(filename_ipv6Match)):
            filename_ipv6Array.append(replaceIpv6Tuple(filename_ipv6Match[i]))
        for i in range(len(filename_macMatch)):
            filename_macArray.append(filename_macMatch[i])
        for i in range(len(filename_imsiMatch)):
            filename_imsiArray.append(filename_imsiMatch[i])
        for i in range(len(filename_imsiHexMatch)):
            filename_imsiHexArray.append(InsertSpaceInTupleImsiHex(str(filename_imsiHexMatch[i])))
        for i in range(len(filename_hostnameMatch)):
            filename_hostnameArray.append(filename_hostnameMatch[i])
        for i in range(len(filename_usernameMatch)):
            filename_usernameArray.append(filename_usernameMatch[i])
        for i in range(len(filename_urlMatch)):
            filename_urlArray.append(filename_urlMatch[i])
        for ipv4 in filename_ipv4Array:
            replacedIpv4Address = checkIfIpv4ExistsAndReplace(ipv4)
            new_filename = item.replace(ipv4, replacedIpv4Address)
            washed_filename = new_filename  # Overwrite current filename with new filename
        for ipv6 in filename_ipv6Array:
            replacedIpv6Address = checkIfIpv6ExistsAndReplace(ipv6)
            new_filename = item.replace(ipv6, str(replacedIpv6Address))
            washed_filename = new_filename  # Overwrite current filename with new filename
        for mac in filename_macArray:
            replacedMacAddress = checkIfMacExistsAndReplace(mac.lower())
            new_filename = item.replace(mac, replacedMacAddress)
            washed_filename = new_filename  # Overwrite current filename with new filename
        for imsi in filename_imsiArray:
            replacedImsiAddress = checkIfImsiExistsAndReplace(imsi)
            new_filename = item.replace(imsi, replacedImsiAddress)
            washed_filename = new_filename  # Overwrite current filename with new filename
        for imsiHex in filename_imsiHexArray:
            replacedImsiHexAddress = checkIfImsiHexExistsAndReplace(imsiHex)
            new_filename = item.replace(imsiHex, replacedImsiHexAddress)
            washed_filename = new_filename  # Overwrite current filename with new filename
        for hostname in filename_hostnameArray:
            if hostname.lower() not in hostnameFound:
                hostnameFound.append(hostname.lower())
            replacedHostnameAddress = checkIfHostnameExistsAndReplace(hostname.lower())
            new_filename = item.replace(hostname, replacedHostnameAddress)
            washed_filename = new_filename  # Overwrite current filename with new filename
        for username in filename_usernameArray:
            if username.lower() not in usernameFound:
                usernameFound.append(username.lower())
            replacedUsernameAddress = checkIfUsernameExistsAndReplace(username.lower())
            new_filename = item.replace(username, str(replacedUsernameAddress))
            washed_filename = new_filename  # Overwrite current filename with new filename
        for url in filename_urlArray:
            if url.lower() not in urlFound:
                urlFound.append(url.lower())
            replacedurlAddress = checkIfUrlExistsAndReplace(url.lower())
            new_filename = item.replace(url, replacedurlAddress)
            washed_filename = new_filename  # Overwrite current filename with new filename
        if current_filename == str(washed_filename):
            return washed_filename
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def unzipFiles():
    try:
        os.chdir(tmpLocation)
        os.getcwd()
        # Walk down the tmp folder and find/unzip all zipped files
        # If there is a zipped file, it will start over to walk into the unzipped files aswell
        for root, dirs, files in os.walk(tmpLocation):
            for item in files:
                itemLocation = root + '/' + item
                # Unzip files ending with .gz, .zip or .tar, then delete the zipped version
                if os.path.exists(itemLocation):
                    if item.endswith(zipEnd):
                        logger.info(current_time() + ' - Unzipping: "' + root + '"/"' + item + '"')
                        newZipName = item.replace('.zip', '')
                        os.system('mkdir "' + str(root) + '"/"' + str(newZipName) + '"')
                        os.system('unzip -o -qq "' + str(root) + '"/"' + str(item) + '" -d "' + str(root) + '"/"' +
                                  str(newZipName) + '"')
                        os.system('rm -rf "' + str(root) + '"/"' + str(item) + '"')
                        unzipFiles()
                    if item.endswith(zipLogEnd):
                        logger.info(current_time() + ' - Unzipping: "' + root + '"/"' + item + '"')
                        newZipName = item.replace('.zip', '')
                        os.system('mkdir "' + str(root) + '"/"' + str(newZipName) + '"')
                        os.system('unzip -o -qq "' + str(root) + '"/"' + str(item) + '" -d "' + str(root) + '"/"' +
                                  str(newZipName) + '"')
                        os.system('rm -rf "' + str(root) + '"/"' + str(item) + '"')
                    if item.endswith(targzEnd) or item.endswith(tarEnd):
                        logger.info(current_time() + ' - Unzipping tar file: "' + root + '"/"' + str(item) + '"')
                        newTarName1 = item.replace('.gz', '')
                        newTarName2 = newTarName1.replace('.tar', '')
                        os.system('mkdir "' + str(root) + '"/"' + newTarName2 + '"')
                        os.system('tar -xf "' + str(root) + '"/"' + str(item) + '" -C "' + str(root) + '"/"' +
                                  str(newTarName2) + '"')
                        os.system('rm -rf "' + str(root) + '"/"' + str(item) + '"')
                        unzipFiles()
                    if item.endswith(targzLogEnd) or item.endswith(tarLogEnd):
                        logger.info(current_time() + ' - Unzipping tar file: "' + root + '"/"' + str(item) + '"')
                        newTarName1 = item.replace('.gz', '')
                        newTarName2 = newTarName1.replace('.tar', '')
                        os.system('mkdir "' + str(root) + '"/"' + newTarName2 + '"')
                        os.system('tar -xf "' + str(root) + '"/"' + str(item) + '" -C "' + str(root) + '"/"' +
                                  str(newTarName2) + '"')
                        os.system('rm -rf "' + str(root) + '"/"' + str(item) + '"')
                    if item.endswith(gzEnd):
                        logger.info(current_time() + ' - Unzipping gz file: "' + root + '"/"' + item + '"')
                        os.system('gzip -fd "' + str(root) + '"/"' + str(item) + '"')
                        unzipFiles()
                    if item.endswith(gzLogEnd):
                        logger.info(current_time() + ' - Unzipping gz file: "' + root + '"/"' + item + '"')
                        os.system('gzip -fd "' + str(root) + '"/"' + str(item) + '"')
        walking_file()
    except Exception as err:
        with open(script_log, 'a') as f:
            f.write(current_time() + ' - An error occurred\n' + traceback.format(err))
        logger.error(current_time() + ' - An error occurred\n' + traceback.format(err))
        quit()

def walking_file():
    try:
        global washed_filename
        tmp_dirs = os.listdir(tmpLocation)
        for folder in tmp_dirs:
            for root, dirs, files in os.walk(folder):
                for i in dirs:
                    washed_filename = ''
                    wash_filename(str(i))
                    if str(i) != washed_filename and washed_filename != '':
                        os.system('mv "' + root + '"/"' + str(i) + '" "' + root + '"/"' + washed_filename + '"')
                        walking_file()
            for root, dirs, files in os.walk(folder):
                for file in files:
                    washed_filename = ''
                    wash_filename(str(file))
                    if str(file) != washed_filename and washed_filename != '':
                        os.system('mv "' + root + '"/"' + file + '" "' + root + '"/"' + washed_filename + '"')
        washFiles()
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def washFiles():
    try:
        with open(script_log, 'a') as f:
            f.write(current_time() + ' - Done unzipping all folders\n')
        logger.info(current_time() + ' - Done unzipping all folders')
        tmp_dirs = os.listdir(tmpLocation)
        for directory in tmp_dirs:
            # Walking down the file structure in the inbox file and washing each file
            for root, dirs, files in os.walk(directory):
                for file in files:
                    logger.info(current_time() + ' - Now washing ' + file)
                    newFileContent = ''
                    with open(root + '/' + file, 'r+') as logFile:
                        for line in logFile:
                            # Going through each line in the file and finding regex mathes
                            currentLine = line.strip()
                            ipv4MatchesInLine = regEx.findall(ipv4Pattern, line)
                            ipv6MatchesInLine = regEx.findall(ipv6Pattern, line)
                            imsiMatchesInLine = regEx.findall(imsiPattern, line)
                            imsiHexMatchesInLine = regEx.findall(imsiHexPattern, line)
                            macMatchesInLine = regEx.findall(macPattern, line)
                            hostnameMatchesInLine = regEx.findall(hostnamePattern, line)
                            usernameMatchesInLine = regEx.findall(usernamePattern, line)
                            urlMatchesInLine = regEx.findall(urlPattern, line)
                            ipv4Array = []
                            ipv6Array = []
                            imsiArray = []
                            imsiHexArray = []
                            macArray = []
                            hostnameArray = []
                            usernameArray = []
                            urlArray = []
                            # Add all matches to their respective array
                            for i in range(len(ipv4MatchesInLine)):
                                ipv4Array.append(replaceCharsInTuple(str(ipv4MatchesInLine[i])))
                            for i in range(len(ipv6MatchesInLine)):
                                ipv6Array.append(replaceIpv6Tuple(ipv6MatchesInLine[i]))
                            for i in range(len(macMatchesInLine)):
                                macArray.append(macMatchesInLine[i])
                            for i in range(len(imsiMatchesInLine)):
                                imsiArray.append(imsiMatchesInLine[i])
                            for i in range(len(imsiHexMatchesInLine)):
                                imsiHexArray.append(InsertSpaceInTupleImsiHex(str(imsiHexMatchesInLine[i])))
                            for i in range(len(hostnameMatchesInLine)):
                                hostnameArray.append(hostnameMatchesInLine[i])
                            for i in range(len(usernameMatchesInLine)):
                                usernameArray.append(usernameMatchesInLine[i])
                            for i in range(len(urlMatchesInLine)):
                                urlArray.append(urlMatchesInLine[i])
                            # Replace all matches found with the other entry in the dictionary
                            for ipv4 in ipv4Array:
                                replacedIpv4Address = checkIfIpv4ExistsAndReplace(ipv4)
                                newLine = currentLine.replace(ipv4, replacedIpv4Address)
                                currentLine = newLine  # Overwrite current line with newLine to reflect changes made
                            for ipv6 in ipv6Array:
                                replacedIpv6Address = checkIfIpv6ExistsAndReplace(ipv6)
                                newLine = currentLine.replace(ipv6, str(replacedIpv6Address))
                                currentLine = newLine  # Overwrite current line with newLine to reflect changes made
                            for mac in macArray:
                                replacedMacAddress = checkIfMacExistsAndReplace(mac.lower())
                                newLine = currentLine.replace(mac, replacedMacAddress)
                                currentLine = newLine  # Overwrite current line with newLine to reflect changes made
                            for imsi in imsiArray:
                                replacedImsiAddress = checkIfImsiExistsAndReplace(imsi)
                                newLine = currentLine.replace(imsi, replacedImsiAddress)
                                currentLine = newLine  # Overwrite current line with newLine to reflect changes made
                            for imsiHex in imsiHexArray:
                                replacedImsiHexAddress = checkIfImsiHexExistsAndReplace(imsiHex)
                                newLine = currentLine.replace(str(imsiHex), str(replacedImsiHexAddress))
                                currentLine = newLine  # Overwrite current line with newLine to reflect changes made
                            for hostname in hostnameArray:
                                if hostname.lower() not in hostnameFound:
                                    hostnameFound.append(hostname.lower())
                                replacedHostnameAddress = checkIfHostnameExistsAndReplace(hostname.lower())
                                newLine = currentLine.replace(hostname, replacedHostnameAddress)
                                currentLine = newLine  # Overwrite current line with newLine to reflect changes made
                            for username in usernameArray:
                                if username.lower() not in usernameFound:
                                    usernameFound.append(username.lower())
                                replacedUsernameAddress = checkIfUsernameExistsAndReplace(username.lower())
                                newLine = currentLine.replace(username, str(replacedUsernameAddress))
                                currentLine = newLine  # Overwrite current line with newLine to reflect changes made
                            for url in urlArray:
                                if url.lower() not in urlFound:
                                    urlFound.append(url.lower())
                                replacedurlAddress = checkIfUrlExistsAndReplace(url.lower())
                                newLine = currentLine.replace(url, replacedurlAddress)
                                currentLine = newLine  # Overwrite current line with newLine to reflect changes made
                            newFileContent += currentLine + '\n'
                        logFile.truncate(0)  # Remove old content of file
                        logFile.seek(0)  # Start writing from index 0
                        logFile.write(newFileContent)
                    logger.info(current_time() + ' - ###Done washing ' + file + '###')
            os.system('touch ' + directory + '/washingreport.txt')
            # Making a file which contains all matches found and what they are changed to
            with open(directory + '/washingreport.txt', 'r+') as reportfile:
                stringToWrite = ('########################\nResult from washing\n########################\n\n' +
                                 'Ipv4 dictionary:\n' + str(ipv4AddressMatches) +
                                 '\n\nIpv6 dictionary:\n' + str(ipv6AddressMatches) +
                                 '\n\nImsi dictionary:\n' + str(imsiMatches) +
                                 '\n\nImsi hex dictionary:\n' + str(imsiHexMatches) +
                                 '\n\nMac address dictionary:\n' + str(macMatches) +
                                 '\n\nUsernames found:\n' + str(usernameFound) +
                                 '\n\nHostnames found:\n' + str(hostnameFound) +
                                 '\n\nUrl found:\n' + str(urlFound) + '\n')
                reportfile.write(stringToWrite)
            # Move all files washed in tmp folder to the respective outbox folder
            os.system('chmod 777 ' + directory)
            os.system('mv ' + directory + ' ' + outboxLocation)
            logger.info('\n\n' + current_time() + '\n########################\nMoved files to: ' + outboxLocation + directory +
                  '\n########################\n')
        with open(script_log, 'a') as f:
            f.write(str(current_time()) + ' - Washing is complete\n')
        logger.info(current_time() + ' - Washing is complete\n')
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def move_files():
    try:
        content_inbox = os.listdir(inboxLocation)
        content_tmp = os.listdir(tmpLocation)
        logger.info(current_time() + ' - Script starting')
        if len(content_inbox) == 0 and len(content_tmp) == 0:
            quit()
        if len(content_inbox) != 0 and len(content_tmp) == 0:
            while True:
                sftp_files = subprocess.check_output('du -H -d1 /nfs/data/inbox/', shell=True).decode()
                time.sleep(2)
                sftp_files_2 = subprocess.check_output('du -H -d1 /nfs/data/inbox/', shell=True).decode()
                if str(sftp_files) == str(sftp_files_2):
                    break
                if str(sftp_files) != str(sftp_files_2):
                    with open(script_log, 'a') as f:
                        f.write(str(current_time()) + ' - Files are still being transferred, will wait to start script\n')
                    logger.warning(current_time() + ' - Files are still being transferred, will wait to start script')
                    time.sleep(2)
            # Make a new directory in /tmp with date and time
            os.system('mkdir ' + tmpDir)
            os.system('cp -r ' + inboxLocation + '* ' + tmpDir)
            os.system('rm -rf ' + inboxLocation + '*')
            os.system('chmod -R 777 ' + tmpDir)
            tmp_folder_content = os.listdir(tmpDir)
            for folder in tmp_folder_content:
                if folder.endswith(zipEnd):
                    logger.info(current_time() + ' - Unzipping: ' + tmpDir + '/"' + folder + '"')
                    newZipName = folder.replace('.zip', '')
                    os.system('mkdir ' + str(tmpDir) + '/"' + str(newZipName) + '"')
                    os.system('unzip -o -qq ' + str(tmpDir) + '/"' + str(folder) + '" -d ' + str(tmpDir) + '/"' +
                              newZipName + '"')
                    os.system('rm -rf ' + str(tmpDir) + '/"' + str(folder) + '"')
                if folder.endswith(gzEnd):
                    logger.info(current_time() + ' - Unzipping gz file: ' + tmpDir + '/"' + folder + '"')
                    os.system('gzip -fd ' + str(tmpDir) + '/"' + str(folder) + '"')
                if folder.endswith(targzEnd) or folder.endswith(tarEnd):
                    logger.info(current_time() + ' - Unzipping tar file: ' + tmpDir + '/"' + folder + '"')
                    newTarName1 = folder.replace('.gz', '')
                    newTarName2 = newTarName1.replace('.tar', '')
                    os.system('mkdir ' + str(tmpDir) + '/"' + newTarName2 + '"')
                    os.system('tar -xf ' + str(tmpDir) + '/"' + str(folder) + '" -C ' + str(tmpDir) + '/"' +
                              newTarName2 + '"')
                    os.system('rm -rf ' + str(tmpDir) + '/"' + folder + '"')
            with open(script_log, 'a') as f:
                f.write(str(current_time()) + ' - Done unzipping root folder\n')
            logger.info(current_time() + ' - Done unzipping root folder')
            check_folder()
        if len(content_tmp) != 0 and len(content_inbox) == 0:
            check_folder()
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

def check_folder():
    try:
        num_gz = subprocess.check_output('find ' + tmpLocation + ' -name "*gz"', shell=True).decode()
        num_gz = num_gz.split('\n')
        num_zipped = subprocess.check_output('find ' + tmpLocation + ' -name "*zip"', shell=True).decode()
        num_zipped = num_zipped.split('\n')
        num_tar = subprocess.check_output('find ' + tmpLocation + ' -name "*tar"', shell=True).decode()
        num_tar = num_tar.split('\n')
        total_zipped = len(num_gz) + len(num_zipped) + len(num_tar) - 3
        if total_zipped >= 500:
            new_limit = total_zipped * 3
            sys.setrecursionlimit(new_limit)
            limit = str(new_limit)
            with open(script_log, 'a') as f:
                f.write(str(current_time()) + ' - Amount of zipped files will exceed max amount of calls, will change '
                                              'recursion limit to: ' + limit + '\n')
            logger.warning(current_time() + ' - Amount of zipped files will exceed max amount of calls, will change '
                                             'recursion limit to: ' + limit)
            unzipFiles()
        if total_zipped <= 500:
            unzipFiles()
    except Exception:
        logger.error(current_time() + ' - An error occurred:\n' + traceback.format_exc())

if __name__ == '__main__':
    move_files()


__all__ = [
    "JobItem",
]
import datetime

import top

class JobItem(top.Table):
    """job_item table ORM.
    """
    _job = top.Job()
    _agent_stocktake = top.AgentStocktake()

    def __init__(self):
        """Toll Outlet Portal job_item table initialiser.
        """
        super(JobItem, self).__init__('job_item')

    @property
    def schema(self):
        return ["id INTEGER PRIMARY KEY",
                "job_id INTEGER",
                "connote_nbr CHAR(30)",
                "item_nbr CHAR(32)",
                "consumer_name CHAR(30)",
                "email_addr CHAR(60)",
                "phone_nbr CHAR(20)",
                "pieces INTEGER",
                "status INTEGER",
                "created_ts TIMESTAMP",
                "pickup_ts TIMESTAMP",
                "pod_name CHAR(40)",
                "identity_type_id INTEGER",
                "identity_type_data CHAR(30)",
                "extract_ts TIMESTAMP",
                "reminder_ts TIMESTAMP",
                "notify_ts TIMESTAMP"]

    def collected_sql(self, business_unit, ignore_pe=False):
        """SQL wrapper to extract the collected items from the "jobitems"
        table.

        **Args:**
            *business_unit*: the id relating to the job.bu_id value.

        **Kwargs:**
            *ignore_pe*: ``boolean`` flag to ignore job items whose
            parent job is Primary Elect (default ``False``)

        **Returns:**
            the SQL string

        """
        sql = """SELECT ji.connote_nbr as 'REF1',
    ji.id as 'JOB_KEY',
    ji.pickup_ts as 'PICKUP_TIME',
    ji.pod_name as 'PICKUP_POD',
    it.description as 'IDENTITY_TYPE',
    ji.identity_type_data as 'IDENTITY_DATA',
    ji.item_nbr as 'ITEM_NBR',
    ag.code as 'AGENT_ID',
    ag.state as 'AGENT_STATE'
FROM job_item as ji, identity_type as it, job as j, agent as ag
WHERE pickup_ts IS NOT null
AND extract_ts IS null
AND ji.identity_type_id = it.id
AND j.agent_id = ag.id
AND (ji.job_id = j.id AND j.bu_id = %d)""" % business_unit

        if ignore_pe:
            sql += """
AND (j.service_code != 3 OR j.service_code IS NULL)"""

        return sql

    def upd_collected_sql(self, id, time):
        """SQL wrapper to update the collected items from the "jobitems"
        table.

        **Args:**
            *id*: the id relating to the ``jobitem.id`` value.

        **Returns:**
            the SQL string

        """
        sql = """UPDATE job_item
SET extract_ts = '%s'
WHERE id = %d""" % (time, id)

        return sql

    def upd_file_based_collected_sql(self, connote, item_nbr, time=None):
        """SQL wrapper to update the collected items from the "jobitems"
        table.

        This variant of the :meth:`upd_collected_sql` method is used
        to close of file-based extractions.

        **Args:**
            *connote*: connote value relating to the ``jobitem.connote_nbr``
            value

            *item_nbr*: connote value relating to the ``jobitem.item_nbr``
            value

        **Kwargs:**
            *time*: override the time to set from the current time

        **Returns:**
            the SQL string

        """
        if time is None:
            time = datetime.datetime.now().isoformat(' ').split('.', 1)[0]

        sql = """UPDATE %(name)s
SET extract_ts = '%(time)s', pickup_ts = '%(time)s'
WHERE connote_nbr = '%(connote)s'
AND item_nbr = '%(item_nbr)s'""" % {'name': self.name,
                                    'time': time,
                                    'connote': connote,
                                    'item_nbr': item_nbr}

        return sql

    def connote_sql(self, connote):
        """SQL wrapper to extract records where job_item.connote_nbr
        is equal to *connote*.

        **Args:**
            connote: Connote value relating to the job_item.connote_nbr
            column.

        **Returns:**
            the SQL string

        """
        sql = """SELECT id
FROM %s
WHERE connote_nbr = '%s'""" % (self.name, connote)

        return sql

    def connote_item_nbr_sql(self, connote, item_nbr):
        """SQL wrapper to extract records where job_item.connote_nbr
        is equal to *connote* and job_item.item_nbr equals *item_nbr*.

        **Args:**
            connote: Connote value relating to the job_item.connote_nbr
            column.

            item_nbr: Item Number value relating to the job_item.item_nbr
            column.

        **Returns:**
            the SQL string

        """
        sql = """SELECT id
FROM %s
WHERE connote_nbr = '%s'
AND item_nbr = '%s'
ORDER BY created_ts DESC""" % (self.name, connote, item_nbr)

        return sql

    def item_number_sql(self, item_nbr):
        """SQL wrapper to extract records where job_item.item_nbr
        is equal to *item_nbr*.

        **Args:**
            item_nbr: Item Number value relating to the job_item.item_nbr
            column.

        **Returns:**
            the SQL string

        """
        sql = """SELECT id
FROM %s
WHERE item_nbr = '%s'""" % (self.name, item_nbr)

        return sql

    def uncollected_sql(self, start_date, uncollected_period):
        """SQL wrapper to extract the job_item records which remain
        uncollected after *uncollected_period* has elapsed.

        **Args:**
            uncollected_period: job_item.notify_ts value that defines
            an uncollected parcel

        **Returns:**
            the SQL string

        """
        sql = """SELECT id
FROM job_item
WHERE (created_ts > '%s' AND notify_ts < '%s')
AND pickup_ts IS NULL
AND (email_addr != '' OR phone_nbr != '')
AND reminder_ts IS NULL""" % (start_date, uncollected_period)

        return sql

    def job_item_agent_details_sql(self, job_item_id):
        """SQL wrapper to extract the agent details against a *job_item_id*.

        SQL also returns additional information relating to the
        *job_item_id* such as:

        * ``jobitem.connote_nbr``
        * ``jobitem.item_nbr``
        * ``jobitem.notify_ts``
        * ``jobitem.created_ts``
        * ``jobitem.email_addr``
        * ``jobitem.phone_nbr``
        * ``jobitem.pickup_ts``

        and the *job.bu_id*.

        **Args:**
            job_item_id: the jobitem.id value to search against

        **Returns:**
            the SQL string

        """
        sql = """SELECT ag.name,
       ag.address,
       ag.suburb,
       ag.postcode,
       ji.connote_nbr,
       ji.item_nbr,
       ji.notify_ts,
       ji.created_ts,
       ji.email_addr,
       ji.phone_nbr,
       ji.pickup_ts,
       j.bu_id
FROM job_item as ji, job as j, agent as ag
WHERE ji.job_id = j.id
AND j.agent_id = ag.id
AND ji.id = %d""" % job_item_id

        return sql

    def update_reminder_ts_sql(self, id, ts=None):
        return self.update_timestamp_sql(id, column='reminder_ts', ts=ts)

    def update_notify_ts_sql(self, id, ts=None):
        return self.update_timestamp_sql(id, column='notify_ts', ts=ts)

    def update_timestamp_sql(self, id, column, ts=None):
        """SQL wrapper to update the ``job_item.reminder_ts`` to *ts*
        timestamp.

        **Args:**
            *id*: integer value relating to the ``job_item.id``

            *column*: the timestamp column to update

        **Kwargs:**
            *ts*: override the current time

        **Returns:**
            the SQL string

        """
        if ts is None:
            ts = datetime.datetime.now().isoformat(' ').split('.', 1)[0]

        sql = """UPDATE %s
SET %s = '%s'
WHERE id = %d
""" % (self.name, column, ts, id)

        return sql

    def connote_base_primary_elect_job(self, connote):
        """SQL wrapper to verify if a *connote* is associated with a primary
        elect job.

        Primary elect jobs are identified by a integer value 3 in the
        ``job.service_code`` column.

        **Args:**
            connote: the jobitem.connote value to search against.

        **Returns:**
            the SQL string

        """
        sql = """SELECT ji.id
FROM job as j, %s as ji
WHERE ji.job_id = j.id
AND ji.connote_nbr = '%s'
AND ji.notify_ts IS NULL
AND (ji.email_addr != '' OR ji.phone_nbr != '')
AND j.service_code = 3""" % (self.name, connote)

        return sql

    def uncollected_jobitems_sql(self,
                                 service_code=3,
                                 bu_ids=None,
                                 delivery_partners=None,
                                 day_range=14):
        """SQL wrapper to extract uncollected Service Code-based jobs.

        Service Code jobs are identified by a integer value in the
        ``job.service_code`` column.

        Query will ignore records that have either white space or
        a spurious ``.`` in the email or phone number columns.

        The *bu_ids* relate to the ``job.bu_id`` column.

        **Kwargs:**
            *service_code*: value relating to the ``job.service_code``
            column (default ``3`` for Primary Elect)

            *bu_ids*: integer based tuple of Business Unit ID's to search
            against (default ``None`` ignores all Business Units)

            *day_range*: number of days from current time to include
            in search (default 14.0 days)

        **Returns:**
            the SQL string

        """
        if bu_ids is None:
            bu_ids = tuple()
        if len(bu_ids) == 1:
            bu_ids = '(%d)' % bu_ids[0]

        if delivery_partners is None:
            delivery_partners = tuple()
        if len(delivery_partners) == 1:
            delivery_partners = "('%s')" % delivery_partners[0]

        now = datetime.datetime.now()
        start_ts = now - datetime.timedelta(days=day_range)
        start_date = start_ts.strftime('%Y-%m-%d %H:%M:%S')

        sql = """SELECT ji.id, ji.connote_nbr, ji.item_nbr
FROM job AS j, %(name)s AS ji, agent AS ag, delivery_partner AS dp
WHERE ji.job_id = j.id
AND j.agent_id = ag.id
AND ag.dp_id = dp.id
AND dp.name IN %(dps)s
AND ji.pickup_ts IS NULL
AND ji.notify_ts IS NULL
AND (ji.email_addr NOT IN ('', '.') OR ji.phone_nbr NOT IN ('', '.'))
AND j.bu_id IN %(bu_ids)s
AND j.service_code = %(sc)d
AND ji.created_ts > '%(start_date)s'""" % {'name': self.name,
                                           'dps': str(delivery_partners),
                                           'bu_ids': str(bu_ids),
                                           'sc': service_code,
                                           'start_date': start_date}

        return sql

    def reference_sql(self,
                      bu_ids,
                      reference_nbr=None,
                      picked_up=False,
                      delivery_partners=None,
                      columns=None,
                      alias='ji'):
        """Extract connote_nbr/item_nbr against *reference_nbr*.

        Query is an ``OR`` against both ``connote_nbr`` and ``item_nbr``.

        **Args:**
            *bu_ids*: integer based tuple of Business Unit ID's to search
            against (default ``None`` ignores all Business Units)

        **Kwargs:**
            *reference_nbr*: parcel ID number as scanned by the agent.  If
            ``None``, then the values from the ``agent_stocktake`` table
            will be used.

            *picked_up*: boolean flag that will extract ``job_items``
            that have been picked up if ``True``. Otherwise, will extract
            ``job_items`` that have not been picked up if ``False``.

            *delivery_partners*: string based list of Delivery Partner
            names to limit result set against.  For example,
            ``['Nparcel', 'Toll']``.  The values supported are as per
            the ``delivery_partner.name`` table set

            *columns*: string prepresentation of the columns to query
            against

            *alias*: table alias (default ``ji``)

        **Returns:**
            the SQL string

        """
        if columns is None:
            columns = self._select_columns(alias)

        if not picked_up:
            pickup_sql = 'IS NULL'
        else:
            pickup_sql = 'IS NOT NULL'

        if len(bu_ids) == 1:
            bu_ids = '(%d)' % bu_ids[0]

        ref = reference_nbr
        if reference_nbr is None:
            ref = self._agent_stocktake.reference_sql()

        dp_sql = str()
        if delivery_partners is not None:
            dps = ', '.join(["'%s'" % x for x in delivery_partners])
            dps = '(%s)' % dps
            if len(delivery_partners) == 1:
                dps = "('%s')" % delivery_partners[0]

            dp_sql = """AND dp.name IN %s AND ag.dp_id = dp.id""" % dps

        union_sql = self.job_based_reference_sql(bu_ids,
                                                 ref,
                                                 picked_up,
                                                 delivery_partners,
                                                 columns)

        sql = """SELECT DISTINCT %(columns)s
FROM %(name)s as %(alias)s,
     job AS j,
     agent AS ag,
     agent_stocktake AS st,
     delivery_partner AS dp
WHERE %(alias)s.job_id = j.id
AND ag.id = st.agent_id
AND j.bu_id IN %(bu_ids)s
AND j.agent_id = ag.id
AND (%(alias)s.connote_nbr IN (%(ref)s)
     OR %(alias)s.item_nbr IN (%(ref)s))
AND %(alias)s.pickup_ts %(pickup_sql)s
%(dp_sql)s
UNION
%(union)s""" % {'columns': columns,
                'bu_ids': str(bu_ids),
                'name': self.name,
                'ref': ref,
                'alias': alias,
                'union': union_sql,
                'pickup_sql': pickup_sql,
                'dp_sql': dp_sql}

        return sql

    def job_based_reference_sql(self,
                                bu_ids,
                                reference_nbr,
                                picked_up=False,
                                delivery_partners=None,
                                columns=None,
                                alias='ji'):
        """Extract connote_nbr/item_nbr against *reference_nbr* matched
        to the ``job.card_ref_nbr``.

        Query is an ``OR`` against both ``connote_nbr`` and ``item_nbr``.

        **Args:**
            *bu_ids*: integer based tuple of Business Unit ID's to search
            against (default ``None`` ignores all Business Units)

            *reference_nbr*: parcel ID number as scanned by the agent

        **Kwargs:**
            *picked_up*: boolean flag that will extract ``job_items``
            that have been picked up if ``True``. Otherwise, will extract
            ``job_items`` that have not been picked up if ``False``.

            *delivery_partners*: string based list of Delivery Partner
            names to limit result set against.  For example,
            ``['Nparcel', 'Toll']``.  The values supported are as per
            the ``delivery_partner.name`` table set

            *columns*: string prepresentation of the columns to query
            against

            *alias*: table alias

        **Returns:**
            the SQL string

        """
        if columns is None:
            columns = self._select_columns(alias)

        pickup_sql = 'AND %s.pickup_ts ' % alias
        if not picked_up:
            pickup_sql += 'IS NULL'
        else:
            pickup_sql += 'IS NOT NULL'

        if len(bu_ids) == 1:
            bu_ids = '(%d)' % bu_ids[0]

        dp_sql = str()
        if delivery_partners is not None:
            dps = ', '.join(["'%s'" % x for x in delivery_partners])
            dps = '(%s)' % dps
            if len(delivery_partners):
                dps = "('%s')" % delivery_partners[0]

            dp_sql = """AND dp.name IN %s AND ag.dp_id = dp.id""" % dps

        sql = """SELECT DISTINCT %(columns)s
FROM %(name)s AS %(alias)s,
     job AS j,
     agent AS ag,
     delivery_partner AS dp,
     agent_stocktake AS st
WHERE %(alias)s.job_id = j.id
AND j.bu_id IN %(bu_ids)s
AND j.agent_id = ag.id
%(dp_sql)s
AND %(alias)s.job_id IN
(
%(sql)s
)
%(pickup_sql)s""" % {'columns': columns,
                     'bu_ids': bu_ids,
                     'name': self.name,
                     'sql': self._job.reference_sql(reference_nbr),
                     'alias': alias,
                     'pickup_sql': pickup_sql,
                     'dp_sql': dp_sql}

        return sql

    def _select_columns(self, alias='ji'):
        """Helper method that captures required columns in the
        uncollected aged report query.

        **Kwargs:**
            *alias*: table alias

        **Returns:**
            the SQL string

        """
        columns = """%(alias)s.id as JOB_ITEM_ID,
       j.bu_id as JOB_BU_ID,
       %(alias)s.connote_nbr as CONNOTE_NBR,
       j.card_ref_nbr as BARCODE,
       %(alias)s.item_nbr as ITEM_NBR,
       j.job_ts as JOB_TS,
       %(alias)s.created_ts as CREATED_TS,
       %(alias)s.notify_ts as NOTIFY_TS,
       %(alias)s.pickup_ts as PICKUP_TS,
       %(alias)s.pieces as PIECES,
       %(alias)s.consumer_name as CONSUMER_NAME,
       ag.dp_code as DP_CODE,
       ag.code as AGENT_CODE,
       ag.name as AGENT_NAME,
       ag.address as AGENT_ADDRESS,
       ag.suburb as AGENT_SUBURB,
       ag.state as AGENT_STATE,
       ag.postcode as AGENT_POSTCODE,
       ag.phone_nbr as AGENT_PHONE_NBR,
       (SELECT DISTINCT ag.dp_code
        FROM agent_stocktake AS st, agent AS aag
        WHERE (%(alias)s.connote_nbr = st.reference_nbr
               OR j.card_ref_nbr = st.reference_nbr
               OR %(alias)s.item_nbr = st.reference_nbr)
        AND st.agent_id = aag.id) AS ST_DP_CODE,
       (SELECT DISTINCT ag.code
        FROM agent_stocktake AS st, agent AS aag
        WHERE (%(alias)s.connote_nbr = st.reference_nbr
               OR j.card_ref_nbr = st.reference_nbr
               OR %(alias)s.item_nbr = st.reference_nbr)
        AND st.agent_id = aag.id) AS ST_AGENT_CODE,
       (SELECT DISTINCT ag.name
        FROM agent_stocktake AS st, agent AS aag
        WHERE (%(alias)s.connote_nbr = st.reference_nbr
               OR j.card_ref_nbr = st.reference_nbr
               OR %(alias)s.item_nbr = st.reference_nbr)
        AND st.agent_id = aag.id) AS ST_AGENT_NAME""" % {'alias': alias}

        return columns

    def non_compliance_sql(self,
                           bu_ids,
                           picked_up=False,
                           delivery_partners=None,
                           alias='ji'):
        """Extract ``job_item`` detail of all items in the ``job_item``
        table that do not exist in the ``agent_stocktake`` table.

        Senarios are based on the *picked_up* flag.  For example, all
        parcels that *have* been picked up or *have not* been picked up.

        **Args:**
            *bu_ids*: integer based tuple of Business Unit ID's to search
            against (default ``None`` ignores all Business Units)

        **Kwargs:**
            *picked_up*: boolean flag that will extract ``job_items``
            that have been picked up if ``True``. Otherwise, will extract
            ``job_items`` that have not been picked up if ``False``.

            *delivery_partners*: string based list of Delivery Partner
            names to limit result set against.  For example,
            ``['Nparcel', 'Toll']``.  The values supported are as per
            the ``delivery_partner.name`` table set

            *alias*: table alias (default ``ji``)

        **Returns:**
            the SQL string

        """
        if bu_ids is None:
            bu_ids = tuple()

        if len(bu_ids) == 1:
            bu_ids = '(%d)' % bu_ids[0]

        columns = self._select_columns(alias)
        col = 'ji.id'

        if not picked_up:
            pickup_sql = 'IS NULL'
        else:
            pickup_sql = 'IS NOT NULL'

        dp_sql = str()
        if delivery_partners is not None:
            dps = ', '.join(["'%s'" % x for x in delivery_partners])
            dps = '(%s)' % dps
            if len(delivery_partners) == 1:
                dps = "('%s')" % delivery_partners[0]

            dp_sql = """AND dp.name IN %s AND ag.dp_id = dp.id""" % dps

        sql = """SELECT DISTINCT %(columns)s
FROM %(name)s AS %(alias)s, job AS j, agent AS ag, delivery_partner AS dp
WHERE %(alias)s.job_id = j.id
AND j.agent_id = ag.id
AND %(alias)s.pickup_ts %(pickup_sql)s
%(dp_sql)s
AND ji.id NOT IN
(%(sql)s)""" % {'columns': columns,
                'alias': alias,
                'name': self.name,
                'pickup_sql': pickup_sql,
                'dp_sql': dp_sql,
                'sql': self.reference_sql(bu_ids=bu_ids,
                                          picked_up=picked_up,
                                          delivery_partners=delivery_partners,
                                          columns=col)}

        return sql

    def total_agent_stocktake_parcel_count_sql(self,
                                               bu_ids,
                                               picked_up=False,
                                               delivery_partners=None,
                                               day_range=7,
                                               alias='ji'):
        """Sum ``agent_stocktake`` based parcel counts per ADP based on
        *picked_up*.

        Query is an ``OR`` against both ``connote_nbr`` and ``item_nbr``.

        **Args:**
            *bu_ids*: integer based tuple of Business Unit ID's to search
            against (default ``None`` ignores all Business Units)

        **Kwargs:**
            *picked_up*: boolean flag that will extract ``job_items``
            that have been picked up if ``True``. Otherwise, will extract
            ``job_items`` that have not been picked up if ``False``.

            *delivery_partners*: string based tuple of Delivery Partner
            names to limit result set against.  For example,
            ``['top', 'toll']``.  The values supported are as per
            the ``delivery_partner.name`` table set

            *day_range*: number of days from current time to include
            in the agent_stocktake table search (default 7 days)

            *alias*: table alias (default ``ji``)

        **Returns:**
            the SQL string

        """
        now = datetime.datetime.now()
        start_ts = now - datetime.timedelta(days=day_range)
        start_date = start_ts.strftime('%Y-%m-%d %H:%M:%S')

        if not picked_up:
            pickup_sql = 'IS NULL'
        else:
            pickup_sql = 'IS NOT NULL'

        if len(bu_ids) == 1:
            bu_ids = '(%d)' % bu_ids[0]

        col = 'st.reference_nbr'
        sql_ref = self.reference_sql(bu_ids=bu_ids,
                                     picked_up=picked_up,
                                     delivery_partners=delivery_partners,
                                     columns=col)

        sql = """SELECT DISTINCT agent.dp_code AS DP_CODE,
        agent.code AS AGENT_CODE,
        agent.name AS AGENT_NAME,
        (SELECT MAX(st.created_ts)
         FROM agent_stocktake AS st, agent AS ag
         WHERE ag.code = agent.code
         AND ag.id = st.agent_id) AS STOCKTAKE_CREATED_TS,
        (SELECT COUNT(DISTINCT ags.reference_nbr)
         FROM agent_stocktake AS ags, agent AS ag
         WHERE ags.agent_id = agent.id
         AND ags.created_ts > '%(start_date)s') AS AGENT_PIECES,
        (SELECT COUNT(ji.id)
         FROM job_item AS ji, job AS j, agent AS ag
         WHERE ag.id = agent.id
         AND j.agent_id = ag.id
         AND j.id = ji.job_id
         AND ji.pickup_ts %(pickup_sql)s) AS TPP_PIECES
FROM agent AS agent, agent_stocktake AS agent_stocktake
WHERE agent_stocktake.agent_id = agent.id
AND reference_nbr != ''
AND agent_stocktake.created_ts > '%(start_date)s'
AND agent_stocktake.reference_nbr IN (%(sql_ref)s)
GROUP BY agent.id,
         agent.dp_code,
         agent.code,
         agent.name,
         agent_stocktake.created_ts,
         agent_stocktake.reference_nbr""" % {'sql_ref': sql_ref,
                                             'start_date': start_date,
                                             'pickup_sql': pickup_sql}

        return sql

    def total_parcel_count_sql(self,
                               picked_up=False,
                               delivery_partners=None,
                               alias='ji'):
        """Sum parcel counts per ADP based on *picked_up*.

        **Kwargs:**
            *picked_up*: boolean flag that will extract ``job_items``
            that have been picked up if ``True``. Otherwise, will extract
            ``job_items`` that have not been picked up if ``False``.

            *delivery_partners*: string based tuple of Delivery Partner
            names to limit result set against.  For example,
            ``['Nparcel', 'Toll']``.  The values supported are as per
            the ``delivery_partner.name`` table set

            *alias*: table alias (default ``ji``)

        **Returns:**
            the SQL string

        """
        if not picked_up:
            pickup_sql = 'IS NULL'
        else:
            pickup_sql = 'IS NOT NULL'

        dp_sql = str()
        if delivery_partners is not None:
            dps = ', '.join(["'%s'" % x for x in delivery_partners])
            dps = '(%s)' % dps
            if len(delivery_partners) == 1:
                dps = "('%s')" % delivery_partners[0]

            dp_sql = """AND dp.name IN %s AND ag.dp_id = dp.id""" % dps

        sql = """SELECT SUM(%(alias)s.pieces)
FROM %(name)s AS %(alias)s, job AS j, agent AS ag, delivery_partner AS dp
WHERE %(alias)s.job_id = j.id
AND j.agent_id = ag.id
%(dp_sql)s
AND %(alias)s.pickup_ts %(pickup_sql)s""" % {'name': self.name,
                                             'dp_sql': dp_sql,
                                             'pickup_sql': pickup_sql,
                                             'alias': alias}

        return sql

    def agent_id_of_aged_parcels(self,
                                 period=7,
                                 delivery_partners=None,
                                 alias='ji'):
        """SQL to provide a distinct list of agents that have an
        aged parcel.

        **Kwargs:**
            *period*: time (in days) from now that is the cut off for
            agent compliance (default 7 days)

            *delivery_partners*: string based tuple of Delivery Partner
            names to limit result set against.  For example,
            ``['top', 'toll']``.  The values supported are as per
            the ``delivery_partner.name`` table set

            *alias*: table alias (default ``ji``)

        **Returns:**
            the SQL string

        """
        now = datetime.datetime.now()
        ts = now - datetime.timedelta(days=period)
        date = ts.strftime('%Y-%m-%d %H:%M:%S')

        dps = delivery_partners
        kwargs = {'period': period,
                  'delivery_partners': dps}
        compliance_sql = self._agent_stocktake.compliance_sql(**kwargs)

        sql = """%(compliance_sql)s
AND ag.id IN
(SELECT DISTINCT(j.agent_id)
 FROM job as j, %(name)s AS %(alias)s, agent AS ag
 WHERE %(alias)s.job_id = j.id
 AND %(alias)s.created_ts < '%(date)s'
 AND %(alias)s.pickup_ts IS NULL)""" % {'compliance_sql': compliance_sql,
                                        'name': self.name,
                                        'alias': alias,
                                        'date': date}

        return sql


import numpy as np
import cv2
import glob
import os
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--rgb_path",     type=str,   help="raw image path", default='./data_augmentation/raw_data/bg_img/select/total/')
# parser.add_argument("--mask_path",     type=str,   help="raw image path", default='./data_augmentation/raw_data/bg_img/select/select_rgb/')
# parser.add_argument("--obj_mask_path",     type=str,   help="raw image path", default='./data_augmentation/raw_data/obj_mask/')
# parser.add_argument("--bg_path",     type=str,   help="raw image path", default='./data_augmentation/raw_data/bg_img/select/select_rgb/')
parser.add_argument("--output_path",     type=str,   help="raw image path", default='./data_augmentation/raw_data/bg_img/save_bg/')

args = parser.parse_args()

if __name__ == '__main__':

    os.makedirs(args.output_path, exist_ok=True)

    rgb_path = os.path.join(args.rgb_path, '*.jpg')
    rgb_list = glob.glob(rgb_path)

    idx = 0
    for rgb_idx in rgb_list:
        print(rgb_idx)
        idx += 1
        file_name = rgb_idx.split('/')[5]
        file_name = file_name.split('.')[0]

        # if not os.path.isfile(args.obj_mask_path + file_name + '.png'):
        rgb_img = cv2.imread(rgb_idx)
        zero_mask = np.zeros(rgb_img.shape[:2], dtype=np.uint8)
        zero_mask = np.expand_dims(zero_mask, axis=-1)

        cv2.imwrite(args.output_path + 'rgb/' + 'bg2_image_idx_{0}_'.format(idx) + '_rgb.jpg', rgb_img)
        cv2.imwrite(args.output_path + 'gt/' + ' bg2_image_idx_{0}_'.format(idx) + '_mask.png', zero_mask)
        # cv2.imwrite(args.obj_mask_path + file_name + '.png', zero_mask)


#!flask/bin/python
from flask import Flask, jsonify, request, json
import requests
import time
from core import test_class
from BeautifulSoup import BeautifulSoup

app = Flask(__name__)

@app.route('/ytc/search', methods=['GET'])
def get_tasks():
    start_time = time.time()
    q1 = request.args.get('q')
    url = 'https://www.google.co.in/search?q='+q1
    response = requests.get(url)
    html = response.content
    soup = BeautifulSoup(html)
    div = soup.find('div', attrs={'id': 'res'})
    jsonArr = []
    for row in div.findAll('cite'):
        jsonArr.append(row.text)

    print("--- %s seconds ---" % (time.time() - start_time))
    return jsonify({'links': jsonArr})

@app.route('/ytc/redosearch', methods=['GET'])
def get_tasks_redo():
    w = 'wiki'
    yt = 'youtuber'
    start_time = time.time()
    q1 = request.args.get('q')
    jsonArr = []
    url = 'https://www.google.co.in/search?q='+q1+w
    response = requests.get(url)
    html = response.content
    soup = BeautifulSoup(html)
    div = soup.find('div', attrs={'id': 'res'})
    for row in div.findAll('cite'):
        jsonArr.append(row.text)

    url = 'https://www.google.co.in/search?q='+q1+yt
    response = requests.get(url)
    html = response.content
    soup = BeautifulSoup(html)
    div = soup.find('div', attrs={'id': 'res'})
    for row in div.findAll('cite'):
        jsonArr.append(row.text)

    jsonArr =  [k for k in jsonArr if 'wikipedia' in k]
    print("--- %s seconds ---" % (time.time() - start_time))
    return jsonify({'links': jsonArr})

tasks = [
    {
        'Approve': 72,
        'Spam': 18.5,
        'Hate Speech': 2.3,
        'Far Right': 0,
        'Harassment' : 7.2
    }
]

@app.route('/pace/predict', methods=['GET'])
def get_pace_data():
    return jsonify(tasks)

@app.route('/pace/message', methods = ['POST'])
def api_message():

    if request.headers['Content-Type'] == 'text/plain':
        return "Text Message: " + request.data

    elif request.headers['Content-Type'] == 'application/json':
        return "JSON Message: " + json.dumps(request.json)

    else:
        return "415 Unsupported Media Type ;)"

@app.route('/test', methods=['GET'])
def get_test():
    f = test_class.Fridge()
    return jsonify(f.in_fridge('apples'))

if __name__ == '__main__':
    app.run(port=5001, debug=True)


import ssl
import re
import traceback
from urllib.request import Request, urlopen
import requests
import json
from careerjet_api import CareerjetAPIClient
from bs4 import BeautifulSoup, SoupStrainer
from concurrent.futures import ThreadPoolExecutor

# ---------------------------------------------------------------------------- #
def getFullJobDesc(jobsDict):
    descriptions = []
    urls = []
    try:
        jobsArr = jobsDict['jobs']

        for element in jobsArr:
            urls.append(str(element['url']))

        print("\nUrls:", urls[0:4])

        def get_url(url):
            return requests.get(url).text

        with ThreadPoolExecutor(max_workers=50) as pool:
            resList = list(pool.map(get_url, urls))
            print("Threads done")

        for element in resList:
            soup = BeautifulSoup(element, 'html.parser')
            print(soup.prettify())
            pageNumberArr = soup.find_all("script")[1]
            pageNumberArr = str(pageNumberArr).replace('",
                                                                                "")
            res = json.loads(pageNumberArr)['description']
            res = res.split('
')
            print("\nInitial parsing finished")
            for part in res:
                if part == ' ' or part == '':
                    res.remove(part)

            res = "".join(res)
            description = res.replace("", "").replace("", "").replace("
", "").replace("
",
                                                                                "").replace(
                "
", "").replace("
", "").replace("Posting Notes:", "")
            print("\nDescription is cleaned")
            descriptions.append(description)
        return descriptions

    except Exception as err:
        traceback.print_exc()
        return " "

# ---------------------------------------------------------------------------- #
if __name__ == "__main__":

    try:
        cj = CareerjetAPIClient("en_US")
        location = 'san francisco'
        keywords = 'data scientist'
        pageNum = '0'

        if keywords is not None or location is not None:
            result_json = cj.search({
                'location': location,
                'keywords': keywords,
                'affid': '5f41b19494be1f6eaab8a755b612e343',
                'user_ip': '11.22.33.44',
                'url': 'http://www.example.com/jobsearch?q=python&l=london',
                'user_agent': 'Mozilla/5.0 (X11; Linux x86_64; rv:31.0) Gecko/20100101 Firefox/31.0',
                'pagesize': '99',
                'page': pageNum,
            })

    except Exception as e:
        print(f'API query failed\nFailed to retrieve jobs from CareerJet API with info: loc={location}, keywords={keywords}, pageNum={pageNum}')
        print(f'Error:\n{e}')

    if len(result_json['jobs']) != 99 or type(result_json) is not dict:
        raise Exception('CareerJet API response did not pass testing. Response:', result_json)

    print('Size of jobs:', len(result_json['jobs']))

    descriptions = getFullJobDesc(result_json)

    print(descriptions)


import threading

import mock
import pytest

from blazingmq import BasicHealthMonitor
from blazingmq import Error
from blazingmq import QueueOptions
from blazingmq import Session
from blazingmq import session_events
from blazingmq.testing import HostHealth

def test_receiving_host_health_events():
    # GIVEN
    spy = mock.MagicMock()
    host_health = BasicHealthMonitor()
    toggled_health_twice = threading.Event()
    host_healthy_events = []

    def callback(*args):
        spy(*args)

        if isinstance(args[0], session_events.HostHealthRestored):
            host_healthy_events.append(args[0])

        if len(host_healthy_events) == 2:
            toggled_health_twice.set()

    # WHEN
    session = Session(callback, host_health_monitor=host_health)
    host_health.set_unhealthy()
    host_health.set_healthy()
    host_health.set_unhealthy()
    host_health.set_healthy()
    toggled_health_twice.wait()
    session.stop()

    # THEN
    assert spy.call_args_list == [
        mock.call(session_events.Connected(None)),
        mock.call(session_events.HostUnhealthy(None)),
        mock.call(session_events.HostHealthRestored(None)),
        mock.call(session_events.HostUnhealthy(None)),
        mock.call(session_events.HostHealthRestored(None)),
        mock.call(session_events.Disconnected(None)),
    ]

def test_enabling_real_host_health_monitoring():
    # GIVEN
    spy = mock.MagicMock()

    def callback(*args):
        spy(*args)

    # WHEN
    session = Session(callback)
    session.stop()

    # THEN
    assert spy.call_args_list == [
        mock.call(session_events.Connected(None)),
        mock.call(session_events.Disconnected(None)),
    ]

def test_disabling_host_health_monitoring():
    # GIVEN
    spy = mock.MagicMock()

    def callback(*args):
        spy(*args)

    # WHEN
    session = Session(callback, host_health_monitor=None)
    session.stop()

    # THEN
    assert spy.call_args_list == [
        mock.call(session_events.Connected(None)),
        mock.call(session_events.Disconnected(None)),
    ]

def test_queue_suspension(unique_queue):
    # GIVEN
    host_health = HostHealth()
    queue_suspended_event_received = threading.Event()

    def on_session_event(event):
        print(event)
        if isinstance(event, session_events.QueueSuspended):
            queue_suspended_event_received.set()

    session = Session(on_session_event, host_health_monitor=host_health)

    session.open_queue(
        unique_queue,
        read=False,
        write=True,
        options=QueueOptions(suspends_on_bad_host_health=True),
    )

    # WHEN
    session.post(unique_queue, b"blah")

    host_health.set_unhealthy()
    queue_suspended_event_received.wait()

    with pytest.raises(Exception) as exc:
        session.post(unique_queue, b"blah")

    session.stop()

    # THEN
    assert exc.type is Error
    assert exc.match("QUEUE_SUSPENDED")


"""
Code for running experiments for
"Asymmetric Tri-training for Debiasing Missing-Not-At-Random Explicit Feedback".
"""
import argparse
import yaml
import warnings

import tensorflow as tf

from trainer import Trainer, Tuner

possible_model_names = ['uniform', 'uniform-at', 'user', 'user-at', 'item', 'item-at',
                        'both', 'both-at', 'nb', 'nb-at', 'nb_true', 'nb_true-at']

parser = argparse.ArgumentParser()
parser.add_argument('--data', '-d', type=str, required=True)
parser.add_argument('--model_name', '-m', type=str,
                    choices=possible_model_names, required=True)
parser.add_argument('--tuning', '-t', action='store_true')

if __name__ == "__main__":
    warnings.filterwarnings("ignore")
    tf.get_logger().setLevel("ERROR")
    args = parser.parse_args()

    # hyper-parameters
    config = yaml.safe_load(open('../config.yaml', 'r'))
    eta = config['eta']
    batch_size = config['batch_size']
    max_iters = config['max_iters']
    pre_iters = config['pre_iters']
    post_steps = config['post_steps']
    post_iters = config['post_iters']
    num_sims = config['num_sims']
    n_trials = config['n_trials']
    model_name = args.model_name
    tuning = args.tuning
    data = args.data

    if tuning:
        tuner = Tuner(data=data, model_name=model_name)
        tuner.tune(n_trials=n_trials)

        print('\n', '=' * 25, '\n')
        print(f'Finished Tuning of {model_name}!')
        print('\n', '=' * 25, '\n')

    trainer = Trainer(data=data, batch_size=batch_size, max_iters=max_iters,
                      pre_iters=pre_iters, post_steps=post_steps,
                      post_iters=post_iters, eta=eta, model_name=model_name)
    trainer.run_simulations(num_sims=num_sims)

    print('\n', '=' * 25, '\n')
    print(f'Finished Running {model_name}!')
    print('\n', '=' * 25, '\n')


import logging as logger
import os

from pyspark.sql import DataFrame

from ddataflow.exceptions import BiggerThanMaxSize
from ddataflow.sampling.default import filter_function
from ddataflow.utils import get_or_create_spark

class DataSource:
    """
    Utility functions at data source level
    """

    def __init__(
        self,
        *,
        name: str,
        config: dict,
        local_data_folder: str,
        snapshot_path: str,
        size_limit,
    ):
        self._name = name
        self._local_data_folder = local_data_folder
        self._snapshot_path = snapshot_path
        self._size_limit = size_limit
        self._config = config
        self._filter = None
        self._source = None

        if "source" in self._config:
            self._source = config["source"]
        else:
            if self._config.get("file-type") == "parquet":
                self._source = lambda spark: spark.read.parquet(self._name)
            else:
                self._source = lambda spark: spark.table(self._name)

        if "filter" in self._config:
            self._filter = self._config["filter"]
        else:
            if self._config.get("default_sampling"):
                self._filter = lambda df: filter_function(df)

    def query(self):
        """
        query with filter unless none is present
        """
        df = self.query_without_filter()

        if self._filter is not None:
            print(f"Filter set for {self._name}, applying it")
            df = self._filter(df)
        else:
            print(f"No filter set for {self._name}")

        return df

    def has_filter(self) -> bool:
        return self._filter is not None

    def query_without_filter(self):
        """
        Go to the raw data source without any filtering
        """
        spark = get_or_create_spark()
        logger.debug(f"Querying without filter source: '{self._name}'")
        return self._source(spark)

    def query_locally(self):
        logger.info(f"Querying locally {self._name}")

        path = self.get_local_path()
        if not os.path.exists(path):
            raise Exception(
                f"""Data source '{self.get_name()}' does not have data in {path}.
            Consider downloading using  the following command:
            ddataflow current_project download_data_sources"""
            )
        spark = get_or_create_spark()
        df = spark.read.parquet(path)

        return df

    def get_dbfs_sample_path(self) -> str:
        return os.path.join(self._snapshot_path, self._get_name_as_path())

    def get_local_path(self) -> str:
        return os.path.join(self._local_data_folder, self._get_name_as_path())

    def _get_name_as_path(self):
        """
        converts the name when it has "/mnt/envents" in the name to a single file in a (flat structure) _mnt_events
        """
        return self.get_name().replace("/", "_")

    def get_name(self) -> str:
        return self._name

    def get_parquet_filename(self) -> str:
        return self._name + ".parquet"

    def estimate_size_and_fail_if_too_big(self):
        """
        Estimate the size of the data source use the _name used in the _config
        It will throw an exception if the estimated size is bigger than the maximum allowed in the configuration
        """

        print("Estimating size of data source: ", self.get_name())
        df = self.query()
        size_estimation = self._estimate_size(df)

        print("Estimated size of the Dataset in GB: ", size_estimation)

        if size_estimation > self._size_limit:
            raise BiggerThanMaxSize(self._name, size_estimation, self._size_limit)

        return df

    def _estimate_size(self, df: DataFrame) -> float:
        """
        Estimates the size of a dataframe in Gigabytes

        Formula:
            number of gigabytes = (N*V*W) / 1024^3
        """

        print(f"Amount of rows in dataframe to estimate size: {df.count()}")
        average_variable_size_bytes = 50
        return (df.count() * len(df.columns) * average_variable_size_bytes) / (
            1024**3
        )


import dearpygui.dearpygui as dpg
from data import extract_data, get_path2
dpg.create_context()

path2data = get_path2("Test_spectrum.syn")

x, y, _ = extract_data(path2data)

# Screen resolution for window scaling. dpi for correct font size.
res_x, res_y = 1920, 1080
dpi = 150

with dpg.font_registry():
    # first argument ids the path to the .ttf or .otf file
    default_font = dpg.add_font("/usr/share/fonts/OTF/CodeNewRomanNerdFont-Bold.otf", 20)

with dpg.window(label="Spectra", width=int(res_x*0.8), height=int(res_y * 0.8)):
    # themes part
    dpg.bind_font(default_font)
    with dpg.theme(tag="spectrum_theme_1"):
        with dpg.theme_component(1):
            dpg.add_theme_color(dpg.mvPlotCol_Line, (231, 0, 230),
                                category=dpg.mvThemeCat_Plots)
            dpg.add_theme_color(dpg.mvPlotCol_Fill, (230, 0, 230, 170),
                                category=dpg.mvThemeCat_Plots)

    # spectrum plot part
    with dpg.plot(label="Spectrum plot", height=int(res_y*0.8), width=-1):
        dpg.add_plot_legend()
        dpg.add_plot_axis(dpg.mvXAxis, label="x")
        dpg.add_plot_axis(dpg.mvYAxis, label="y", tag="yaxis")

        dpg.add_shade_series(list(x), list(y), label="Synthetic spectra",
                             parent="yaxis", tag="syn_spectrum")
        # apply theme
        dpg.bind_item_theme(dpg.last_item(), "spectrum_theme_1")

dpg.create_viewport(title='Extractor', width=int(res_x*0.9),
                    height=int(res_y*0.9))

dpg.setup_dearpygui()
dpg.show_viewport()
dpg.start_dearpygui()
dpg.destroy_context()


import os
import sys
import warnings
import platform
import datetime as dt
import numpy as np
import matplotlib.pyplot as mp
import matplotlib.dates as md

def dmy2ymd(dmy):
    return dt.datetime.strptime(
        str(dmy, encoding='utf-8'),
        '%d-%m-%Y').date().strftime('%Y-%m-%d')

def read_data(filename):
    dates, closing_prices = np.loadtxt(
        filename, delimiter=',', usecols=(1, 6),
        unpack=True, dtype=np.dtype('M8[D], f8'),
        converters={1: dmy2ymd})
    return dates, closing_prices

def calc_returns(N, closing_prices):
    returns = np.diff(closing_prices) / closing_prices[:-1]
    weights = np.hanning(N)
    smr = np.convolve(weights, returns, 'valid')
    return returns, smr

def fit_polys(fit_x, fit_y, fit_d, poly_x):
    fit_p = np.polyfit(fit_x, fit_y, fit_d)
    poly_y = np.polyval(fit_p, poly_x)
    return poly_y

def find_inters(fit_x, fit_y1, fit_y2,
                fit_d, min_x, max_x):
    fit_p1 = np.polyfit(fit_x, fit_y1, fit_d)
    fit_p2 = np.polyfit(fit_x, fit_y2, fit_d)
    fit_p3 = np.polysub(fit_p1, fit_p2)
    roots = np.roots(fit_p3)
    reals = roots[np.isreal(roots)].real
    inters = []
    for real in reals:
        if min_x < real and real < max_x:
            inters.append([real,
                           np.polyval(fit_p1, real)])
    inters.sort()
    inters = np.array(inters)
    return inters

def init_chart(N, first_day, last_day):
    mp.gcf().set_facecolor(np.ones(3) * 240 / 255)
    mp.title('Smoothing Returns (%d Days)' % N, fontsize=20)
    mp.xlabel('Trading Days From %s To %s' % (
        first_day.astype(md.datetime.datetime).strftime(
            '%d %b %Y'),
        last_day.astype(md.datetime.datetime).strftime(
            '%d %b %Y')), fontsize=14)
    mp.ylabel('Returns Of Stock Price', fontsize=14)
    ax = mp.gca()
    ax.xaxis.set_major_locator(
        md.WeekdayLocator(byweekday=md.MO))
    ax.xaxis.set_minor_locator(md.DayLocator())
    ax.xaxis.set_major_formatter(
        md.DateFormatter('%d %b %Y'))

    mp.tick_params(which='both', top=True, right=True,
                   labelright=True, labelsize=10)
    mp.grid(linestyle=':')

def draw_returns(dates, bhp_returns, vale_returns):
    dates = dates.astype(md.datetime.datetime)
    mp.plot(dates, bhp_returns, '-', c='orangered',
            alpha=0.25, label='BHP Returns')
    mp.plot(dates, vale_returns, '-', c='dodgerblue',
            alpha=0.25, label='VALE Returns')

def draw_smrs(N, dates, bhp_smrs, vale_smrs):
    dates = dates.astype(md.datetime.datetime)
    mp.plot(dates, bhp_smrs, '-', c='orangered',
            alpha=0.75, label='BHP SMR-%d' % N)
    mp.plot(dates, vale_smrs, '-', c='dodgerblue',
            alpha=0.75, label='VALE SMR-%d)' % N)

def draw_polys(N, dates, bhp_polys, vale_polys, degree):
    dates = dates.astype(md.datetime.datetime)
    mp.plot(dates, bhp_polys, '-', c='orangered',
            label='BHP SMR-%d Polynomial (%d)' % (
                N, degree))
    mp.plot(dates, vale_polys, '-', c='dodgerblue',
            label='VALE SMR-%d Polynomal (%d)' % (
                N, degree))

def draw_inters(inters):
    dates, inters = np.hsplit(inters, 2)
    dates = dates.astype(int).astype(
        'M8[D]').astype(md.datetime.datetime)
    mp.scatter(dates, inters, marker='X', s=120,
               c='firebrick',
               label='Intersection of SMRs', zorder=3)

    mp.gcf().autofmt_xdate()
    mp.legend()

def show_chart():
    mng = mp.get_current_fig_manager()
    if 'Windows' in platform.system():
        mng.window.state('zoomed')
    else:
        mng.resize(*mng.window.maxsize())
    mp.show()

def main(argc, argv, envp):
    warnings.filterwarnings('ignore',
                            category=np.RankWarning)
    dates, bhp_closing_prices = read_data('bhp.csv')
    dates, vale_closing_prices = read_data('vale.csv')
    N = 8
    bhp_returns, bhp_smrs = calc_returns(
        N, bhp_closing_prices)
    vale_returns, vale_smrs = calc_returns(
        N, vale_closing_prices)
    days = dates[N - 1:-1].astype(int)
    degree = 5
    bhp_ploys = fit_polys(days, bhp_smrs, degree, days)
    vale_ploys = fit_polys(days, vale_smrs, degree, days)
    inters = find_inters(days, bhp_smrs, vale_smrs,
                         degree, days[0], days[-1])
    init_chart(N, dates[0], dates[-2])
    draw_returns(dates[:-1], bhp_returns, vale_returns)
    draw_smrs(N, dates[N - 1:-1], bhp_smrs, vale_smrs)
    draw_polys(N, dates[N - 1:-1], bhp_ploys, vale_ploys,
               degree)
    draw_inters(inters)
    show_chart()
    return 0

if __name__ == '__main__':
    sys.exit(main(len(sys.argv), sys.argv, os.environ))


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"View module showing all messages concerning discarded beads"
from    typing               import cast
from    bokeh                import layouts
from    model.plots          import PlotState
from    taskview.plots       import PlotView, CACHE_TYPE, TaskPlotCreator
from    utils                import initdefaults
from    view.base            import stretchout
from    ._widgets            import QualityControlWidgets
from    ._plots              import QualityControlPlots
from    ._model              import QualityControlModelAccess

class _StateDescriptor:
    def __get__(self, inst, owner):
        return getattr(inst, '_state').state if inst else self

    @staticmethod
    def setdefault(inst, value):
        "sets the default value"
        getattr(inst, '_ctrl').display.updatedefaults("qc.state", state = PlotState(value))

    def __set__(self, inst, value):
        getattr(inst, '_ctrl').display.update("qc.state", state = PlotState(value))

class QCPlotState:
    "qc plot state"
    state = PlotState.active
    name  = "qc.state"
    @initdefaults(frozenset(locals()))
    def __init__(self, **_):
        pass

class QualityControlPlotCreator(TaskPlotCreator[QualityControlModelAccess, None]):
    "Creates plots for discard list"
    _plotmodel: None
    _model:     QualityControlModelAccess
    _RESET:     frozenset = frozenset()
    state = cast(PlotState, _StateDescriptor())

    def __init__(self, ctrl):
        super().__init__(ctrl)
        self._widgets = QualityControlWidgets(ctrl, self._model)
        self._plots   = QualityControlPlots  (ctrl, self._model)
        self._state   = QCPlotState()
        ctrl.display.add(self._state)
        self.addto(ctrl)

    def observe(self, ctrl):
        "observes the model"
        super().observe(ctrl)
        self._plots  .observe(ctrl)
        self._widgets.observe(self, ctrl)

    def _addtodoc(self, ctrl, doc, *_):
        "returns the figure"
        mode    = self.defaultsizingmode()
        widgets = self._widgets.addtodoc(self, ctrl, mode)
        grid    = self._plots.addtodoc(self._ctrl, doc, mode)
        out     = layouts.row(grid, widgets, **mode)
        self.__resize(ctrl, out)
        return stretchout(out)

    def _reset(self, cache:CACHE_TYPE):
        self._widgets.reset(cache)
        self._plots.reset(cache)

    def __resize(self, ctrl, sizer):
        figtb   = ctrl.theme.get("theme", "figtbheight")
        borders = ctrl.theme.get("theme", "borders")

        sizer.update(**self.defaulttabsize(ctrl))

        widg  = sizer.children[1]
        width = max(i.width for i in widg.children)
        for i in widg.children:
            i.width = width
        widg.children[-1].height = (
            sizer.height - sum(i.height for i in widg.children[:-1])-figtb
        )
        widg.update(width  = width, height = sizer.height)

        sizer.children[0].update(width = sizer.width-width, height = sizer.height)
        sizer.children[0].children[1].update(
            width  = sizer.width-width-borders,
            height = sizer.height
        )
        plots = sizer.children[0].children[1].children[1].children
        for i in plots:
            i[0].update(
                plot_width  = sizer.children[0].children[1].width,
                plot_height = (sizer.height-figtb)//len(plots)
            )

class QualityControlView(PlotView[QualityControlPlotCreator]):
    "a widget with all discards messages"
    TASKS       = 'datacleaning', 'extremumalignment'
    PANEL_NAME  = 'Quality Control'
    def ismain(self, ctrl):
        "Cleaning and alignment, ... are set-up by default"
        self._ismain(ctrl, tasks = self.TASKS)