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

"""
GitHub key downloader.
We use this for academic research on SSH keys entropy.
"""

import os
import sys
import inspect
import resource

from requests.auth import HTTPBasicAuth

currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
parentdir = os.path.dirname(currentdir)
sys.path.insert(0, parentdir)

import requests
import logging
import coloredlogs
import traceback
import json
import argparse
import re
import math
import random
import time
import shutil
import multiprocessing
import time
import signal
import utils
import databaseutils
import collections
import threading
from threading import Lock as Lock
import types
import Queue
from blessed import Terminal
from cmd2 import Cmd
from evt_dequeue import EvtDequeue
import dbutil

from github_base import AccessResource, RateLimitHit
from database import GitHubKey, GitHubUser as GitHubUserDb, GitHubUserKeys
from database import Base as DB_Base
from sqlalchemy.orm import scoped_session
import sqlalchemy as salch
from trace_logger import Tracelogger

import gc
import mem_top
from pympler.tracker import SummaryTracker
from collections import OrderedDict, namedtuple
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey

logger = logging.getLogger(__name__)
coloredlogs.install(level=logging.INFO)

# GitHub User ID & Name
GitHubUser = namedtuple('GitHubUser', ['user_id', 'user_name', 'user_type', 'user_url'])

class DownloadJob(object):
    """
    Represents link to download
    """

    TYPE_USERS = 1
    TYPE_KEYS = 2

    __slots__ = ['url', 'type', 'user', 'fail_cnt', 'last_fail', 'priority', 'time_added']

    def __init__(self, url=None, jtype=TYPE_USERS, user=None, priority=0, time_added=None, *args, **kwargs):
        self.url = url
        self.type = jtype
        self.user = user
        self.fail_cnt = 0
        self.last_fail = 0
        self.priority = priority
        self.time_added = time.time() if time_added is None else time_added

    def to_json(self):
        js = collections.OrderedDict()
        js['url'] = self.url
        js['type'] = self.type
        js['fail_cnt'] = self.fail_cnt
        js['last_fail'] = self.last_fail
        js['priority'] = self.priority
        js['time_added'] = self.time_added
        if self.user is not None:
            js['user_id'] = self.user.user_id
            js['user_name'] = self.user.user_name
            js['user_type'] = self.user.user_type
            js['user_url'] = self.user.user_url
        return js

    @classmethod
    def from_json(cls, js):
        tj = cls()
        tj.url = js['url']
        tj.type = js['type']
        tj.fail_cnt = js['fail_cnt']
        tj.last_fail = js['last_fail']
        tj.priority = utils.defvalkey(js, 'priority', 0)
        tj.time_added = utils.defvalkey(js, 'time_added', 0)
        if 'user_id' in js:
            user_url = js['user_url'] if 'user_url' in js else None
            tj.user = GitHubUser(user_id=js['user_id'], user_name=js['user_name'], user_type=js['user_type'], user_url=user_url)
        return tj

    @staticmethod
    def cmp(self, other):
        """
        Comparator
        :param self:
        :param other:
        :return:
        """
        # Inside the category: fail cnt, time added.
        if self.type == other.type:
            if self.fail_cnt == other.fail_cnt:
                if self.time_added == other.time_added:
                    return int(other.priority - self.priority)
                else:
                    return int(self.time_added - other.time_added)
            else:
                return int(self.fail_cnt - other.fail_cnt)
        else:
            # Outside the category - priority ordering. Higher the priority, sooner will be picked
            if self.priority == other.priority:
                return int(self.time_added - other.time_added)
            else:
                return int(other.priority - self.priority)

    def __cmp__(self, other):
        """
        Compare operation for priority queue.
        :param other:
        :return:
        """
        return self.cmp(self, other)

class GitHubLoader(Cmd):
    """
    GitHub SSH keys loader
    """
    prompt = '$> '

    LINK_FACTOR = 70
    USER_URL = 'https://api.github.com/users/%s'
    USERS_URL = 'https://api.github.com/users?since=%s'
    KEYS_URL = 'https://api.github.com/users/%s/keys'
    KEYS_ID_URL = 'https://api.github.com/user/%s/keys'

    def __init__(self, attempts=5, threads=1, state=None, state_file=None, config_file=None, audit_file=None,
                 max_mem=None, users_only=False, merge=False, update_keys=False, *args, **kwargs):

        Cmd.__init__(self, *args, **kwargs)
        self.t = Terminal()
        self.trace_logger = Tracelogger(logger=logger)

        self.attempts = int(attempts)
        self.total = None
        self.terminate = False
        self.since_id = 0
        self.last_users_count = None
        self.user_lock = Lock()
        self.processed_user_set = set()
        self.processed_user_set_lock = Lock()

        self.max_mem = max_mem
        self.merge = merge
        self.update_keys = update_keys
        self.users_only = users_only
        self.users_per_page = 30
        self.users_bulk_load_pages = 500
        self.user_load_bulk = 5000
        self.user_refill_lock = Lock()
        self.state = state
        self.state_file_path = state_file
        self.rate_limit_reset = None
        self.rate_limit_remaining = None

        self.config = None
        self.config_file = config_file

        self.audit_file = audit_file
        self.audit_records_buffered = []
        self.audit_lock = Lock()

        self.stop_event = threading.Event()
        self.threads = int(threads)
        self.link_queue = Queue.PriorityQueue()  # Store links to download here
        self.worker_threads = []

        self.state_thread = None
        self.state_thread_lock = Lock()

        self.resources_list = []
        self.resources_queue = Queue.PriorityQueue()
        self.local_data = threading.local()

        self.new_users_events = EvtDequeue()
        self.new_keys_events = EvtDequeue()

        self.db_config = None
        self.engine = None
        self.session = None

        self.mem_tracker = None

    def signal_handler(self, signal, frame):
        """
        Signal handler - terminate gracefully
        :param signal:
        :param frame:
        :return:
        """
        logger.info('CTRL+C pressed')
        self.trigger_stop()

    def trigger_stop(self):
        """
        Sets terminal conditions to true
        :return:
        """
        self.terminate = True
        self.stop_event.set()

    def trigger_quit(self):
        """
        Terminal condition & file change
        :return:
        """
        self.trigger_stop()
        utils.try_touch('.github-quit')

    #
    # CMD handlers
    #

    def do_quit(self, arg):
        self.trigger_quit()
        logger.info('Waiting for thread termination')

        time.sleep(1)
        logger.info('Quitting')
        return Cmd.do_quit(self, arg)

    def do_reset(self, line):
        print('\033c')

    def do_gc(self, line):
        gc.collect()

    def do_mem_top(self, line):
        print(mem_top.mem_top())

    def do_mem_track_init(self, line):
        self.mem_tracker = SummaryTracker()

    def do_mem_track_diff(self, line):
        print(self.mem_tracker.print_diff())

    def do_mem_track_deinit(self, line):
        self.mem_tracker = None

    def do_mem(self, line):
        print('Memory usage: %s kB' % resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)

    def do_state(self, line):
        js = self.state_gen()

        if line is None or len(line) == 0:
            del js['link_queue']
            del js['resource_stats']
        elif line == '1':
            del js['link_queue']

        print(json.dumps(js, indent=2))

    def do_deq_enable(self, line):
        self.new_keys_events.disabled = False
        self.new_users_events.disabled = False

    def do_deq_disable(self, line):
        self.new_keys_events.disabled = True
        self.new_users_events.disabled = True

    #
    # Init
    #

    def init_config(self):
        """
        Loads config & state files
        :return:
        """
        if self.state_file_path is not None and os.path.exists(self.state_file_path):
            with open(self.state_file_path, 'r') as fh:
                self.state = json.load(fh, object_pairs_hook=OrderedDict)
                logger.info('State loaded: %s' % os.path.abspath(self.state_file_path))

        with open(self.config_file, 'r') as fh:
            self.config = json.load(fh, object_pairs_hook=OrderedDict)
            logger.info('Config loaded: %s' % os.path.abspath(self.config_file))

            if 'since_id' in self.config:
                self.since_id = self.config['since_id']

            # Process resources - randomized
            if 'res' in self.config:
                res_tmp = self.config['res']
                random.shuffle(res_tmp)
                for idx, res in enumerate(res_tmp):
                    r = AccessResource(usr=res['usr'], token=res['token'], idx=idx)
                    self.resources_list.append(r)
                    self.resources_queue.put(r)
                    logger.info('Resource %02d loaded: %s' % (idx, r.usr))
            else:
                # unauth
                r = AccessResource(usr=None, token=None)
                self.resources_list.append(r)
                self.resources_queue.put(r)

    def init_db(self):
        """
        Initializes database engine & session.
        Has to be done on main thread.
        :return:
        """
        self.db_config = databaseutils.process_db_config(self.config['db'])

        from sqlalchemy import create_engine
        from sqlalchemy.orm import sessionmaker, scoped_session
        self.engine = create_engine(self.db_config.constr, pool_recycle=3600)
        self.session = scoped_session(sessionmaker(bind=self.engine))

        # Make sure tables are created
        DB_Base.metadata.create_all(self.engine)

    def init_workers(self):
        """
        Initialize worker threads
        :return:
        """
        logger.info('Starting %d working threads' % self.threads)
        for idx in range(self.threads):
            t = threading.Thread(target=self.work_thread_main, args=(idx, ))
            self.worker_threads.append(t)

        # Kick-off all threads
        for t in self.worker_threads:
            t.start()

        logger.info('Worker threads started')

    def cli(self):
        """
        CLI thread
        :return:
        """
        logger.info('CLI thread started')
        self.cmdloop()
        logger.info('Terminating CLI thread')

    #
    # Operation
    #

    def work(self):
        """
        Main thread work method
        :return:
        """
        # Interrupt signals
        signal.signal(signal.SIGINT, self.signal_handler)

        self.init_config()
        self.init_db()

        # Resume last state
        self.state_resume()

        # Monitor threads.
        self.state_thread = threading.Thread(target=self.state_main, args=())
        self.state_thread.start()

        # If there is no link to process - create from since.
        if self.link_queue.qsize() == 0:
            self.kickoff_links()

        # Worker threads
        self.init_workers()

        logger.info('Main thread started %s %s %s' % (os.getpid(), os.getppid(), threading.current_thread()))

        # CLI thread
        cli_thread = threading.Thread(target=self.cli, args=())
        cli_thread.setDaemon(True)
        cli_thread.start()

        # Join on workers
        self.after_loop()
        logger.info('Terminating main thread')
        return None

    def after_loop(self, wait_for_state=True):
        """
        After work loop finishes
        :return:
        """
        logger.info('Waiting termination of slave threads')

        # Wait here for termination of all workers and monitors.
        try:
            for t in self.worker_threads:
                t.join()

            if wait_for_state:
                self.state_thread.join()
        except:
            logger.error('Exception during thread join')
            logger.error(traceback.format_exc())

        logger.info('All threads terminates, last state save')
        self.state_save()

    def work_thread_main(self, idx):
        """
        Worker thread main loop
        :return:
        """
        self.local_data.idx = idx
        logger.info('Working thread %d started' % idx)

        while not self.terminate and not self.stop_event.is_set():
            self.interruptible_sleep_delta(0.1)

            # Get credential to process link with
            resource = self.resource_allocate()
            if resource is None:
                continue

            # We have resource, now get the job
            job = None
            try:
                job = self.link_queue.get(True, timeout=1.0)
            except Queue.Empty:
                self.resource_return(resource)
                continue

            # If job last fail is too recent - put again back to queue
            if time.time() - job.last_fail < 3.0:
                self.link_queue.put(job)  # re-insert to the back of the queue for later processing
                self.resource_return(resource)
                continue

            # Job processing starts here - fetch data page with the resource.
            js_data = None
            try:
                self.local_data.job = job
                self.local_data.resource = resource
                js_data, headers, raw_response = self.load_page_local()

            except RateLimitHit as e:
                logger.error('[%d] Rate limit hit: %s, failcnt: %d, res: %s, exception: %s'
                             % (idx, job.url, job.fail_cnt, resource.usr, e))
                continue

            except Exception as e:
                logger.error('[%d] Exception in processing job: %s, failcnt: %d, res: %s, exception: %s'
                             % (idx, job.url, job.fail_cnt, resource.usr, e))

                self.on_job_failed(job)
                continue

            finally:
                self.resource_return(resource)
                self.local_data.resource = None
                self.local_data.last_usr = resource.usr
                self.local_data.last_remaining = resource.remaining

            # Process downloaded data here.
            try:
                if js_data is None:
                    self.audit_log('404', job.url, jtype=job.type, job=job)
                    self.flush_audit()
                    continue

                if job.type == DownloadJob.TYPE_USERS:
                    self.process_users_data(job, js_data, headers, raw_response)
                else:
                    self.process_keys_data(job, js_data, headers, raw_response)

            except Exception as e:
                logger.error('[%d] Unexpected exception, processing type %s, link %s: cnt: %d, res: %s, %s'
                             % (idx, job.type, job.url, job.fail_cnt, resource.usr, e))

                self.trace_logger.log(e)
                self.on_job_failed(job)
            finally:
                self.local_data.resource = None
                self.local_data.job = None
                self.local_data.last_usr = None
                self.local_data.last_remaining = None
                resource = None
                job = None
                headers = None
                raw_response = None

        pass
        logger.info('Terminating worker thread %d' % idx)

    def on_job_failed(self, job):
        """
        If job failed, this teaches it how to behave
        :param job:
        :return:
        """
        job.fail_cnt += 1
        job.last_fail = time.time()

        # if failed too many times - log and discard.
        if job.fail_cnt > 35:
            logger.warning('Job failed too many times %s' % job.url)
            self.audit_log('too-many-fails', job.url, jtype=job.type, job=job)
            self.flush_audit()
        else:
            self.link_queue.put(job)  # re-insert to the queue for later processing

    def load_page_local(self):
        """
        Loads page stored in thread local
        :return:
        """

        auth = None
        resource = self.local_data.resource
        if resource.usr is not None:
            auth = HTTPBasicAuth(resource.usr, resource.token)

        job = self.local_data.job

        res = requests.get(job.url, timeout=10, auth=auth)
        headers = res.headers

        resource.reset_time = utils.try_float(headers.get('X-RateLimit-Reset'))
        resource.remaining = utils.try_int(headers.get('X-RateLimit-Remaining'))
        resource.last_used = time.time()
        resource.used_cnt += 1

        if res.status_code == 403 and resource.remaining is not None and resource.remaining < 10:
            resource.fail_cnt += 1
            raise RateLimitHit

        if res.status_code == 404:
            resource.fail_cnt += 1
            logger.warning('URL not found: %s' % job.url)
            return None, None, None

        if res.status_code // 100 != 2:
            resource.fail_cnt += 1
            res.raise_for_status()

        data = res.content
        if data is None:
            resource.fail_cnt += 1
            raise Exception('Empty response')

        js = json.loads(data, object_pairs_hook=OrderedDict)
        return js, headers, res

    def resource_allocate(self, blocking=True, timeout=1.0):
        """
        Takes resource from the pool.
        If the resource has low remaining credit, thread is suspended to re-charge.
        :return: resource or None if not available in the time
        """
        try:
            resource = self.resources_queue.get(True, timeout=1.0)
            if resource.remaining is not None and resource.remaining <= self.threads + 2:
                sleep_sec = resource.reset_time - time.time() if resource.reset_time is not None else 15*60
                sleep_sec += 120  # extra 2 minutes to avoid problems with resources

                logger.info('Rate limit exceeded on resource %s, remaining: %s, sleeping till: %s, it is %d seconds, '
                            '%d minutes'
                            % (resource.usr, resource.remaining, resource.reset_time, sleep_sec, sleep_sec / 60.0))
                self.sleep_interruptible(time.time() + sleep_sec)
                logger.info('Resource sleep finished %s' % resource.usr)

                # Reset estimations, needs to be refreshed
                resource.remaining = None
                resource.reset_time = None

            return resource

        except Queue.Empty:
            return None

    def resource_return(self, res):
        """
        Returns resource to the pool
        :param res:
        :return:
        """
        self.resources_queue.put(res)

    def sleep_interruptible(self, until_time):
        """
        Interruptible sleep - sleep until given time.
        :param until_time:
        :return:
        """
        while time.time() <= until_time:
            time.sleep(1.0)
            if self.terminate or self.stop_event.is_set():
                return

    def interruptible_sleep_delta(self, sleep_time):
        """
        Sleeps the current thread for given amount of seconds, stop event terminates the sleep - to exit the thread.
        :param sleep_time:
        :return:
        """
        if sleep_time is None:
            return

        sleep_time = float(sleep_time)

        if sleep_time == 0:
            return

        sleep_start = time.time()
        while not self.stop_event.is_set() and not self.terminate:
            time.sleep(0.1)
            if time.time() - sleep_start >= sleep_time:
                return

    #
    # Parser and processing logic
    #

    def kickoff_links(self):
        """
        Kick off the scrapping by adding initial links to the queue
        :return:
        """
        if self.update_keys:
            self.fill_user_key_links()

        else:
            job = DownloadJob(url=self.USERS_URL % self.since_id, jtype=DownloadJob.TYPE_USERS)
            self.link_queue.put(job)
            logger.info('Kickoff link added: %s' % job.url)

    def process_users_data(self, job, js, headers, raw_response):
        """
        Process user data - produce keys links + next user link
        :param job:
        :param js:
        :param headers:
        :param raw_response:
        :return:
        """
        max_id = 0
        github_users = []
        cur_time = int(time.time())
        for user in js:
            if 'id' not in user:
                logger.error('Field ID not found in user')
                continue

            github_user = GitHubUser(user_id=int(user['id']), user_name=user['login'],
                                     user_type=user['type'], user_url=user['url'])
            github_users.append(github_user)

            if github_user.user_id > max_id:
                max_id = github_user.user_id

            if self.users_only:
                continue

            key_url = '%s/keys' % github_user.user_url
            new_job = DownloadJob(url=key_url, jtype=DownloadJob.TYPE_KEYS, user=github_user,
                                  priority=random.randint(0, 1000), time_added=cur_time)
            self.link_queue.put(new_job)

        # Link with the maximal user id
        users_url = self.USERS_URL % max_id
        new_job = DownloadJob(url=users_url, jtype=DownloadJob.TYPE_USERS, time_added=cur_time)

        # Optimizing the position of this link in the link queue
        queue_size = self.link_queue.qsize()
        queue_size_max = self.LINK_FACTOR * self.threads
        fill_up_ratio = queue_size / float(queue_size_max)

        # Key jobs are uniformly distributed on priorities 0...1000.
        # To increase queue size pick priority closer to 1000, do decrease, closer to 0
        priority = random.randint(0, 500)
        if queue_size < queue_size_max:
            priority = int((1 - fill_up_ratio) * 5000) + 500
        if queue_size > 3*queue_size_max:
            priority = 0

        new_job.priority = priority
        lucky_one = False
        with self.user_lock:
            if self.since_id < max_id:
                self.since_id = max_id
                self.link_queue.put(new_job)
                lucky_one = True

                # Bulk user optimisation - add more users at once, multithreading
                if self.users_only:
                    self.bulk_user_only_load(max_id=max_id, cur_time=cur_time, priority=priority)

        logger.info('[%02d, usr=%20s, remaining=%5s] Processed users link %s, Next since: %3s. ResQSize: %4d, '
                    'LQSize: %4d, fill-up: %0.4f, priority: %4s, ram: %s kB, new=%s, New users: [%s]'
                    % (self.local_data.idx, self.local_data.last_usr, self.local_data.last_remaining,
                       len(github_users)+1, max_id, self.resources_queue.qsize(),
                       queue_size, fill_up_ratio, priority,
                       resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, lucky_one,
                       ', '.join([str(x.user_name) for x in github_users])))

        # Store all users.
        self.store_users_list(github_users)

    def bulk_user_only_load(self, max_id, cur_time, priority):
        """
        Fills the queue with multiple user links
        :param max_id:
        :param cur_time:
        :param priority:
        :return:
        """
        new_jobs = []
        for page_idx in range(1, self.users_bulk_load_pages + 1):
            users_url = self.USERS_URL % (max_id + self.users_per_page * page_idx)
            new_job = DownloadJob(url=users_url, jtype=DownloadJob.TYPE_USERS,
                                  time_added=cur_time, priority=priority - page_idx)
            new_jobs.append(new_job)

        self.since_id = max_id + self.users_per_page * self.users_bulk_load_pages
        logger.info('[%02d] Bulk load, new max=%s' % (self.local_data.idx, self.since_id))

        for job in new_jobs:
            self.link_queue.put(job)

    def store_users_list(self, users):
        """
        Stores all user in the list
        :param users
        :return:
        """
        # Handling gaps in the user space ID. With user-only optimization it causes
        # overlaps.
        reduced_by = 0
        with self.processed_user_set_lock:
            ids = [user.user_id for user in users]
            ids_ok = []
            for id in ids:
                if id in self.processed_user_set:
                    reduced_by += 1
                    continue
                self.processed_user_set.add(id)
                ids_ok.append(id)
            users = [user for user in users if user.user_id in ids_ok]

        # Bulk user load
        s = self.session()
        id_list = sorted([user.user_id for user in users])
        db_users = s.query(GitHubUserDb).filter(GitHubUserDb.id.in_(id_list)).all()
        db_user_map = {user.id: user for user in db_users}

        for user in users:
            self.new_users_events.insert()

            # Store user to the DB
            try:
                db_user = utils.defvalkey(db_user_map, key=user.user_id)
                self.store_user(user, s, db_user=db_user, db_user_loaded=True)

            except Exception as e:
                logger.warning('[%02d] Exception in storing user %s' % (self.local_data.idx, e))
                self.trace_logger.log(e)
                logger.info('[%02d] idlist: %s' % (self.local_data.idx, id_list))
                self.trigger_quit()
                break

        try:
            s.commit()
            # logger.info('[%02d] Commited, reduced by: %s' % (self.local_data.idx, reduced_by))
        except Exception as e:
            logger.warning('[%02d] Exception in storing bulk users' % self.local_data.idx)
            logger.warning(traceback.format_exc())
            logger.info('[%02d] idlist: %s' % (self.local_data.idx, id_list))
            self.trigger_quit()
        finally:
            utils.silent_close(s)

    def process_keys_data(self, job, js, headers, raw_response):
        """
        Processing key loaded data
        :param job:
        :param js:
        :param headers:
        :param raw_response:
        :return:
        """
        js_empty = js is None or len(js) == 0

        # Expect failures, commit everything before
        if self.merge and not js_empty:
            try:
                s = self.session()
                s.commit()
            except Exception as e:
                logger.warning('Could not pre-commit: %s' % e)

        # Store each key.
        key_ids = []
        for key in js:
            s = None
            self.new_keys_events.insert()

            try:
                s = self.session()
                self.store_key(job.user, key, s)
                key_ids.append(int(key['id']))
                s.commit()

                self.assoc_key(job.user.user_id, key['id'], s)
                s.commit()

                s.flush()        # writes changes to DB
                s.expunge_all()  # removes objects from session

            except Exception as e:
                logger.warning('Exception in storing key %s' % e)
                self.trace_logger.log(e)

            finally:
                utils.silent_close(s)
                s = None

        # Deassoc lost keys
        try:
            s = self.session()
            self.deassoc_lost_keys(job.user.user_id, key_ids, s)
            s.commit()

        except Exception as e:
            logger.warning('Exception in deassoc for users %s : %s' % (job.user.user_id, e))
            self.trace_logger.log(e)

        finally:
            utils.silent_close(s)
            s = None

        self.on_keys_processed()

    def on_keys_processed(self):
        """
        Event called on keys have been processed.
        Currently used to re-fill the link queue in the update key scenario.
        :return:
        """
        if not self.update_keys:
            return

        with self.user_refill_lock:
            qsize = self.link_queue.qsize()
            if qsize < 30:
                self.fill_user_key_links()

    def store_user(self, user, s, db_user=None, db_user_loaded=False):
        """
        Stores username to the database.
        :param user:
        :return:
        """
        type_id = 0
        if user.user_type == 'User':
            type_id = 1
        elif user.user_type == 'Organization':
            type_id = 2

        try:
            if not db_user_loaded:
                db_user = s.query(GitHubUserDb).filter(GitHubUserDb.id == user.user_id).one_or_none()
            if db_user is not None:
                db_user.date_last_check = salch.func.now()
                db_user.usr_type = type_id
                s.merge(db_user)
                return 0

        except Exception as e:
            self.trace_logger.log(e)
            logger.warning('User query problem: %s' % e)

        # Store a new user here
        try:
            db_user = GitHubUserDb()
            db_user.id = user.user_id
            db_user.username = user.user_name
            db_user.usr_type = type_id
            s.add(db_user)
            return 0

        except Exception as e:
            self.trace_logger.log(e)
            logger.warning('[%02d] Exception during user store: %s' % (self.local_data.idx, e))
            if db_user_loaded:
                raise
            return 1

    def load_existing_key(self, key, s):
        """
        Loads existing key if exists
        :param key:
        :param s:
        :return:
        """
        key_id = int(key['id'])
        return s.query(GitHubKey).filter(GitHubKey.id == key_id).one_or_none()

    def store_key(self, user, key, s):
        """
        Stores user key to the database.
        :param user:
        :param key:
        :param s: current DB session
        :return:
        """

        # Loading phase
        existing_key = None
        try:
            if self.merge or self.update_keys:
                existing_key = self.load_existing_key(key, s)

        except Exception as e:
            logger.warning('Exception: %s' % e)

        # Storing phase
        try:
            if existing_key is not None:
                existing_key.date_last_check = salch.func.now()
                s.merge(existing_key)
                return existing_key

            key_id = int(key['id'])
            key_raw = key['key']

            key_type, key_val = [utils.strip(x) for x in key_raw.split(' ', 1)]

            db_key = GitHubKey()
            db_key.id = key_id
            db_key.key_id = key_id
            db_key.key_type = key_type
            db_key.key_user_found = user.user_name
            db_key.key_user_id_found = user.user_id
            db_key.text_raw = key_raw

            if key_type == 'ssh-rsa':
                try:
                    key_obj = utils.load_ssh_pubkey(key_raw)
                    if isinstance(key_obj, RSAPublicKey):
                        db_key.key_size = key_obj.key_size
                        numbers = key_obj.public_numbers()
                        db_key.key_modulus_hex = '%x' % numbers.n
                        db_key.key_exponent = numbers.e
                except Exception as e:
                    logger.info('Exception during processing the key[%s]: %s' % (key_type, e))

            s.add(db_key)
            return db_key

        except Exception as e:
            utils.silent_rollback(s)
            logger.warning('Exception during key store: %s' % e)
            return 1

    def assoc_key(self, user_id, key_id, s):
        """
        Association user <-> key
        :param user_id:
        :param key_id:
        :param s:
        :return:
        """
        try:
            uassoc = None
            if self.merge:
                uassoc = s.query(GitHubUserKeys)\
                    .filter(GitHubUserKeys.user_id == user_id)\
                    .filter(GitHubUserKeys.key_id == key_id)\
                    .one_or_none()

            if uassoc is not None:
                uassoc.lost_at = None
                return

            if uassoc is None:
                uassoc = GitHubUserKeys()
                uassoc.user_id = user_id
                uassoc.key_id = key_id
                uassoc.fount_at = salch.func.now()
                uassoc.lost_at = None
                s.add(uassoc)

            return 0

        except Exception as e:
            utils.silent_rollback(s)
            logger.warning('Exception during key assoc: %s' % e)
            return 1

    def deassoc_lost_keys(self, user_id, key_id, s):
        """
        Mark all lost keys as lost for the given user (not in key_id)
        :param user_id:
        :param key_id:
        :param s:
        :return:
        """
        try:
            q = s.query(GitHubUserKeys)\
                .filter(GitHubUserKeys.user_id == user_id)

            if len(key_id) > 0:
                q = q.filter(GitHubUserKeys.key_id.notin_(key_id))

            q.update({GitHubUserKeys.lost_at: salch.func.now()}, synchronize_session='fetch')
            return 0

        except Exception as e:
            utils.silent_rollback(s)
            logger.warning('Exception during key deassoc: %s' % e)
            return 1

    def flush_state(self):
        """
        Flushes state/config to the state file
        :return:
        """
        self.state['since_id'] = self.since_id
        self.state['rate_limit_remaining'] = self.rate_limit_remaining
        self.state['rate_limit_reset'] = self.rate_limit_reset
        utils.flush_json(self.state, self.state_file_path)

    #
    # DB
    #

    def fill_user_key_links(self):
        """
        Loads next X users from the database, advances since_id
        :return:
        """
        # self.since_id
        s = self.session()
        try:
            db_users = s.query(GitHubUserDb)\
                .filter(GitHubUserDb.id > self.since_id)\
                .order_by(GitHubUserDb.id)\
                .limit(self.user_load_bulk)\
                .all()

            for user in db_users:
                key_url = self.KEYS_ID_URL % user.id
                github_user = GitHubUser(user_id=user.id, user_name=user.username,
                                         user_type=user.usr_type, user_url=self.USER_URL % user.username)
                new_job = DownloadJob(url=key_url, jtype=DownloadJob.TYPE_KEYS, user=github_user,
                                      priority=random.randint(0, 1000), time_added=time.time())
                self.link_queue.put(new_job)

                if user.id > self.since_id:
                    self.since_id = user.id

        except Exception as e:
            logger.warning('Exception in loading users: %s' % e)
            utils.silent_rollback(s)

        finally:
            utils.silent_close(s)

    #
    # Auditing - errors, problems for further analysis
    #

    def audit_log(self, evt=None, link=None, jtype=None, job=None):
        """
        Appends audit log to the buffer. Lock protected.
        :param evt:
        :param link:
        :return:
        """
        log = collections.OrderedDict()
        log['time'] = time.time()
        log['evt'] = evt
        log['jtype'] = jtype
        log['link'] = link

        if job is not None and isinstance(job, DownloadJob):
            log['job'] = job.to_json()

        with self.audit_lock:
            self.audit_records_buffered.append(log)

    def flush_audit(self):
        """
        Flushes audit logs to the JSON append only file.
        Routine protected by the lock (no new audit record can be inserted while holding the lock)
        :return:
        """
        if self.audit_file is None:
            self.audit_records_buffered = []
            return

        self.audit_lock.acquire()
        try:
            if len(self.audit_records_buffered) == 0:
                return
            with open(self.audit_file, 'a') as fa:
                for x in self.audit_records_buffered:
                    fa.write(json.dumps(x) + "\n")
            self.audit_records_buffered = []
        except Exception as e:
            logger.error('Exception in audit log dump %s' % e)
        finally:
            self.audit_lock.release()

    #
    # State save / resume
    #

    def state_main(self):
        """
        State thread - periodical dump of the queues.
        :return:
        """
        logger.info('State thread started %s %s %s' % (os.getpid(), os.getppid(), threading.current_thread()))
        try:
            while not self.stop_event.is_set() and not self.terminate:
                # Dump stats each x seconds
                # Sleep is here because of dumping the state for the last time just before program quits.
                self.interruptible_sleep_delta(2)
                self.state_save()

                # Check memory conditions
                self.state_ram_check()

        except Exception as e:
            self.trace_logger.log(e)
            logger.error('Exception in state: %s' % e)

        finally:
            pass

        logger.info('State loop terminated')

    def state_ram_check(self):
        """
        Checks memory terminating conditions
        :return:
        """

        if self.max_mem is None:
            return

        cur_ram = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        if cur_ram <= self.max_mem:
            return

        logger.warning('Maximum memory threshold reached: %s kB = %s MB, threshold = %s kB'
                       % (cur_ram, cur_ram / 1024.0, self.max_mem))
        self.trigger_stop()

    def state_gen(self):
        """
        Dumps state
        :return:
        """
        try:
            js_q = collections.OrderedDict()
            js_q['gen'] = time.time()
            js_q['link_size'] = self.link_queue.qsize()
            js_q['since_id'] = self.since_id
            js_q['memory'] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss

            # Dequeues
            self.new_users_events.maintain()
            self.new_keys_events.maintain()

            users_in_5min = self.new_users_events.under_limit(5*60)
            keys_in_5min = self.new_keys_events.under_limit(5*60)

            js_q['users_dequeue_size'] = self.new_users_events.len()
            js_q['keys_dequeue_size'] = self.new_keys_events.len()
            js_q['users_5min'] = users_in_5min
            js_q['keys_5min'] = keys_in_5min
            js_q['users_1min'] = users_in_5min / 5.0
            js_q['keys_1min'] = keys_in_5min / 5.0

            # link queue structure
            qdata = list(self.link_queue.queue)
            qdata.sort(cmp=DownloadJob.cmp)
            js_q['link_structure'] = ''.join(['.' if x.type == DownloadJob.TYPE_KEYS else 'U' for x in qdata])

            # Stats.
            js_q['resource_stats'] = [x.to_json() for x in list(self.resources_list)]

            # Finally - the queue
            js_q['link_queue'] = [x.to_json() for x in qdata]
            return js_q

        except Exception as e:
            self.trace_logger.log(e)
            logger.error('Exception in state: %s', e)

    def state_save(self):
        """
        saves the state
        :return:
        """
        try:
            js_q = self.state_gen()
            utils.flush_json(js_q, self.state_file_path)

        except Exception as e:
            self.trace_logger.log(e)
            logger.error('Exception in state: %s', e)

    def state_resume(self):
        """
        Attempts to resume the queues from the monitoring files
        :return:
        """
        try:
            if self.state is None:
                return

            if 'since_id' in self.state:
                self.since_id = self.state['since_id']

            if 'link_queue' in self.state:
                for rec in self.state['link_queue']:
                    job = DownloadJob.from_json(rec)
                    self.link_queue.put(job)
                logger.info('Link queue resumed, entries: %d' % len(self.state['link_queue']))

        except Exception as e:
            self.trace_logger.log(e)
            logger.warning('Exception in resuming the state %s' % e)
            logger.error('State resume failed, exiting')
            sys.exit(1)

def main():
    args_src = sys.argv
    parser = argparse.ArgumentParser(description='Downloads GitHub SSH keys')
    parser.add_argument('-c', dest='config', default=None, help='JSON config file')
    parser.add_argument('-s', dest='status', default=None, help='JSON status file')
    parser.add_argument('-t', dest='threads', default=1, help='Number of download threads to use')
    parser.add_argument('--max-mem', dest='max_mem', default=None, type=int,
                        help='Maximal memory threshold in kB when program terminates itself')
    parser.add_argument('--users-only', dest='users_only', default=False, action='store_const', const=True,
                        help='Load only users list')
    parser.add_argument('--merge', dest='merge', default=False, action='store_const', const=True,
                        help='Merge DB operation - merge instead of add. slower, updates if exists')
    parser.add_argument('--update-keys', dest='update_keys', default=False, action='store_const', const=True,
                        help='Update keys for existing users')

    args = parser.parse_args(args=args_src[1:])
    config_file = args.config

    audit_file = os.path.join(os.getcwd(), 'audit.json')
    state_file = args.status if args.status is not None else os.path.join(os.getcwd(), 'state.json')
    if os.path.exists(state_file):
        utils.file_backup(state_file, backup_dir='.')

    if os.path.exists('.github-quit'):
        os.remove('.github-quit')

    sys.argv = [args_src[0]]
    logger.info('GitHub loader started, args: %s' % args)
    l = GitHubLoader(state_file=state_file, config_file=config_file, audit_file=audit_file, threads=args.threads,
                     max_mem=args.max_mem, users_only=args.users_only, merge=args.merge, update_keys=args.update_keys)
    l.work()
    sys.argv = args_src

# Launcher
if __name__ == "__main__":
    main()

{'create_engine': 'sqlalchemy.create_engine', 'sessionmaker': 'sqlalchemy.orm.sessionmaker', 'scoped_session': 'sqlalchemy.orm.scoped_session'}


import random
import torch
from torch import nn
import torch.nn.functional as F
from torch.autograd import Variable
from models.encoders import GRUEncoder, TransformerEncoder
from models.decoders import GRUDecoder, TransformerDecoder

class EncoderDecoder(nn.Module):
    """
    A Base Seq2Seq Model w/ an Encoder-Decoder Architecture
    Returns log probabilities of shape [batch_size, seq_len, tgt_vocab_size]
    """

    def __init__(self, encoder, decoder, generator):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.generator = generator

    def forward(self, src, tgt, src_mask, tgt_mask, src_lengths=None, trg_lengths=None):

        # pass the src sequence through the Encoder
        encoder_outputs, encoder_final = self.encode(
            src, src_mask, src_lengths)

        if isinstance(self.decoder, GRUDecoder):
            # last hidden state of the Encoder is used as the initial hidden state of the Decoder
            # [num_layers, batch, 2 * hidden_size]]
            encoder_final = encoder_final[:self.decoder.num_layers]
        else:
            encoder_final = None

        # pass the tgt sequence through the Decoder
        decoder_output, _ = self.decode(trg=tgt, encoder_outputs=encoder_outputs,
                                        src_mask=src_mask, trg_mask=tgt_mask,
                                        encoder_final=encoder_final, decoder_hidden=None)
        logits = self.generator(decoder_output)
        return F.log_softmax(logits, dim=-1)

    def encode(self, src, src_mask, src_lengths):
        """
        Encode the src sequence using the Encoder

        Arguments:
            src: Source sequence tensor [batch_size, seq_len]
            src_mask: Mask for src sequence [batch_size, 1, seq_len]
            src_lengths: lenght of each example in batch

        Returns:
            A Tensor of shape [batch_size, seq_len, hidden_size]
        """
        return self.encoder(src, src_mask, src_lengths)

    def decode(self, trg, encoder_outputs, src_mask, trg_mask, encoder_final=None, decoder_hidden=None):
        """
        Decode the target sequence given the outputs from the Encoder

        Arguments:
            tgt: Target sequence tensor [batch_size, seq_len]
            encoder_outputs: Output Tensor from the Encoder [batch_size, seq_len, hidden_size]
            encoder_final: Final hidden state from the Encoder (only for GRUEncoder)
            src_mask: Mask for src sequence [batch_size, 1, seq_len]
            tgt_mask: Mask for the tgt sequence [batch_size, seq_len, seq_len] (only for TransformerDecoder)
            decoder_hidden: decoder hidden state

        Returns:
            A Tensor of shape [batch_size, seq_len, hidden_size]
        """
        return self.decoder(trg, encoder_outputs, src_mask, trg_mask, encoder_final, hidden=decoder_hidden)

def make_seq2seq_model(params):
    """ Make a Seq2Seq moddel with given parameters """

    if params.model_type == "GRU":
        encoder = GRUEncoder(src_vocab_size=params.src_vocab_size, embed_size=params.embed_size,
                             hidden_size=params.hidden_size, input_dropout_p=params.input_dropout,
                             num_layers=params.n_layers_enc, dropout_p=params.layer_dropout)

        decoder = GRUDecoder(trg_vocab_size=params.tgt_vocab_size, embed_size=params.embed_size,
                             hidden_size=params.hidden_size, attention=params.attention,
                             input_dropout_p=params.input_dropout,
                             dropout_p=params.layer_dropout, device=params.device,
                             num_layers=params.n_layers_dec)
    else:
        encoder = TransformerEncoder(params.embed_size, params.src_vocab_size, params.hidden_size, params.n_layers_enc,
                                     params.num_heads, params.max_length, d_ff=params.d_ff, input_dropout=params.input_dropout,
                                     layer_dropout=params.layer_dropout, attention_dropout=params.attention_dropout,
                                     relu_dropout=params.relu_dropout)

        # create the Transformer Decoder
        decoder = TransformerDecoder(params.embed_size, params.tgt_vocab_size, params.hidden_size, params.n_layers_dec,
                                     params.num_heads, params.max_length, d_ff=params.d_ff, input_dropout=params.input_dropout,
                                     layer_dropout=params.layer_dropout, attention_dropout=params.attention_dropout,
                                     relu_dropout=params.relu_dropout)

    # standard linear + softmax generation step
    generator = nn.Linear(params.hidden_size,
                          params.tgt_vocab_size, bias=False)

    if params.tgt_emb_prj_weight_sharing:
        # share the weight matrix between the target embedding and final softmax layer
        if params.model_type == "Transformer":
            generator.weight = decoder.embeddings.embedding.weight
        else:
            generator.weight = decoder.embed.weight

    if params.emb_src_tgt_weight_sharing:
        # Share the weight matrix between source & target word embeddings
        # this can only happen if you have a shared vocab
        assert params.src_vocab_size == params.tgt_vocab_size, "To share word embedding between src and tgt, the vocab sizes must be the same"
        if params.model_type == "Transformer":
            encoder.embeddings.embedding.weight = decoder.embeddings.embedding.weight
        else:
            encoder.embed.weight = decoder.embed.weight

    # define the EncoderDecoer Model
    model = EncoderDecoder(encoder, decoder, generator).to(params.device)

    # weight initialization of the parameters of the model: Xavier Initialization
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)

    return model


# -*- coding: utf-8 -*-
"""
pytorch_neural_network_binary_classification.ipynb
## Neural Network Classification with PyTorch
## 1. Make Custom Classification Dataset and Get It Ready
"""

from sklearn.datasets import make_circles

# Number of samples, let's get 1000 samples
n_samples = 1000

# Create Circles
X, y = make_circles(n_samples=n_samples,
                    noise=0.03,
                    random_state=42)

# Make DataFrame of Circle Data
import pandas as pd

dataFrame = pd.DataFrame({'X1': X[:, 0],
                          'X2': X[:, 1],
                          'Label': y})

dataFrame.head(10)

from matplotlib import pyplot as plt
plt.scatter(x=X[:, 0],
            y=X[:, 1],
            c=y)

"""
## 1.1 Turn Data into Tensors and Create Train & Test Splits
"""

# Turn data into tensors
import torch
X = torch.from_numpy(X).type(torch.float)
y = torch.from_numpy(y).type(torch.float)

# Split data into training and tests sets
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

"""
## 2. Building a Model

To do so, we want to:

1. Setup device agonistic code so our code will run on an accelerator (GPU) if there is one
2. Construct a model (by subclassing nn.Module)
3. Define a loss function and optimizer
4. Create a training and test loop
"""

import torch
from torch import nn

device = 'cuda' if torch.cuda.is_available() else 'cpu'

"""
Now we've setup device agnostic code, let's create a model that:

1. Subclasses nn.Module (almost all models in PyTorch subclass nn.Module)
2. Create 2 nn.Linear() layers that are capable of handling the shapes of our data
3. Defines a forward() method that outlines the forward pass (or forward computation) of the model
4. Instatiate an instance of our model class and send it to the target device
"""

# 1. Construct a model that subclasses nn.Module
class CircleModel(nn.Module):
  def __init__(self):
    super().__init__()
    # 2. Create 2 nn.Linear layers capable of handling the shapes of our data
    self.input_layer = nn.Linear(in_features=2,
                                 out_features=64,
                                 bias=True) # takes in 2 features and upscales to 64 features

    self.first_hidden_layer = nn.Linear(in_features=64,
                                  out_features=1,
                                  bias=True) # takes in 64 features from previous layer and outputs a single feature (same shape as y)

  # 3. Define a forward() method that outlines the forward pass
  def forward(self, x:torch.Tensor) -> torch.Tensor:
    return self.first_hidden_layer(self.input_layer(x)) # x -> input_layer ->  first_hidden_layer -> output

# 4. Instantiate an instance of our model class and send it to the target device
model = CircleModel().to(device)
model

# Let's replicate the model above using nn.Sequential()
sequential_model = nn.Sequential(
    nn.Linear(in_features=2, out_features=64, bias=True),
    nn.Linear(in_features=64, out_features=1, bias=True)
).to(device)

sequential_model

# Setup the Loss Function and Optimizer
loss_func = nn.BCEWithLogitsLoss() # BCEWithLogitsLoss = sigmoid activation function built-in
optimizer = torch.optim.SGD(params=model.parameters(),
                            lr=0.01)

# Evaluation Function
def accuracy_func(y_pred, y_true):
  correct = torch.eq(y_true, y_pred).sum().item()
  return ((correct/len(y_true))*100)

"""
## 3. Build a Train & Test Loop
"""

torch.cuda.manual_seed(42)
torch.manual_seed(42)
epochs = 100

for epoch in range(epochs):

  # Train Mode
  model.train()

  # Send the Data Target Device
  X_train, X_test = X_train.to(device), X_test.to(device)
  y_train, y_test = y_train.to(device), y_test.to(device)

  # 1. Forward Pass
  train_logits = model.forward(X_train).squeeze()
  train_pred = torch.round(torch.sigmoid(train_logits)) # turn logits -> pred probs -> pred labels

  # 2. Calculate Train Loss/Accuracy
  train_loss = loss_func(train_logits, y_train)
  train_accuracy = accuracy_func(train_pred, y_train)

  # 3. Zero Grad with Optimizer
  optimizer.zero_grad()

  # 4. Backpropagation with Loss
  train_loss.backward()

  # 5. Optimizer step (gradient descent)
  optimizer.step()

  # Test Mode
  model.eval()
  with torch.inference_mode():
    # 1. Farward Pass
    test_logits = model.forward(X_test).squeeze()
    test_pred = torch.round(torch.sigmoid(test_logits))

    # 2. Calculate Test Loss/Accuracy
    test_loss = loss_func(test_logits, y_test)
    test_accuracy = accuracy_func(test_pred, y_test)

  if epoch % 5 == 0:
    print(f'Epoch: {epoch} | Training Accuracy {train_accuracy:.2f} | Training Loss: {train_loss:.2f} | Testing Accuracy {test_accuracy:.2f} | Testing Loss: {test_loss:.2f}')

"""
## From the metrics it looks like our model isn't learning anything...
"""

import requests
from pathlib import Path

# Download helper functions
if Path("helper_functions.py").is_file():
  print("helper_functions.py already exists, skipping download")
else:
  print("Downloading helper_functions.py")
  request = requests.get("https://raw.githubusercontent.com/UygarKAYA/DeepLearning/main/utils/helper_functions.py")
  with open("helper_functions.py", "wb") as f:
    f.write(request.content)

from utils.helper_functions import plot_decision_boundary

# Plot decision boundary of the model
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.title("Train")
plot_decision_boundary(model, X_train, y_train)
plt.subplot(1, 2, 2)
plt.title("Test")
plot_decision_boundary(model, X_test, y_test)

"""
## 4. Improving a model (from a model perspective)

* Add more layers - give the model more chances to learn about patterns in the * data
* Add more hidden units - go from 5 hidden units to 10 hidden units
* Fit for longer
* Changing the activation functions
* Change the learning rate
* Change the loss function

"""

class CircleModelV1(nn.Module):
  def __init__(self):
    super().__init__()
    self.input_layer = nn.Linear(in_features=2, out_features=10)
    self.first_hidden_layer = nn.Linear(in_features=10, out_features=10)
    self.second_hidden_layer = nn.Linear(in_features=10, out_features=1)

  def forward(self, x:torch.Tensor) -> torch.Tensor:
    return self.second_hidden_layer(self.first_hidden_layer(self.input_layer(x)))

improvedModel = CircleModelV1().to(device)
improvedModel

# Create a Loss Function
loss_func = nn.BCEWithLogitsLoss()
# Create a Optimizer
optimizer = torch.optim.SGD(params=improvedModel.parameters(),
                            lr=0.001)

torch.manual_seed(42)
torch.cuda.manual_seed(42)

epochs = 250
for epoch in range(epochs):
  # Training Step
  improvedModel.train()

  train_logits = improvedModel.forward(X_train).squeeze()
  train_prediction = torch.round(torch.sigmoid(train_logits))

  train_loss = loss_func(train_logits, y_train)
  train_accuracy = accuracy_func(train_prediction, y_train)

  optimizer.zero_grad()

  train_loss.backward()

  optimizer.step()

  # Testing Step
  improvedModel.eval()
  with torch.inference_mode():
    test_logits = improvedModel.forward(X_test).squeeze()
    test_predictions = torch.round(torch.sigmoid(test_logits))

    test_loss = loss_func(test_logits, y_test)
    test_accuracy = accuracy_func(test_predictions, y_test)

  if epoch % 25 == 0:
    print(f'Epoch: {epoch} | Training Accuracy {train_accuracy:.2f} | Training Loss: {train_loss:.2f} | Testing Accuracy {test_accuracy:.2f} | Testing Loss: {test_loss:.2f}')

from utils.helper_functions import plot_decision_boundary

# Plot decision boundary of the model
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.title("Train")
plot_decision_boundary(improvedModel, X_train, y_train)
plt.subplot(1, 2, 2)
plt.title("Test")
plot_decision_boundary(improvedModel, X_test, y_test)

"""
## 4.1 The missing piece: non-linearity
### 4.1.1 Building a model with non-linearity
"""

# Build a model with non-linear activation functions

class NonLinearModel(nn.Module):
  def __init__(self):
    super().__init__()
    self.input_layer = nn.Linear(in_features=2, out_features=10)
    self.first_hidden_layer = nn.Linear(in_features=10, out_features=10)
    self.second_hidden_layer = nn.Linear(in_features=10, out_features=1)
    self.relu = nn.ReLU()

  def forward(self, x:torch.Tensor) -> torch.Tensor:
    return self.second_hidden_layer(self.relu(self.first_hidden_layer(self.relu(self.input_layer(x)))))

improvedModelV1 = NonLinearModel().to(device)
improvedModelV1

loss_func = nn.BCEWithLogitsLoss()
optimizer = torch.optim.SGD(params=improvedModelV1.parameters(),
                            lr=0.01)

torch.manual_seed(42)
torch.cuda.manual_seed(42)

epochs = 250
for epoch in range(epochs):
  # Training Step
  improvedModelV1.train()

  train_logits = improvedModelV1.forward(X_train).squeeze()
  train_prediction = torch.round(torch.sigmoid(train_logits))

  train_loss = loss_func(train_logits, y_train)
  train_accuracy = accuracy_func(train_prediction, y_train)

  optimizer.zero_grad()

  train_loss.backward()

  optimizer.step()

  # Testing Step
  improvedModelV1.eval()
  with torch.inference_mode():
    test_logits = improvedModelV1.forward(X_test).squeeze()
    test_predictions = torch.round(torch.sigmoid(test_logits))

    test_loss = loss_func(test_logits, y_test)
    test_accuracy = accuracy_func(test_predictions, y_test)

  if epoch % 25 == 0:
    print(f'Epoch: {epoch} | Training Accuracy {train_accuracy:.2f} | Training Loss: {train_loss:.2f} | Testing Accuracy {test_accuracy:.2f} | Testing Loss: {test_loss:.2f}')

from utils.helper_functions import plot_decision_boundary

# Plot decision boundary of the model
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.title("Train")
plot_decision_boundary(improvedModelV1, X_train, y_train)
plt.subplot(1, 2, 2)
plt.title("Test")
plot_decision_boundary(improvedModelV1, X_test, y_test)


################
import sys
import os
import logging
from configparser import ConfigParser
from lasair_zooniverse import lasair_zooniverse_class
from lasair_zooniverse import lasair_object
import time
import numpy as np
################################

#base_dir = os.getenv('LASAIR_CONFIG_PATH')
base_dir = '/Users/dwright/dev/Zooniverse_SLSN'
config = ConfigParser()
config.read('%s/config.ini' % (base_dir))
log = logging.getLogger("main-logger")

while(1):
    try:
        lasair_zoo = lasair_zooniverse_class(config.get('APP', 'KAFKA_SERVER'), config.get('APP', 'ENDPOINT'))
        objectIds = lasair_zoo.query_lasair_topic(config.get('APP','GROUP_ID'), config.get('APP','TOPIC'))

        if (config.get('APP','RECORDS_LIMIT')) != 'None':
            max_limit = config.getint('APP','RECORDS_LIMIT')
            objectIds = objectIds[:max_limit]
        proto_subjects = []
        for object_id in objectIds:
            lasair_zoo.wget_objectdata(object_id, config.get('APP','URL'), config.get('APP','DATA_DIR'))
            proto_subject = lasair_zoo.produce_proto_subject(object_id, config.get('APP','DATA_DIR'))
            if (proto_subject != None):
                proto_subjects.append(proto_subject)

        #lasair_zoo.create_subjects_and_link_to_project(proto_subjects,config.get('APP','PROJECT_ID'), config.get('APP','WORKFLOW_ID'), None)
        time.sleep(config.getint('APP','SLEEP_TIME')) #sleep for one day
    except Exception:
        log.exception("Error in main block")


from portal.plugins.kinesisresponder.kinesis_responder import KinesisResponder
import json
import urllib
from django.conf import settings
from s3_mixin import S3Mixin, FileDoesNotExist
from vs_mixin import VSMixin
import logging
from gnmvidispine.vs_item import VSItem, VSNotFound
from datetime import datetime
import portal.plugins.gnmatomresponder.constants as const
import re

logger = logging.getLogger(__name__)

#Still need: holding image. this is more likely to come from the launch detection side than the atom side.

extract_extension = re.compile(r'^(?P<basename>.*)\.(?P<extension>[^\.]+)$')
multiple_underscore_re = re.compile(r'_{2,}')
make_filename_re = re.compile(r'[^\w\d\.]')

class MasterImportResponder(KinesisResponder, S3Mixin, VSMixin):
    def get_project_collection(self, content):
        """
        Gets a populated VSCollection reference to the project collection mentioned in content
        :param content: parsed json message including the 'projectId' key
        :return:
        """
        from .exceptions import NotAProjectError

        try:
            project_collection = self.get_collection_for_id(content['projectId'])
        except VSNotFound:
            try:
                logger.warning(u"Invalid parent collection {0} specified for atom {1} ({2}). Falling back to default.".format(
                    content['projectId'], content['atomId'],
                    content.get('title','(unknown title)').encode("UTF-8","backslashescape")
                ))
            except UnicodeDecodeError:
                pass
            except UnicodeEncodeError:
                pass
            project_collection = None
        except NotAProjectError:
            try:
                logger.warning(u"Collection {0} specified for atom {1} ({2}) is not a Project. Falling back to default.".format(
                    content['projectId'], content['atomId'],
                    content.get('title','(unknown title)').encode("UTF-8","backslashescape")
                ))
            except UnicodeDecodeError:
                pass
            except UnicodeEncodeError:
                pass
            project_collection = None

        if project_collection is None:
            collection_id = getattr(settings,'ATOM_RESPONDER_DEFAULT_PROJECTID',None)
            if collection_id is None:
                raise RuntimeError("Unable to get a project ID for atom {0}, and no default is set".format(content['atomId']))
            else:
                project_collection = self.get_collection_for_id(collection_id)

        return project_collection

    def update_pluto_record(self, item_id, project_id):
        import traceback
        try:
            from portal.plugins.gnm_masters.signals import master_external_update

            master_external_update.send(sender=self.__class__, item_id=item_id, project_id=project_id)
        except ImportError as e:
            logger.error("Unable to signal master update: {0}".format(e))
        except Exception as e:
            logger.error("An error happened when outputting master_external_create signal: {0}".format(traceback.format_exc()))

    def get_or_create_master_item(self, atomId, title, project_collection, user):
        master_item = self.get_item_for_atomid(atomId)

        created = False

        if master_item is None:
            if title is None:
                raise RuntimeError("Title field not set for atom {0}.".format(atomId))
            if user is None:
                logger.warning("User field not set for atom {0}.".format(atomId))
                user_to_set="unknown_user"
            else:
                user_to_set=user
            master_item = self.create_placeholder_for_atomid(atomId,
                                                             title=title,
                                                             user=user_to_set,
                                                             parent=project_collection
                                                             )
            logger.info("Created item {0} for atom {1}".format(master_item.name, atomId))
            created = True
        return master_item, created

    def process(self,record, approx_arrival, attempt=0):
        """
        Process a message from the kinesis stream.  Each record is a JSON document which contains keys for atomId, s3Key,
        projectId.  This will find an item with the given atom ID or create a new one, get a signed download URL from
        S3 for the media and then instruct Vidsipine to import it.
        Rather than wait for the job to complete here, we return immediately and rely on receiving a message from VS
        when the job terminates.
        :param record: JSON document in the form of a string
        :param approx_arrival:
        :param attempt: optional integer showing how many times this has been retried
        :return:
        """
        from media_atom import request_atom_resend, HttpError
        from tasks import timed_request_resend
        content = json.loads(record)

        logger.info(content)

        #We get two types of message on the stream, one for incoming xml the other for incoming media.
        if content['type'] == const.MESSAGE_TYPE_MEDIA or content['type'] == const.MESSAGE_TYPE_RESYNC_MEDIA:
            project_collection = self.get_project_collection(content)
            if 'user' in content:
                atom_user = content['user']
            else:
                atom_user = None
            master_item, created = self.get_or_create_master_item(content['atomId'], content['title'], project_collection, atom_user)

            return self.import_new_item(master_item, content, parent=project_collection)
        elif content['type'] == const.MESSAGE_TYPE_PAC:
            logger.info("Got PAC form data message")
            record = self.register_pac_xml(content)
            self.ingest_pac_xml(record)
            logger.info("PAC form data message complete")
        elif content['type'] == const.MESSAGE_TYPE_PROJECT_ASSIGNED:
            logger.info("Got project (re-)assignment message: {0}".format(content))

            master_item = self.get_item_for_atomid(content['atomId'])
            if master_item is not None:
                logger.info("Master item for atom already exists at {0}, assigning".format(master_item.name))
                self.assign_atom_to_project(content['atomId'], content['commissionId'], content['projectId'], master_item)
            else:
                logger.warning("No master item exists for atom {0}.  Requesting a re-send from media atom tool".format(content['atomId']))
                try:
                    request_atom_resend(content['atomId'], settings.ATOM_TOOL_HOST, settings.ATOM_TOOL_SECRET)
                except HttpError as e:
                    if e.code == 404:
                        if attempt >= 10:
                            logger.error("{0}: still nothing after 10 attempts. Giving up.".format(content['atomId']))
                            raise
                        logger.warning("{0}: Media atom tool responded with a 404 on attempt {1}: {2}. Retrying in 60s.".format(content['atomId'], attempt, e.content))
                        timed_request_resend.apply_async(args=(record, approx_arrival),
                                                         kwargs={'attempt': attempt+1},
                                                         countdown=60)
                    else:
                        logger.exception("{0}: Could not request resync".format(content['atomId']))

            logger.info("Project (re-)assignment complete")
        else:
            raise ValueError("Unrecognised message type: {0}".format(content['type']))

    def register_pac_xml(self, content):
        """
        Start the import of new PAC data by registering it in the database.
        :param content: JSON message content as received from atom tool
        :return: the database model instance
        """
        from models import PacFormXml

        (record, created) = PacFormXml.objects.get_or_create(atom_id=content['atomId'], defaults={'received': datetime.now()})
        if not created:
            logger.info("PAC form xml had already been delivered for {0}, over-writing".format(content['atomId']))

        record.completed = None
        record.status = "UNPROCESSED"
        record.last_error = ""
        record.pacdata_url = "s3://{bucket}/{path}".format(bucket=content['s3Bucket'], path=content['s3Path'])
        record.save()
        return record

    def check_for_old_finished_jobs(self, vs_item_id):
        from models import ImportJob

        jobs = ImportJob.objects.filter(item_id=vs_item_id).filter(status='FINISHED').count()

        return jobs > 0

    def check_key(self, key, vs_item_id):
        from models import ImportJob

        jobs = ImportJob.objects.filter(item_id=vs_item_id).filter(s3_path=key).count()

        return jobs > 0

    def check_for_processing(self, vs_item_id):
        from models import ImportJob

        jobs = ImportJob.objects.filter(item_id=vs_item_id).filter(processing=True).count()

        return jobs > 0

    def import_new_item(self, master_item, content, parent=None):
        from models import ImportJob, PacFormXml
        from pac_xml import PacXmlProcessor
        from mock import MagicMock
        if not isinstance(master_item, VSItem) and not isinstance(master_item, MagicMock): raise TypeError #for intellij
        from portal.plugins.kinesisresponder.sentry import inform_sentry

        vs_item_id = master_item.get("itemId")

        if vs_item_id is None:
            vs_item_id = master_item.name

        old_finished_jobs = self.check_for_old_finished_jobs(vs_item_id)

        old_key = self.check_key(content['s3Key'], vs_item_id)

        if old_finished_jobs is True and old_key is True:
            logger.info('A job for item {0} has already been successfully completed. Aborting.'.format(vs_item_id))
            inform_sentry('A job for item {0} has already been successfully completed. Aborting.'.format(vs_item_id), {
                "master_item": master_item,
                "content": content,
                "parent": parent
            })
            return

        processing_job = self.check_for_processing(vs_item_id)

        if processing_job is True:
            logger.info('Job for item {0} already in progress. Aborting.'.format(vs_item_id))
            inform_sentry('Job for item {0} already in progress. Aborting.'.format(vs_item_id), {
                "master_item": master_item,
                "content": content,
                "parent": parent
            })
            return

        safe_title = content.get('title','(unknown title)').encode("UTF-8","backslashescape").decode("UTF-8")

        #using a signed URL is preferred, but right now VS seems to have trouble ingesting it.
        #so, we download instead and ingest that. get_s3_signed_url is left in to make it simple to switch back
        #download_url = self.get_s3_signed_url(bucket=settings.ATOM_RESPONDER_DOWNLOAD_BUCKET, key=content['s3Key'])
        downloaded_path = self.download_to_local_location(bucket=settings.ATOM_RESPONDER_DOWNLOAD_BUCKET,
                                                          key=content['s3Key'],
                                                          #this is converted to a safe filename within download_to_local_location
                                                          filename=content.get('title', None)) #filename=None => use s3key instead

        download_url = "file://" + urllib.quote(downloaded_path)

        logger.info(u"{n}: Ingesting atom with title '{0}' from media atom with ID {1}".format(safe_title, content['atomId'], n=master_item.name))
        try:
            logger.info(u"{n}: Download URL is {0}".format(download_url, n=master_item.name))
        except UnicodeEncodeError:
            pass
        except UnicodeDecodeError:
            pass

        job_result = master_item.import_to_shape(uri=download_url,
                                                 essence=True,
                                                 shape_tag=getattr(settings,"ATOM_RESPONDER_SHAPE_TAG","lowres"),
                                                 priority=getattr(settings,"ATOM_RESPONDER_IMPORT_PRIORITY","HIGH"),
                                                 jobMetadata={'gnm_source': 'media_atom'},
                                                 )
        logger.info(u"{0} Import job is at ID {1}".format(vs_item_id, job_result.name))

        master_item.set_metadata({const.GNM_ASSET_FILENAME: downloaded_path})

        if parent is not None:
            logger.info(u"{0}: Adding to collection {1}".format(vs_item_id, parent.name))
            master_item.name = vs_item_id
            parent.addToCollection(master_item)
            logger.info(u"{0}: Done".format(vs_item_id))
        else:
            logger.error(u"{0}: No parent collection specified for item!".format(vs_item_id))

        self.update_pluto_record(vs_item_id, parent.name)
        #make a note of the record. This is to link it up with Vidispine's response message.
        record = ImportJob(item_id=vs_item_id,
                           job_id=job_result.name,
                           status='STARTED',
                           started_at=datetime.now(),
                           user_email=content.get('user',"Unknown user"),
                           atom_id=content['atomId'],
                           atom_title=content.get('title', "Unknown title"),
                           s3_path=content['s3Key'],
                           processing=True)
        previous_attempt = record.previous_attempt()
        if previous_attempt:
            record.retry_number = previous_attempt.retry_number+1
            logger.info(u"{0} Import job is retry number {1}".format(vs_item_id, record.retry_number))
        record.save()

        try:
            logger.info(u"{n}: Looking for PAC info that has been already registered".format(n=vs_item_id))
            pac_entry = PacFormXml.objects.get(atom_id=content['atomId'])
            logger.info(u"{n}: Found PAC form information at {0}".format(pac_entry.pacdata_url,n=vs_item_id))
            proc = PacXmlProcessor(self.role_name, self.session_name)
            proc.link_to_item(pac_entry, master_item)
        except PacFormXml.DoesNotExist:
            logger.info(u"{n}: No PAC form information has yet arrived".format(n=vs_item_id))

    def ingest_pac_xml(self, pac_xml_record):
        """
        Master process to perform import of pac data
        :param pac_xml_record: instance of PacFormXml model
        :return:
        """
        from pac_xml import PacXmlProcessor

        vsitem = self.get_item_for_atomid(pac_xml_record.atom_id)
        if vsitem is None:
            logger.warning("No item could be found for atom ID {0}, waiting for it to arrive".format(pac_xml_record.atom_id))
            return

        proc = PacXmlProcessor(self.role_name,self.session_name)

        #this process will call out to Pluto to do the linkup once the data has been received
        return proc.link_to_item(pac_xml_record, vsitem)

    def assign_atom_to_project(self, atomId, commissionId, projectId, master_item):
        """
        (re)-assigns the given master to a project.
        If the project is not associated with the given commission, warns but does not fail.
        :param atomId:
        :param commissionId:
        :param projectId:
        :return:
        """

        current_project_id = master_item.get(const.PARENT_COLLECTION)
        if current_project_id is not None:
            logger.warning("Re-assigning master {0} to project {1} so removing from {2}".format(master_item.name, projectId, current_project_id))
            current_project_ref = self.get_collection_for_id(current_project_id)
            current_project_ref.removeFromCollection(master_item)

        new_project_ref = self.get_collection_for_id(projectId)
        expected_commission_id = new_project_ref.get(const.PARENT_COLLECTION)
        if expected_commission_id!=commissionId:
            logger.warning("Project {0} belongs to commission {1}, but media atom thinks it belongs to commission {2}. Continuing anyway".format(projectId, expected_commission_id, commissionId))

        logger.info("Adding master {0} to collection {1}".format(master_item.name, new_project_ref.name))
        new_project_ref.addToCollection(master_item)
        logger.info("Setting up project fields for {0}".format(master_item.name))
        self.set_project_fields_for_master(master_item,parent_project=new_project_ref)
        logger.info("Telling gnm_masters about updates for {0}".format(master_item.name))
        self.update_pluto_record(master_item.name, projectId)


from typing import Any, Tuple, Optional, List

import pytorch_lightning as pl

from torch import Tensor
from torch.optim import SGD, Optimizer, Adam
from torch.optim.lr_scheduler import MultiStepLR

from .config.kdfinetune import Config
from src.data import get_partial_datamodule, BaseDataModule
from .utils import (
    get_loss_function,
    compute_accuracy,
    evaluate_benign_accuracy,
    evaluate_backdoor_success_rate
)
from .kdbackdoor import KDBackdoorModel

_config = Config()

class KDFinetuneModel(pl.LightningModule):
    name: str = "kdfinetune"

    def __init__(
        self, *,
        kd_checkpoint_path: str,
        train_partial_rate: float = _config.train_partial_rate,
        test_partial_rate: float = _config.test_partial_rate,
        loss_function: str = _config.loss_function,
        lr: float = _config.lr,
        momentum: float = _config.momentum,
        weight_decay: float = _config.weight_decay,
        epochs: int = _config.epochs,
        milestones: List[int] = _config.milestones,
        gamma: float = _config.gamma,
        **datamodule_kwargs
    ) -> None:
        super().__init__()

        self.save_hyperparameters()

        kd_model = KDBackdoorModel.load_from_checkpoint(kd_checkpoint_path)
        self._network = kd_model._teacher_network
        self._backdoor = kd_model._backdoor_network
        self._target_label = kd_model.hparams.target_label
        self._datamodule = get_partial_datamodule(
            kd_model.datamodule,
            train_partial_rate=train_partial_rate,
            test_partial_rate=test_partial_rate,
            **datamodule_kwargs
        )

        self._loss_function = get_loss_function(loss_function)

    def forward(self, x: Tensor) -> Any:
        return self._network(x)

    def training_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> Any:
        x, y = batch
        pred_y = self._network(x)
        loss = self._loss_function(pred_y, y)
        acc = compute_accuracy(pred_y, y)

        self.log_dict({
            "finetune_train_loss": loss,
            "finetune_train_acc": acc
        }, on_epoch=True)

        return loss

    def on_train_epoch_end(self, unused: Optional = None) -> None:
        test_acc = evaluate_benign_accuracy(
            model=self._network,
            dataloader=self.datamodule.test_dataloader(),
            device=self.device
        )
        self.log_dict({
            "finetune_test_acc": test_acc
        }, on_step=False, on_epoch=True)

        self.log_dict({
            "finetune_teacher_backdoor_success_rate": evaluate_backdoor_success_rate(
                model=self._network,
                backdoor=self._backdoor,
                dataloader=self.datamodule.test_dataloader(),
                target_label=self._target_label,
                device=self.device
            )
        })

    def configure_optimizers(self) -> Optimizer:
        optimizer = SGD(
            self._network.parameters(),
            lr=self.hparams.lr,
            momentum=self.hparams.momentum,
            weight_decay=self.hparams.weight_decay
        )
        scheduler = MultiStepLR(
            optimizer,
            milestones=self.hparams.milestones,
            gamma=self.hparams.gamma
        )

        return [optimizer], [scheduler]

    @property
    def datamodule(self) -> BaseDataModule:
        return self._datamodule


from tkinter import *
from PIL import ImageTk, Image
from tkinter import messagebox
import subprocess

# widget =GUI elements: button, textbox,labels, images
# windows= serves as a container to hold or contain these widgets

window = Tk()  # instantiate an instance of a window
window.geometry("1000x650")
window.title("ToodleDo - Sign In")

window.geometry("1000x650")
img = ImageTk.PhotoImage(Image.open("To do list Home webpage.png"))
label = Label(window, image=img)
label.pack()
label.place(x=0, y=0)

task = StringVar()
task_entry = Entry(window,
                   font=("Cambria", 15)
                   )
task_entry.pack(side=LEFT)
task_entry.place(x=612, y=185)
task_entry.focus()

task_list = []

def addTask():
    task = task_entry.get()
    task_entry.delete(0, END)

    if task:
        with open("tasklist.txt", 'a') as taskfile:
            taskfile.write(f"\n{task}")
        task_list.append(task)
        listbox.insert(END, task)

def deleteTask():
    global task_list
    task = str(listbox.get(ANCHOR))
    if task in task_list:
        task_list.remove(task)
        with open("tasklist.txt", 'w') as taskfile:
            for task in task_list:
                taskfile.write(task + "/n")

        listbox.delete(ANCHOR)

def openTaskFile():
    try:
        global task_list
        with open("task.png", "r") as taskfile:
            tasks = taskfile.readlines()

        for task in tasks:
            if task != '\n':
                task_list.append(task)
                listbox.insert(END, task)

    except:
        file = open('tasklist.txt', 'w')
        file.close()

# listbox
frame1 = Frame(window, bd=3, width=295, height=295, bg="#32405b")
frame1.pack(pady=(160, 0))
frame1.place(x=623, y=242)

listbox = Listbox(frame1, font=('arial', 12), width=28, height=15, bg="#32405b", fg="white", cursor="hand2",
                  selectbackground="#5a95ff")
listbox.pack(side=LEFT, fill=BOTH, padx=2)
scrollbar = Scrollbar(frame1)
scrollbar.pack(side=RIGHT, fill=BOTH)

listbox.config(yscrollcommand=scrollbar.set)
scrollbar.config(command=listbox.yview)

openTaskFile()

submit_button = Button(window,
                       text="ADD",
                       command=addTask,
                       font=("Jacques Francois", 10),
                       fg="#E7C6FF",
                       bg="#F72585",
                       activeforeground="#E7C6FF",
                       activebackground="#7209B7",
                       state=ACTIVE)
submit_button.pack(side=RIGHT)
submit_button.place(x=842, y=186)

Delete_icon = PhotoImage(file="delete.png")
button = Button(window, image=Delete_icon, command=deleteTask)
button.pack(side=BOTTOM, pady=13)
button.place_configure(x=735, y=540)

window.mainloop()


import cv2
import fitz
import numpy as np
import pytesseract
from PIL import Image

wfname = '../../test/test_resources/DOC090519-09052019083949.pdf'

PATH = '/home/cecilia/Documenti/images/'

# async def main_body():
#     with open(wfname, 'rb') as f:
#         file = File(name = wfname.split('/')[-1], body = f.read(), type = '')
#     res = extract_file(file)
#     async for r in res:
#         print(r['text'])
#         print()
#         print()
#
#
# if __name__ == '__main__':
#     loop = asyncio.get_event_loop()
#     loop.run_until_complete(asyncio.wait([main_body()]))

# with open(wfname, 'rb') as f:
#     file = f

# get grayscale image
def get_grayscale(image):
    return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

# noise removal
def remove_noise(image):
    return cv2.medianBlur(image, 3)

# thresholding
def thresholding(image):
    return cv2.threshold(image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1]

# dilation
def dilate(image):
    kernel = np.ones((5, 5), np.uint8)
    return cv2.dilate(image, kernel, iterations=1)

# erosion
def erode(image):
    kernel = np.ones((2, 2), np.uint8)
    return cv2.erode(image, kernel, iterations=1)

# opening - erosion followed by dilation
def opening(image):
    kernel = np.ones((5, 5), np.uint8)
    return cv2.morphologyEx(image, cv2.MORPH_OPEN, kernel)

# canny edge detection
def canny(image):
    return cv2.Canny(image, 100, 200)

# skew correction
def deskew(image):
    coords = np.column_stack(np.where(image > 0))
    angle = cv2.minAreaRect(coords)[-1]
    if angle < -45:
        angle = -(90 + angle)

    else:
        angle = -angle
    (h, w) = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
    return rotated

# template matching
def match_template(image, template):
    return cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)

doc = fitz.open("pdf", open(wfname, 'rb').read())
print(doc)
for i, page in enumerate(doc):
    text = page.getText()
    print(text)
    if not text:
        images = page.getImageList()

        print(images)

        pm = fitz.Pixmap(doc, images[0][0])  # prende la prima immagine

        size = [pm.width, pm.height]
        img = Image.frombytes("L", [pm.width, pm.height], pm.samples)

        img = img.rotate(360 - page.rotation, expand=True)

        img = img.convert('RGB')

        img = np.array(img)

        img = cv2.resize(img, None, fx=3, fy=3)

        cv2.imwrite(PATH+'img1.png', img)

        image = get_grayscale(img)
        # cv2.imshow('greyscale', image)
        cv2.imwrite(PATH+'greyscale.png', image)

        image = remove_noise(img)
        # cv2.imshow('remove noise', image)
        cv2.imwrite(PATH+'removenoise.png', image)

        #
        # image = thresholding(img)
        # cv2.imshow('thresholding', img)

        image = dilate(img)
        # cv2.imshow('dilate', image)
        cv2.imwrite(PATH+'dilate.png', image)

        image = erode(img)
        # cv2.imshow('erode', image)
        cv2.imwrite(PATH+'erode.png', image)

        image = opening(img)
        # cv2.imshow('opening', image)
        cv2.imwrite(PATH+'opening.png', image)

        image = canny(img)
        # cv2.imshow('canny', image)
        cv2.imwrite(PATH+'canny.png', image)

        img_new = remove_noise(img)
        img_new = erode(img_new)
        cv2.imwrite(PATH+'final.png', img_new)

        config = "--psm 6 --oem 1"# --user-words /home/cecilia/Documenti/images/vocab.txt"
        text = pytesseract.image_to_string(img_new, config=config, lang='ita')
        print(text)

        print()
        print()
        print('*'*20)
        print()
        print()

        text = pytesseract.image_to_string(img, config=config, lang='ita')
        print(text)

        # image = deskew(img)
        # cv2.imshow('deskew', image)

        # cv2.waitKey(0)

        # plt.imshow(img)
        # plt.show()

        # img = cv2.imread("book_page.jpg")

        # print(img.show())

# for i, page in enumerate(doc):
#     text = page.getText()
#
#     logger.info(f"Pagina {i}")
#     try:
#         if not text:  # da passare a tesseract
#             images = page.getImageList()
#
#             pm = fitz.Pixmap(doc, images[0][0])  # prende la prima immagine
#             logger.info(f"{images} {len(pm.samples)}")
#             size = [pm.width, pm.height]
#             img = Image.frombytes("L", [pm.width, pm.height], pm.samples)
#             img = img.rotate(360 - page.rotation, expand=True)
#
#             text = await loop.run_in_executor(POOL, functools.partial(pytesseract.image_to_string, lang="ita"),
#                                               img)
#             yield dict(i=i, text=text, filename = file.name )
#         else:
#             yield dict(i=i, text=text, filename = file.name)
#     except Exception as inst:
#         logging.exception(inst)


import argparse

from utils.constants import STRFTIME_COLON_SEPARATED, STRFTIME_DASH_SEPARATED
from evaluators.gnss_evaluator import GNSSEvaluator
from graphs import GraphsBuilder
from query import DBQuery
from utils.utils import initialize_logger, convert_strftime_format

parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument(
    "-s",
    "--start_time",
    help="The beginning of the data's time frame, format for example: -s \"2020-09-04 08:14:00\"\n"
         "Note: This field is mandatory!",
    dest="start_time",
    required=True,
)
parser.add_argument(
    "-e",
    "--end_time",
    help="The end of the data's time frame, format for example: -e \"2020-09-04 08:25:40\"\n"
         "Note: This field is mandatory",
    dest="end_time",
    required=True,
)
parser.add_argument(
    "-u",
    "--upload_graphs",
    help="A flag for whether to upload the graphs to google drive or not, default is False",
    dest="upload_graphs",
    default=False,
)
parser.add_argument(
    "-m",
    "--min_speed_threshold",
    help="The minimum speed threshold that used to ignore \"noises\", default = 10",
    dest="min_speed_threshold",
    default=10,
    type=float
)
parser.add_argument(
    "-r",
    "--speed_error_threshold",
    help="The threshold for marking data rows with speed error, default = 0.3",
    dest="speed_error_threshold",
    default=0.3,
    type=float
)
parser.add_argument(
    "-t",
    "--time_error_threshold",
    help="The threshold for marking data rows with time error, default = 0.15",
    dest="time_error_threshold",
    default=0.15,
    type=float
)
args = parser.parse_args()

if __name__ == '__main__':
    logger = initialize_logger()
    logger.info("Starting to evaluate data for time range {} -> {}".format(args.start_time, args.end_time))

    # Get the data
    run_data_by_time_range = DBQuery().get_run_data_by_time_range(start_time=args.start_time, end_time=args.end_time)

    # Evaluate GNSS data for each data row
    evaluated_run_data = GNSSEvaluator(
        min_speed_threshold=args.min_speed_threshold,
        speed_error_threshold=args.speed_error_threshold,
        time_error_threshold=args.time_error_threshold
    ).evaluate(run_data_by_time_range)

    converted_start_time = convert_strftime_format(args.start_time, STRFTIME_COLON_SEPARATED, STRFTIME_DASH_SEPARATED)
    converted_end_time = convert_strftime_format(args.end_time, STRFTIME_COLON_SEPARATED, STRFTIME_DASH_SEPARATED)

    # Build graphs according to evaluated data
    graphs_builder = GraphsBuilder(converted_start_time, converted_end_time, args.upload_graphs)
    graphs_builder.plot_long_to_lat(evaluated_run_data)
    graphs_builder.plot_time_to_speed(evaluated_run_data)
    graphs_builder.plot_time_to_time_delta(evaluated_run_data)

    logger.info("Done evaluating data".format(args.start_time, args.end_time))


import pandas as pd
from os import getenv
import sqlalchemy

heroku_secret_key = getenv('HEROKU_LINK')
db = sqlalchemy.create_engine(heroku_secret_key)

sql_1 = """CREATE TABLE users(
    id SERIAL PRIMARY KEY,
    name TEXT UNIQUE,
    password TEXT,
    role INTEGER
);"""

db.execute(sql_1)
sql_2 = """CREATE TABLE words(
    id SERIAL PRIMARY KEY,
    class TEXT,
    kanji TEXT,
    hiragana TEXT,
    roomaji TEXT,
    english TEXT
);"""

db.execute(sql_2)

input_data = pd.read_csv('words.csv', delimiter=',')

input_data.to_sql('words', db, if_exists='append', index=False)


import socket
from PyQt5.QtCore import QThread, pyqtSignal
from urllib.error import URLError
from ..models.api import API

class LoadItemsThread(QThread):
    progress_signal = pyqtSignal(API, list, int)
    error_signal = pyqtSignal(Exception)
    finished_signal = pyqtSignal(list)

    def __init__(self, api_collections, extent, start_time, end_time, query,
                 on_progress=None, on_error=None, on_finished=None):
        QThread.__init__(self)
        self.current_page = 0

        self.api_collections = api_collections
        self.extent = extent
        self.start_time = start_time
        self.end_time = end_time
        self.query = query
        self.on_progress = on_progress
        self.on_error = on_error
        self.on_finished = on_finished
        self._current_collections = []

        self.progress_signal.connect(self.on_progress)
        self.error_signal.connect(self.on_error)
        self.finished_signal.connect(self.on_finished)

    def run(self):
        try:
            all_items = []
            for api_collection in self.api_collections:
                self.current_page = 0
                api = api_collection['api']
                collections = api_collection['collections']
                self._current_collections = collections

                items = api.search_items(collections,
                                         self.extent,
                                         self.start_time,
                                         self.end_time,
                                         self.query,
                                         on_next_page=self.on_next_page)
                all_items.extend(items)
            self.finished_signal.emit(all_items)
        except URLError as e:
            self.error_signal.emit(e)
        except socket.timeout as e:
            self.error_signal.emit(e)

    def on_next_page(self, api):
        self.current_page += 1
        self.progress_signal.emit(
            api,
            self._current_collections,
            self.current_page
        )


import asyncio
import logging
from datetime import timedelta
from typing import Optional
from unittest.mock import AsyncMock, Mock

from azure.servicebus.aio import ServiceBusClient, ServiceBusReceiver, ServiceBusSender
from azure.servicebus.amqp import AmqpAnnotatedMessage
from azure.servicebus import ServiceBusReceivedMessage, ServiceBusMessage
from azure.servicebus._common.utils import utc_now

from .consumer_app import ConsumerApp

# Based on https://github.com/Azure/azure-sdk-for-python/blob/1356c85959f4ba182a92491f6d99d016411c8ab1/sdk/servicebus/azure-servicebus/tests/mocks.py#L18
class MockReceivedMessage(ServiceBusReceivedMessage):
    def __init__(self, prevent_renew_lock=False, exception_on_renew_lock=False, **kwargs):
        self._lock_duration = kwargs.get("lock_duration", 2)

        # Create AmqpAnnotatedMessage from value_body as per:
        # https://github.com/Azure/azure-sdk-for-python/blob/28e19bd7d8221479b9091f70aa6a84d40c3549a0/sdk/servicebus/azure-servicebus/samples/async_samples/send_and_receive_amqp_annotated_message_async.py#L49-L60
        # value_body = kwargs.pop("value_body", None)
        data_body = kwargs.pop("data_body", None)
        value_message = AmqpAnnotatedMessage(
            data_body=data_body,
        )
        self._raw_amqp_message = value_message
        self.message_id = "todo-message-id"

        self._received_timestamp_utc = utc_now()
        self.locked_until_utc = self._received_timestamp_utc + timedelta(seconds=self._lock_duration)
        self._settled = False
        self._receiver = kwargs.pop("receiver", None)

        self._prevent_renew_lock = prevent_renew_lock
        self._exception_on_renew_lock = exception_on_renew_lock

    @property
    def _lock_expired(self):
        if self.locked_until_utc and self.locked_until_utc <= utc_now():
            return True
        return False

    @property
    def locked_until_utc(self):
        return self._locked_until_utc

    @locked_until_utc.setter
    def locked_until_utc(self, value):
        self._locked_until_utc = value

class SentMessage:
    topic_name: str
    message: ServiceBusMessage

    def __init__(self, topic_name: str, message: ServiceBusMessage) -> None:
        self.topic_name = topic_name
        self.message = message

class MockServiceBusClientBuilder:
    _topics: dict  # key: topic name, value: (dict keyed on subscription name, value: list of messages)
    _topic_subscription_receivers = dict[str, ServiceBusReceiver]  # keyed on <topic_name>|<subscription_name>
    sentMessages: list[SentMessage]

    def __init__(self):
        self._topics = {}
        self._topic_subscription_receivers = {}
        self.sentMessages = []

    def add_messages_for_topic_subscription(self, topic_name: str, subscription_name: str, messages: list[str]):
        topic = self._topics.get(topic_name)
        if topic is None:
            topic = {}
            self._topics[topic_name] = topic

        topic_subscription = topic.get(subscription_name)
        if topic_subscription is not None:
            raise Exception(f"Messages already added for topic {topic_name} and subscription {subscription_name}")

        topic[subscription_name] = messages
        return self

    def get_subscription_receiver(self, topic_name, subscription_name, auto_lock_renewer=None):
        key = f"{topic_name}|{subscription_name}"
        receiver = self._topic_subscription_receivers.get(key)
        if not receiver is None:
            return receiver

        # else create a new receiver
        topic = self._topics.get(topic_name)
        if topic is None:
            raise Exception(f"No messages added for topic {topic_name}")
        messages = topic.get(subscription_name)
        if messages is None:
            raise Exception(f"No messages added for topic {topic_name} and subscription {subscription_name}")

        receiver = AsyncMock(spec=ServiceBusReceiver)
        receiver._running = False

        async def receive_messages(max_message_count=None, max_wait_time=None):
            nonlocal messages
            logging.info("In receive_messages")
            if messages == []:
                logging.info(f"No messages to return - sleeping (max_wait_time: {max_wait_time})")
                await asyncio.sleep(max_wait_time or 1)
                return []

            messages_to_return = [MockReceivedMessage(data_body=message, receiver=receiver) for message in messages]
            messages = []
            logging.info(f"returning messages: {messages_to_return}")
            return messages_to_return

        receiver.receive_messages = receive_messages

        # save receiver (enables us to retrieve the receiver in the test code to make assertions on it)
        self._topic_subscription_receivers[key] = receiver

        return receiver

    def get_topic_sender(self, topic_name):
        sender = Mock(spec=ServiceBusSender)

        async def send_messages(message: ServiceBusMessage):
            self.sentMessages.append(SentMessage(topic_name, message))

        sender.send_messages = AsyncMock(side_effect=send_messages)
        return sender

    def build(self):
        mock_sb_client = AsyncMock(spec=ServiceBusClient)
        mock_sb_client.get_subscription_receiver = Mock(side_effect=self.get_subscription_receiver)
        mock_sb_client.get_topic_sender = Mock(side_effect=self.get_topic_sender)
        return mock_sb_client

async def run_app_with_timeout(app: ConsumerApp, timeout_seconds: int = 0.1, filter: Optional[list[str]] = None):
    async def cancel_after_n_seconds(n):
        await asyncio.sleep(n)
        app.cancel()

    logging.info("Calling app.run...")
    await asyncio.gather(app.run(filter=filter), cancel_after_n_seconds(timeout_seconds))