"""!

@brief Cluster analysis algorithm: X-Means
@details Implementation based on papers @cite article::xmeans::1, @cite article::xmeans::mndl

@authors Andrei Novikov (pyclustering@yandex.ru)
@date 2014-2020
@copyright BSD-3-Clause

"""

import copy
import numpy

from enum import IntEnum
from math import log

from pyclustering.cluster.encoder import type_encoding
from pyclustering.cluster.kmeans import kmeans
from pyclustering.cluster.center_initializer import kmeans_plusplus_initializer

from pyclustering.core.metric_wrapper import metric_wrapper
from pyclustering.core.wrapper import ccore_library

import pyclustering.core.xmeans_wrapper as wrapper

from pyclustering.utils import distance_metric, type_metric

class splitting_type(IntEnum):
    """!
    @brief Enumeration of splitting types that can be used as splitting creation of cluster in X-Means algorithm.

    """

    ## Bayesian information criterion (BIC) to approximate the correct number of clusters.
    ## Kass's formula is used to calculate BIC:
    ## \f[BIC(\theta) = L(D) - \frac{1}{2}pln(N)\f]
    ##
    ## The number of free parameters \f$p\f$ is simply the sum of \f$K - 1\f$ class probabilities, \f$MK\f$ centroid coordinates, and one variance estimate:
    ## \f[p = (K - 1) + MK + 1\f]
    ##
    ## The log-likelihood of the data:
    ## \f[L(D) = n_jln(n_j) - n_jln(N) - \frac{n_j}{2}ln(2\pi) - \frac{n_jd}{2}ln(\hat{\sigma}^2) - \frac{n_j - K}{2}\f]
    ##
    ## The maximum likelihood estimate (MLE) for the variance:
    ## \f[\hat{\sigma}^2 = \frac{1}{N - K}\sum\limits_{j}\sum\limits_{i}||x_{ij} - \hat{C}_j||^2\f]
    BAYESIAN_INFORMATION_CRITERION = 0

    ## Minimum noiseless description length (MNDL) to approximate the correct number of clusters @cite article::xmeans::mndl.
    ## Beheshti's formula is used to calculate upper bound:
    ## \f[Z = \frac{\sigma^2 \sqrt{2K} }{N}(\sqrt{2K} + \beta) + W - \sigma^2 + \frac{2\alpha\sigma}{\sqrt{N}}\sqrt{\frac{\alpha^2\sigma^2}{N} + W - \left(1 - \frac{K}{N}\right)\frac{\sigma^2}{2}} + \frac{2\alpha^2\sigma^2}{N}\f]
    ##
    ## where \f$\alpha\f$ and \f$\beta\f$ represent the parameters for validation probability and confidence probability.
    ##
    ## To improve clustering results some contradiction is introduced:
    ## \f[W = \frac{1}{n_j}\sum\limits_{i}||x_{ij} - \hat{C}_j||\f]
    ## \f[\hat{\sigma}^2 = \frac{1}{N - K}\sum\limits_{j}\sum\limits_{i}||x_{ij} - \hat{C}_j||\f]
    MINIMUM_NOISELESS_DESCRIPTION_LENGTH = 1

class xmeans:
    """!
    @brief Class represents clustering algorithm X-Means.
    @details X-means clustering method starts with the assumption of having a minimum number of clusters,
             and then dynamically increases them. X-means uses specified splitting criterion to control
             the process of splitting clusters. Method K-Means++ can be used for calculation of initial centers.

             CCORE implementation of the algorithm uses thread pool to parallelize the clustering process.

    Here example how to perform cluster analysis using X-Means algorithm:
    @code
        from pyclustering.cluster import cluster_visualizer
        from pyclustering.cluster.xmeans import xmeans
        from pyclustering.cluster.center_initializer import kmeans_plusplus_initializer
        from pyclustering.utils import read_sample
        from pyclustering.samples.definitions import SIMPLE_SAMPLES

        # Read sample 'simple3' from file.
        sample = read_sample(SIMPLE_SAMPLES.SAMPLE_SIMPLE3)

        # Prepare initial centers - amount of initial centers defines amount of clusters from which X-Means will
        # start analysis.
        amount_initial_centers = 2
        initial_centers = kmeans_plusplus_initializer(sample, amount_initial_centers).initialize()

        # Create instance of X-Means algorithm. The algorithm will start analysis from 2 clusters, the maximum
        # number of clusters that can be allocated is 20.
        xmeans_instance = xmeans(sample, initial_centers, 20)
        xmeans_instance.process()

        # Extract clustering results: clusters and their centers
        clusters = xmeans_instance.get_clusters()
        centers = xmeans_instance.get_centers()

        # Print total sum of metric errors
        print("Total WCE:", xmeans_instance.get_total_wce())

        # Visualize clustering results
        visualizer = cluster_visualizer()
        visualizer.append_clusters(clusters, sample)
        visualizer.append_cluster(centers, None, marker='*', markersize=10)
        visualizer.show()
    @endcode

    Visualization of clustering results that were obtained using code above and where X-Means algorithm allocates four clusters.
    @image html xmeans_clustering_simple3.png "Fig. 1. X-Means clustering results (data 'Simple3')."

    By default X-Means clustering algorithm uses Bayesian Information Criterion (BIC) to approximate the correct number
    of clusters. There is an example where another criterion Minimum Noiseless Description Length (MNDL) is used in order
    to find optimal amount of clusters:
    @code
        from pyclustering.cluster import cluster_visualizer
        from pyclustering.cluster.xmeans import xmeans, splitting_type
        from pyclustering.cluster.center_initializer import kmeans_plusplus_initializer
        from pyclustering.utils import read_sample
        from pyclustering.samples.definitions import FCPS_SAMPLES

        # Read sample 'Target'.
        sample = read_sample(FCPS_SAMPLES.SAMPLE_TARGET)

        # Prepare initial centers - amount of initial centers defines amount of clusters from which X-Means will start analysis.
        random_seed = 1000
        amount_initial_centers = 3
        initial_centers = kmeans_plusplus_initializer(sample, amount_initial_centers, random_state=random_seed).initialize()

        # Create instance of X-Means algorithm with MNDL splitting criterion.
        xmeans_mndl = xmeans(sample, initial_centers, 20, splitting_type=splitting_type.MINIMUM_NOISELESS_DESCRIPTION_LENGTH, random_state=random_seed)
        xmeans_mndl.process()

        # Extract X-Means MNDL clustering results.
        mndl_clusters = xmeans_mndl.get_clusters()

        # Visualize clustering results.
        visualizer = cluster_visualizer(titles=['X-Means with MNDL criterion'])
        visualizer.append_clusters(mndl_clusters, sample)
        visualizer.show()
    @endcode

    @image html xmeans_clustering_mndl_target.png "Fig. 2. X-Means MNDL clustering results (data 'Target')."

    As in many others algorithms, it is possible to specify metric that should be used for cluster analysis, for
    example, Chebyshev distance metric:
    @code
        # Create instance of X-Means algorithm with Chebyshev distance metric.
        chebyshev_metric = distance_metric(type_metric.CHEBYSHEV)
        xmeans_instance = xmeans(sample, initial_centers, max_clusters_amount, metric=chebyshev_metric).process()
    @endcode

    @see center_initializer

    """

    def __init__(self, data, initial_centers=None, kmax=20, tolerance=0.001, criterion=splitting_type.BAYESIAN_INFORMATION_CRITERION, ccore=True, **kwargs):
        """!
        @brief Constructor of clustering algorithm X-Means.

        @param[in] data (array_like): Input data that is presented as list of points (objects), each point should be represented by list or tuple.
        @param[in] initial_centers (list): Initial coordinates of centers of clusters that are represented by list: `[center1, center2, ...]`,
                    if it is not specified then X-Means starts from the random center.
        @param[in] kmax (uint): Maximum number of clusters that can be allocated.
        @param[in] tolerance (double): Stop condition for each iteration: if maximum value of change of centers of clusters is less than tolerance than algorithm will stop processing.
        @param[in] criterion (splitting_type): Type of splitting creation (by default `splitting_type.BAYESIAN_INFORMATION_CRITERION`).
        @param[in] ccore (bool): Defines if C++ pyclustering library should be used instead of Python implementation.
        @param[in] **kwargs: Arbitrary keyword arguments (available arguments: `repeat`, `random_state`, `metric`, `alpha`, `beta`).

        Keyword Args:

            - repeat (unit): How many times K-Means should be run to improve parameters (by default is `1`).
               With larger `repeat` values suggesting higher probability of finding global optimum.
            - random_state (int): Seed for random state (by default is `None`, current system time is used).
            - metric (distance_metric): Metric that is used for distance calculation between two points (by default
               euclidean square distance).
            - alpha (double): Parameter distributed [0.0, 1.0] for alpha probabilistic bound \f$Q\left(\alpha\right)\f$.
               The parameter is used only in case of MNDL splitting criterion, in all other cases this value is ignored.
            - beta (double): Parameter distributed [0.0, 1.0] for beta probabilistic bound \f$Q\left(\beta\right)\f$.
               The parameter is used only in case of MNDL splitting criterion, in all other cases this value is ignored.

        """

        self.__pointer_data = numpy.array(data)
        self.__clusters = []
        self.__random_state = kwargs.get('random_state', None)
        self.__metric = copy.copy(kwargs.get('metric', distance_metric(type_metric.EUCLIDEAN_SQUARE)))

        if initial_centers is not None:
            self.__centers = numpy.array(initial_centers)
        else:
            self.__centers = kmeans_plusplus_initializer(data, 2, random_state=self.__random_state).initialize()

        self.__kmax = kmax
        self.__tolerance = tolerance
        self.__criterion = criterion
        self.__total_wce = 0.0
        self.__repeat = kwargs.get('repeat', 1)
        self.__alpha = kwargs.get('alpha', 0.9)
        self.__beta = kwargs.get('beta', 0.9)

        self.__ccore = ccore and self.__metric.get_type() != type_metric.USER_DEFINED
        if self.__ccore is True:
            self.__ccore = ccore_library.workable()

        self.__verify_arguments()

    def process(self):
        """!
        @brief Performs cluster analysis in line with rules of X-Means algorithm.

        @return (xmeans) Returns itself (X-Means instance).

        @see get_clusters()
        @see get_centers()

        """

        if self.__ccore is True:
            self.__process_by_ccore()

        else:
            self.__process_by_python()

        return self

    def __process_by_ccore(self):
        """!
        @brief Performs cluster analysis using CCORE (C/C++ part of pyclustering library).

        """

        ccore_metric = metric_wrapper.create_instance(self.__metric)

        result = wrapper.xmeans(self.__pointer_data, self.__centers, self.__kmax, self.__tolerance, self.__criterion,
                                self.__alpha, self.__beta, self.__repeat, self.__random_state,
                                ccore_metric.get_pointer())

        self.__clusters = result[0]
        self.__centers = result[1]
        self.__total_wce = result[2][0]

    def __process_by_python(self):
        """!
        @brief Performs cluster analysis using python code.

        """

        self.__clusters = []
        while len(self.__centers) <= self.__kmax:
            current_cluster_number = len(self.__centers)

            self.__clusters, self.__centers, _ = self.__improve_parameters(self.__centers)
            allocated_centers = self.__improve_structure(self.__clusters, self.__centers)

            if current_cluster_number == len(allocated_centers):
                break
            else:
                self.__centers = allocated_centers

        self.__clusters, self.__centers, self.__total_wce = self.__improve_parameters(self.__centers)

    def predict(self, points):
        """!
        @brief Calculates the closest cluster to each point.

        @param[in] points (array_like): Points for which closest clusters are calculated.

        @return (list) List of closest clusters for each point. Each cluster is denoted by index. Return empty
                 collection if 'process()' method was not called.

        An example how to calculate (or predict) the closest cluster to specified points.
        @code
            from pyclustering.cluster.xmeans import xmeans
            from pyclustering.samples.definitions import SIMPLE_SAMPLES
            from pyclustering.utils import read_sample

            # Load list of points for cluster analysis.
            sample = read_sample(SIMPLE_SAMPLES.SAMPLE_SIMPLE3)

            # Initial centers for sample 'Simple3'.
            initial_centers = [[0.2, 0.1], [4.0, 1.0], [2.0, 2.0], [2.3, 3.9]]

            # Create instance of X-Means algorithm with prepared centers.
            xmeans_instance = xmeans(sample, initial_centers)

            # Run cluster analysis.
            xmeans_instance.process()

            # Calculate the closest cluster to following two points.
            points = [[0.25, 0.2], [2.5, 4.0]]
            closest_clusters = xmeans_instance.predict(points)
            print(closest_clusters)
        @endcode

        """
        nppoints = numpy.array(points)
        if len(self.__clusters) == 0:
            return []

        self.__metric.enable_numpy_usage()

        npcenters = numpy.array(self.__centers)
        differences = numpy.zeros((len(nppoints), len(npcenters)))
        for index_point in range(len(nppoints)):
            differences[index_point] = self.__metric(nppoints[index_point], npcenters)

        self.__metric.disable_numpy_usage()

        return numpy.argmin(differences, axis=1)

    def get_clusters(self):
        """!
        @brief Returns list of allocated clusters, each cluster contains indexes of objects in list of data.

        @return (list) List of allocated clusters.

        @see process()
        @see get_centers()
        @see get_total_wce()

        """

        return self.__clusters

    def get_centers(self):
        """!
        @brief Returns list of centers for allocated clusters.

        @return (list) List of centers for allocated clusters.

        @see process()
        @see get_clusters()
        @see get_total_wce()

        """

        return self.__centers

    def get_cluster_encoding(self):
        """!
        @brief Returns clustering result representation type that indicate how clusters are encoded.

        @return (type_encoding) Clustering result representation.

        @see get_clusters()

        """

        return type_encoding.CLUSTER_INDEX_LIST_SEPARATION

    def get_total_wce(self):
        """!
        @brief Returns sum of Euclidean Squared metric errors (SSE - Sum of Squared Errors).
        @details Sum of metric errors is calculated using distance between point and its center:
                 \f[error=\sum_{i=0}^{N}euclidean_square_distance(x_{i}-center(x_{i}))\f]

        @see process()
        @see get_clusters()

        """

        return self.__total_wce

    def __search_optimial_parameters(self, local_data):
        """!
        @brief Split data of the region into two cluster and tries to find global optimum by running k-means clustering
                several times (defined by 'repeat' argument).

        @param[in] local_data (list): Points of a region that should be split into two clusters.

        @return (tuple) List of allocated clusters, list of centers and total WCE (clusters, centers, wce).

        """
        optimal_wce, optimal_centers, optimal_clusters = float('+inf'), None, None

        for _ in range(self.__repeat):
            candidates = 5
            if len(local_data) < candidates:
                candidates = len(local_data)

            local_centers = kmeans_plusplus_initializer(local_data, 2, candidates, random_state=self.__random_state).initialize()

            kmeans_instance = kmeans(local_data, local_centers, tolerance=self.__tolerance, ccore=False, metric=self.__metric)
            kmeans_instance.process()

            local_wce = kmeans_instance.get_total_wce()
            if local_wce < optimal_wce:
                optimal_centers = kmeans_instance.get_centers()
                optimal_clusters = kmeans_instance.get_clusters()
                optimal_wce = local_wce

        return optimal_clusters, optimal_centers, optimal_wce

    def __improve_parameters(self, centers, available_indexes=None):
        """!
        @brief Performs k-means clustering in the specified region.

        @param[in] centers (list): Cluster centers, if None then automatically generated two centers using center initialization method.
        @param[in] available_indexes (list): Indexes that defines which points can be used for k-means clustering, if None then all points are used.

        @return (tuple) List of allocated clusters, list of centers and total WCE (clusters, centers, wce).

        """

        if available_indexes and len(available_indexes) == 1:
            index_center = available_indexes[0]
            return [available_indexes], self.__pointer_data[index_center], 0.0

        local_data = self.__pointer_data
        if available_indexes:
            local_data = [self.__pointer_data[i] for i in available_indexes]

        local_centers = centers
        if centers is None:
            clusters, local_centers, local_wce = self.__search_optimial_parameters(local_data)
        else:
            kmeans_instance = kmeans(local_data, local_centers, tolerance=self.__tolerance, ccore=False, metric=self.__metric).process()

            local_wce = kmeans_instance.get_total_wce()
            local_centers = kmeans_instance.get_centers()
            clusters = kmeans_instance.get_clusters()

        if available_indexes:
            clusters = self.__local_to_global_clusters(clusters, available_indexes)

        return clusters, local_centers, local_wce

    def __local_to_global_clusters(self, local_clusters, available_indexes):
        """!
        @brief Converts clusters in local region define by 'available_indexes' to global clusters.

        @param[in] local_clusters (list): Local clusters in specific region.
        @param[in] available_indexes (list): Map between local and global point's indexes.

        @return Global clusters.

        """

        clusters = []
        for local_cluster in local_clusters:
            current_cluster = []
            for index_point in local_cluster:
                current_cluster.append(available_indexes[index_point])

            clusters.append(current_cluster)

        return clusters

    def __improve_structure(self, clusters, centers):
        """!
        @brief Check for best structure: divides each cluster into two and checks for best results using splitting criterion.

        @param[in] clusters (list): Clusters that have been allocated (each cluster contains indexes of points from data).
        @param[in] centers (list): Centers of clusters.

        @return (list) Allocated centers for clustering.

        """

        allocated_centers = []
        amount_free_centers = self.__kmax - len(centers)

        for index_cluster in range(len(clusters)):
            # solve k-means problem for children where data of parent are used.
            (parent_child_clusters, parent_child_centers, _) = self.__improve_parameters(None, clusters[index_cluster])

            # If it's possible to split current data
            if len(parent_child_clusters) > 1:
                # Calculate splitting criterion
                parent_scores = self.__splitting_criterion([clusters[index_cluster]], [centers[index_cluster]])
                child_scores = self.__splitting_criterion([parent_child_clusters[0], parent_child_clusters[1]], parent_child_centers)

                split_require = False

                # Reallocate number of centers (clusters) in line with scores
                if self.__criterion == splitting_type.BAYESIAN_INFORMATION_CRITERION:
                    if parent_scores < child_scores:
                        split_require = True

                elif self.__criterion == splitting_type.MINIMUM_NOISELESS_DESCRIPTION_LENGTH:
                    # If its score for the split structure with two children is smaller than that for the parent structure,
                    # then representing the data samples with two clusters is more accurate in comparison to a single parent cluster.
                    if parent_scores > child_scores:
                        split_require = True

                if (split_require is True) and (amount_free_centers > 0):
                    allocated_centers.append(parent_child_centers[0])
                    allocated_centers.append(parent_child_centers[1])

                    amount_free_centers -= 1
                else:
                    allocated_centers.append(centers[index_cluster])

            else:
                allocated_centers.append(centers[index_cluster])

        return allocated_centers

    def __splitting_criterion(self, clusters, centers):
        """!
        @brief Calculates splitting criterion for input clusters.

        @param[in] clusters (list): Clusters for which splitting criterion should be calculated.
        @param[in] centers (list): Centers of the clusters.

        @return (double) Returns splitting criterion. High value of splitting criterion means that current structure is
                 much better.

        @see __bayesian_information_criterion(clusters, centers)
        @see __minimum_noiseless_description_length(clusters, centers)

        """

        if self.__criterion == splitting_type.BAYESIAN_INFORMATION_CRITERION:
            return self.__bayesian_information_criterion(clusters, centers)

        elif self.__criterion == splitting_type.MINIMUM_NOISELESS_DESCRIPTION_LENGTH:
            return self.__minimum_noiseless_description_length(clusters, centers)

        else:
            assert 0

    def __minimum_noiseless_description_length(self, clusters, centers):
        """!
        @brief Calculates splitting criterion for input clusters using minimum noiseless description length criterion.

        @param[in] clusters (list): Clusters for which splitting criterion should be calculated.
        @param[in] centers (list): Centers of the clusters.

        @return (double) Returns splitting criterion in line with bayesian information criterion.
                Low value of splitting cretion means that current structure is much better.

        @see __bayesian_information_criterion(clusters, centers)

        """

        score = float('inf')

        W = 0.0
        K = len(clusters)
        N = 0.0

        sigma_square = 0.0

        alpha = self.__alpha
        alpha_square = alpha * alpha
        beta = self.__beta

        for index_cluster in range(0, len(clusters), 1):
            Ni = len(clusters[index_cluster])
            if Ni == 0:
                return float('inf')

            Wi = 0.0
            for index_object in clusters[index_cluster]:
                Wi += self.__metric(self.__pointer_data[index_object], centers[index_cluster])

            sigma_square += Wi
            W += Wi / Ni
            N += Ni

        if N - K > 0:
            sigma_square /= (N - K)
            sigma = sigma_square ** 0.5

            Kw = (1.0 - K / N) * sigma_square
            Ksa = (2.0 * alpha * sigma / (N ** 0.5)) * (alpha_square * sigma_square / N + W - Kw / 2.0) ** 0.5
            UQa = W - Kw + 2.0 * alpha_square * sigma_square / N + Ksa

            score = sigma_square * K / N + UQa + sigma_square * beta * ((2.0 * K) ** 0.5) / N

        return score

    def __bayesian_information_criterion(self, clusters, centers):
        """!
        @brief Calculates splitting criterion for input clusters using bayesian information criterion.

        @param[in] clusters (list): Clusters for which splitting criterion should be calculated.
        @param[in] centers (list): Centers of the clusters.

        @return (double) Splitting criterion in line with bayesian information criterion.
                High value of splitting criterion means that current structure is much better.

        @see __minimum_noiseless_description_length(clusters, centers)

        """

        scores = [float('inf')] * len(clusters)     # splitting criterion
        dimension = len(self.__pointer_data[0])

        # estimation of the noise variance in the data set
        sigma_sqrt = 0.0
        K = len(clusters)
        N = 0.0

        for index_cluster in range(0, len(clusters), 1):
            for index_object in clusters[index_cluster]:
                sigma_sqrt += self.__metric(self.__pointer_data[index_object], centers[index_cluster])

            N += len(clusters[index_cluster])

        if N - K > 0:
            sigma_sqrt /= (N - K)
            p = (K - 1) + dimension * K + 1

            # in case of the same points, sigma_sqrt can be zero (issue: #407)
            sigma_multiplier = 0.0
            if sigma_sqrt <= 0.0:
                sigma_multiplier = float('-inf')
            else:
                sigma_multiplier = dimension * 0.5 * log(sigma_sqrt)

            # splitting criterion
            for index_cluster in range(0, len(clusters), 1):
                n = len(clusters[index_cluster])

                L = n * log(n) - n * log(N) - n * 0.5 * log(2.0 * numpy.pi) - n * sigma_multiplier - (n - K) * 0.5

                # BIC calculation
                scores[index_cluster] = L - p * 0.5 * log(N)

        return sum(scores)

    def __verify_arguments(self):
        """!
        @brief Verify input parameters for the algorithm and throw exception in case of incorrectness.

        """
        if len(self.__pointer_data) == 0:
            raise ValueError("Input data is empty (size: '%d')." % len(self.__pointer_data))

        if len(self.__centers) == 0:
            raise ValueError("Initial centers are empty (size: '%d')." % len(self.__pointer_data))

        if self.__tolerance < 0:
            raise ValueError("Tolerance (current value: '%d') should be greater or equal to 0." %
                             self.__tolerance)

        if self.__repeat <= 0:
            raise ValueError("Repeat (current value: '%d') should be greater than 0." %
                             self.__repeat)

        if self.__alpha < 0.0 or self.__alpha > 1.0:
            raise ValueError("Parameter for the probabilistic bound Q(alpha) should in the following range [0, 1] "
                             "(current value: '%f')." % self.__alpha)

        if self.__beta < 0.0 or self.__beta > 1.0:
            raise ValueError("Parameter for the probabilistic bound Q(beta) should in the following range [0, 1] "
                             "(current value: '%f')." % self.__beta)


import pygame
from pygame.surface import *
from pygame.sprite import Sprite
from pygame.sprite import RenderUpdates as SpriteGroup
from pygame.sprite import spritecollide
from pygame.sprite import spritecollideany
from pygame.rect import Rect
from random import *
from config import *
from log import *

screen = None

def createFrames(image):
    fr_width = image.get_height()
    fr_size = fr_width, fr_width
    frames = []
    for frame_no in range(0, image.get_width(), fr_width):
        frame = pygame.Surface(fr_size)
        frame.blit(image, (0,0), ((frame_no,0), fr_size))
        frame.set_colorkey(PUCE)
        frames.append(frame)
    return frames

def initPygame():
    pygame.init()
    global screen

    screen = pygame.display.set_mode(RESOLUTION, FLAGS)

class Widget(Sprite):
    """Use Widget class for better movement tracking

       Widget class inherits from Sprite class.
       Test cases for class Widget

       >>> from widget import *
       >>> import pygame
       >>> s = pygame.surface.Surface((30,30))
       >>> w = Widget(s, (0,0,30,30), (0,0))
       >>> w.rect
       <rect(0, 0, 30, 30)>
       >>> w.update()
       >>> w.rect
       <rect(0, 0, 30, 30)>
       >>> w.getMovement()
       [0, 0]

       >>> w.setX(1)
       >>> w.getX()
       1

       >>> w.setY(4)
       >>> w.getY()
       4

       >>> w.setMovement((3,5))
       >>> w.getMovement()
       (3, 5)

       >>> w.getName()
       (0, 0)

       >>> w.setPosition((5,7))
       >>> w.getPosition()
       (5, 7)
    """

    def __init__(self, image, rect, name=''):
        """Instantiate a widget with a given surface,
           rectangle, and (x,y) movement pair.
        """
        Sprite.__init__(self)
        self.movement = [0, 0]
        self.rect = Rect(rect)
        self.lastRect = self.rect
        self.image = image
        self.name = name
        self.frames = []
        self.frameIndex = 0
        self.frameRate = 1
        self.timeDelay = WIDGETFRAMES
        self.lastUpdate = 0
        self.world = None
        self.undone = False
        self.id = self.rect.top + self.rect.left +\
                  self.rect.width + self.rect.height

    def attachToWorld(self, world):
        self.world = world
        self.id = self.world.curWidgetID
        self.world.curWidgetID += 1

    def startAnimation(self, frames, startIndex, frameRate):
        self.frames = frames
        self.frameIndex = startIndex
        self.frameRate = frameRate
        self.image = self.frames[startIndex]
        self.lastUpdate = self.timeDelay

    def __str__(self):
        return str(self.rect.left) + str(self.rect.top) + str(self.id)

    def setMovement(self, vector):
        """Set movement with a pair"""
        if(self.movement != [0,0]
           and vector == [0,0]):
            self.world.dirtyGroup.add(self)
        self.movement = vector

    def getMovement(self):
        """Return movement as a pair"""
        return self.movement

    def setStop(self):
        """Set movement to 0"""
        self.setMovement([0,0])

    def setY(self, y):
        """Set y-component of movement"""
        self.movement[1] = y

    def setX(self, x):
        """Set x-component of movement"""
        self.movement[0] = x

    def getX(self):
        """Get x-component of movement"""
        return self.movement[0]
    def getY(self):
        """Set y-component of movement"""
        return self.movement[1]

    def setPosition(self, pair):
        """Set x and y coords of Widget"""
        self.rect.topleft = pair

    def getPosition(self):
        """Get x and y coords of Widget"""
        return self.rect.topleft

    def update(self):
        """Move sprite according to its movement vector"""
        # Widget needs to be animated
        if (len(self.frames) > 0):
            if self.lastUpdate <= 0:
                self.frameIndex = (self.frameIndex+1)%(len(self.frames))
                self.image = self.frames[self.frameIndex]
                self.lastUpdate = self.timeDelay
                self.world.dirtyGroup.add(self)
            else:
                self.lastUpdate -= 1

        elif(self.getMovement != [0,0]):
            self.lastRect = Rect(self.rect)
            self.rect.move_ip(self.movement)
            self.world.dirtyGroup.add(self)

    def undoUpdate(self):
        """Widget returns to state prior to last update()"""
        self.rect = self.lastRect

    def getShadow(self):
        shadow = Sprite()
        shadow.rect = self.lastRect.move(0,0)
        return shadow

    def getName(self):
        """Get name of Widget"""
        return self.name

class WorldlessWidget(Widget):
    def update(self):
        """Move sprite according to its movement vector"""
        # Widget needs to be animated
        if (len(self.frames) > 0):
            if self.lastUpdate <= 0:
                self.frameIndex = (self.frameIndex+1)%(len(self.frames))
                self.image = self.frames[self.frameIndex]
                self.lastUpdate = self.timeDelay

        self.lastRect = Rect(self.rect)
        self.rect.move_ip(self.movement)


import asyncpg

from app.main import create_app
from app.settings import Settings
from tests.test_currencies import test_currencies_success

settings = Settings()

async def test_rate_success(aiohttp_client, loop):
    currencies = await test_currencies_success(aiohttp_client, loop)
    currency = currencies['results']['currencies'][0]['id']
    client = await aiohttp_client(create_app)

    conn = await asyncpg.connect(settings.dsn)
    try:
        resp = await client.get(f'/rate/{currency}/')
        assert resp.status == 200
        result = await resp.json()
        assert 'current_rate' in result
        assert 'avg_volume' in result

        resp = await client.get(f'/rate/?id={currency}')
        assert resp.status == 200
        result = await resp.json()
        assert 'current_rate' in result
        assert 'avg_volume' in result
    finally:
        await conn.close()

async def test_rate_not_found(aiohttp_client, loop):
    curency_name = 'BCH'
    client = await aiohttp_client(create_app)
    conn = await asyncpg.connect(settings.dsn)
    try:
        currency = await conn.fetchval('select id from currency where name=$1', curency_name)
        await conn.execute('delete from currency where name=$1', curency_name)
        resp = await client.get(f'/rate/{currency}/')
        assert resp.status == 404
        result = await resp.json()
        assert 'error' in result
        assert 'detail' in result
        assert 'not found transaction' == result['detail']

    finally:
        await conn.close()

async def test_rate_error(aiohttp_client, loop):
    curency_name = 'BCH'
    client = await aiohttp_client(create_app)
    resp = await client.get(f'/rate/{curency_name}/')
    assert resp.status == 400


#!/usr/bin/env python

# This is more of a work in progress, but this script will
# test the code for creating vespagrams with our curved wavefront correction.

import obspy
import numpy as np
import time
import matplotlib.pyplot as plt

import circ_array as c
from circ_beam import Vespagram_Lin, Vespagram_PWS, Baz_vespagram_PWS, Baz_vespagram_Lin
from array_plotting import plotting

# parameters
# phase of interest
phase = 'SKS'
phases = ['SKS','SKKS','ScS','Sdiff','sSKS','sSKKS','PS','SKKKS','pSKS']

# frequency band
fmin = 0.13
fmax = 0.26

st = obspy.read('./data/19970525/*SAC')

# get array metadata
event_time = c.get_eventtime(st)
geometry = c.get_geometry(st)
distances = c.get_distances(st,type='deg')
mean_dist = np.mean(distances)
stations = c.get_stations(st)

# get travel time information and define a window
Target_phase_times, time_header_times = c.get_predicted_times(st,phase)

avg_target_time = np.mean(Target_phase_times)
min_target = int(np.nanmin(Target_phase_times, axis=0))
max_target = int(np.nanmax(Target_phase_times, axis=0)) + 100

stime = event_time + min_target
etime = event_time + max_target

# trim the stream
# Normalise and cut seismogram around defined window
st = st.copy().trim(starttime=stime, endtime=etime)
st = st.normalize()

# get predicted slownesses and backazimuths
predictions = c.pred_baz_slow(
    stream=st, phases=phases, one_eighty=True)

# find the line with the predictions for the phase of interest
row = np.where((predictions == phase))[0]
P, S, BAZ, PRED_BAZ_X, PRED_BAZ_Y, PRED_AZ_X, PRED_AZ_Y, DIST, TIME = predictions[row, :][0]

# make the box around the prediction to search over
smin=float(S)-2
smax=float(S)+6
s_step=0.1

# filter
st = st.filter('bandpass', freqmin=fmin, freqmax=fmax,
                  corners=4, zerophase=True)

# get the traces and phase traces
Traces = c.get_traces(st)
Phase_traces = c.get_phase_traces(st)

# get sampleing rate
sampling_rate=st[0].stats.sampling_rate

# slowness vespagrams
vesp_lin = Vespagram_Lin(traces=Traces, sampling_rate=sampling_rate, geometry=geometry,
                         distance=mean_dist, baz=float(BAZ), smin=smin, smax=smax, s_space=s_step)

fig = plt.figure(figsize=(10,8))
ax = fig.add_subplot(211)
p = plotting(ax = ax)

p.plot_vespagram(vespagram=vesp_lin, ymin=smin, ymax=smax, y_space=s_step, tmin=min_target, tmax=max_target,
                 sampling_rate=sampling_rate, title="Slow - Lin", predictions=predictions, type='slow',
                 envelope=True)

vesp_pws = Vespagram_PWS(traces=Traces, phase_traces=Phase_traces, sampling_rate=sampling_rate, geometry=geometry,
                         distance=mean_dist, baz=float(BAZ), smin=smin, smax=smax, s_space=s_step, degree=2)

ax = fig.add_subplot(212)
p = plotting(ax = ax)
p.plot_vespagram(vespagram=vesp_pws, ymin=smin, ymax=smax, y_space=s_step, tmin=min_target, tmax=max_target,
                 sampling_rate=sampling_rate, title="Slow - PWS", predictions=predictions, type='slow',
                 envelope=True)
plt.tight_layout()

plt.show()
# backazimuth vespagrams

bmin=float(BAZ)-30
bmax=float(BAZ)+30
b_step=0.05

vesp_lin = Baz_vespagram_Lin(traces=Traces, sampling_rate=sampling_rate, geometry=geometry,
                         distance=mean_dist, slow=float(S), bmin=bmin, bmax=bmax, b_space=b_step)

fig = plt.figure(figsize=(10,8))
ax = fig.add_subplot(211)
p = plotting(ax = ax)

p.plot_vespagram(vespagram=vesp_lin, ymin=bmin, ymax=bmax, y_space=b_step, tmin=min_target, tmax=max_target,
                 sampling_rate=sampling_rate, title="Baz - Lin", predictions=predictions, type='baz',
                 envelope=True)

vesp_pws = Baz_vespagram_PWS(traces=Traces, phase_traces=Phase_traces, sampling_rate=sampling_rate, geometry=geometry,
                         distance=mean_dist, slow=float(S), bmin=bmin, bmax=bmax, b_space=b_step, degree=2)

ax = fig.add_subplot(212)
p = plotting(ax = ax)
p.plot_vespagram(vespagram=vesp_pws, ymin=bmin, ymax=bmax, y_space=b_step, tmin=min_target, tmax=max_target,
                 sampling_rate=sampling_rate, title="Baz - PWS", predictions=predictions, npeaks=5, type='baz',
                 envelope=True)
plt.tight_layout()
plt.show()


"""
Meta manager. Defines complex workflow in terms of lower level managers

For usage example see tests

"""
import re
from time import time
import logging

from collections import defaultdict, OrderedDict as odict
from copy import copy, deepcopy

import yaml

from shub_workflow.base import WorkFlowManager

from .utils import get_scheduled_jobs_specs

logger = logging.getLogger(__name__)

_STARTING_JOB_RE = re.compile("--starting-job(?:=(.+))?")

class GraphManager(WorkFlowManager):

    jobs_graph = {}
    base_failed_outcomes = ('failed', 'killed by oom', 'cancelled', 'cancel_timeout', 'memusage_exceeded',
                            'cancelled (stalled)')

    def __init__(self):
        self.__failed_outcomes = list(self.base_failed_outcomes)
        # Ensure jobs are traversed in the same order as they went pending.
        self.__pending_jobs = odict()
        self.__running_jobs = odict()
        self._available_resources = {}  # map resource : ammount
        self._acquired_resources = defaultdict(list)  # map resource : list of (job, ammount)
        self.__tasks = {}
        super(GraphManager, self).__init__()
        self.__start_time = defaultdict(time)
        for task in self.configure_workflow() or ():
            self._add_task(task)

    @property
    def description(self):
        return f"Workflow manager for {self.name!r}"

    def _add_task(self, task):
        assert task.task_id not in self.jobs_graph,\
            "Workflow inconsistency detected: task %s referenced twice." % task.task_id
        self.jobs_graph[task.task_id] = task.as_jobgraph_dict()
        self.__tasks[task.task_id] = task
        for ntask in task.get_next_tasks():
            self._add_task(ntask)

    def configure_workflow(self):
        raise NotImplementedError("configure_workflow() method need to be implemented.")

    def on_start(self):
        if not self.jobs_graph:
            self.argparser.error('Jobs graph configuration is empty.')
        if not self.args.starting_job and not self.args.resume_from_jobid:
            self.argparser.error('You must provide either --starting-job or --resume-from-jobid.')
        self._fill_available_resources()
        ran_tasks = self._maybe_setup_resume()
        self._setup_starting_jobs(ran_tasks)
        self.workflow_loop_enabled = True
        logger.info("Starting '%s' workflow", self.name)

    def _get_starting_jobs_from_resumed_job(self):
        starting_jobs = []
        job = self.get_project().jobs.get(self.args.resume_from_jobid)
        next_option_is_task = False
        for option in job.metadata.get('job_cmd'):
            if next_option_is_task:
                starting_jobs.append(option)
            else:
                m = _STARTING_JOB_RE.match(option)
                if m:
                    task = m.groups()[0]
                    if m:
                        starting_jobs.append(task)
                    else:
                        next_option_is_task = True
        return starting_jobs

    def _maybe_setup_resume(self):
        ran_tasks = []
        if self.args.resume_from_jobid:
            # fill tasks job ids
            logger.info("Will Resume from job (%s)", self.args.resume_from_jobid)
            for _, name, jobid in get_scheduled_jobs_specs(self, [self.args.resume_from_jobid]):
                mname, taskid = name.split('/')
                assert mname == self.name, "Resuming from wrong manager job: %s" % self.args.resume_from_jobid
                self.__tasks[taskid].append_jobid(jobid)
                ran_tasks.append(taskid)
        return ran_tasks

    def _setup_starting_jobs(self, ran_tasks, candidates=None):
        candidates = candidates or self.args.starting_job
        if not candidates:  # resuming
            candidates = self._get_starting_jobs_from_resumed_job()
        for taskid in candidates:
            if taskid in ran_tasks:
                logger.info("Task %s already done %s.", taskid, tuple(self.__tasks[taskid].get_scheduled_jobs()))
                next_tasks = [t.task_id for t in self.__tasks[taskid].get_next_tasks()]
                if next_tasks:
                    self._setup_starting_jobs(ran_tasks, next_tasks)
            else:
                self._add_initial_pending_job(taskid)
                logger.info("Resuming at task %s", taskid)

    def _fill_available_resources(self):
        """
        Ensure there are enough starting resources in order every job
        can run at some point
        """
        for job in self.jobs_graph.keys():
            for required_resources in self.__tasks[job].get_required_resources():
                for resource, req_amount in required_resources.items():
                    old_amount = self._available_resources.get(resource, 0)
                    if old_amount < req_amount:
                        logger.info("Increasing available resources count for %r"
                                    " from %r to %r.  Old value was not enough"
                                    " for job %r to run.",
                                    resource, old_amount, req_amount, job)
                        self._available_resources[resource] = req_amount

    def get_job(self, job, pop=False):
        if job not in self.jobs_graph:
            self.argparser.error('Invalid job: %s. Available jobs: %s' % (job, repr(self.jobs_graph.keys())))
        if pop:
            return self.jobs_graph.pop(job)
        return self.jobs_graph[job]

    def _add_initial_pending_job(self, job):
        wait_for = self.get_job(job).get('wait_for', [])
        self._add_pending_job(job, wait_for=tuple(wait_for))

    def _add_pending_job(self, job, wait_for=(), is_retry=False):
        if job in self.args.skip_job:
            return
        if job in self.__tasks:
            task = self.__tasks[job]
            parallelization = task.get_parallel_jobs()
        else:
            task_id = self.get_job(job).get('origin', job)
            task = self.__tasks[task_id]
            parallelization = 1
        if parallelization == 1:
            self.__pending_jobs[job] = {
                'wait_for': set(wait_for),
                'is_retry': is_retry,
                'wait_time': task.wait_time,
            }
        else:
            # Split parallelized task into N parallel jobs.
            basejobconf = self.get_job(job, pop=True)
            for i in range(parallelization):
                job_unit = "%s_%i" % (job, i)
                job_unit_conf = deepcopy(basejobconf)
                job_unit_conf['origin'] = job
                job_unit_conf['index'] = i

                for _, nextjobs in job_unit_conf.get('on_finish', {}).items():
                    if i != 0:  # only job 0 will conserve finish targets
                        for nextjob in copy(nextjobs):
                            if nextjob != 'retry':
                                if nextjob in self.jobs_graph:
                                    self.get_job(nextjob).setdefault('wait_for', []).append(job_unit)
                                    if nextjob in self.__pending_jobs:
                                        self.__pending_jobs[nextjob]['wait_for'].add(job_unit)
                                else:
                                    for i in range(parallelization):
                                        nextjobp = "%s_%i" % (job, i)
                                        self.get_job(nextjobp).get('wait_for', []).append(job_unit)
                                        if nextjobp in self.__pending_jobs:
                                            self.__pending_jobs[nextjobp]['wait_for'].add(job_unit)
                                nextjobs.remove(nextjob)
                self.jobs_graph[job_unit] = job_unit_conf
                self.__pending_jobs[job_unit] = {
                    'wait_for': set(wait_for),
                    'is_retry': is_retry,
                    'origin': job,
                    'wait_time': task.wait_time,
                }
            for other, oconf in self.jobs_graph.items():
                if job in oconf.get('wait_for', []):
                    oconf['wait_for'].remove(job)
                    if other in self.__pending_jobs:
                        self.__pending_jobs[other]['wait_for'].discard(job)
                    for i in range(parallelization):
                        job_unit = "%s_%i" % (job, i)
                        oconf['wait_for'].append(job_unit)
                        if other in self.__pending_jobs:
                            self.__pending_jobs[other]['wait_for'].add(job_unit)

    def add_argparser_options(self):
        super(GraphManager, self).add_argparser_options()
        self.argparser.add_argument('--jobs-graph', help='Define jobs graph_dict on command line', default='{}')
        self.argparser.add_argument('--starting-job', action='append', default=[],
                                    help='Set starting jobs. Can be given multiple times.')
        self.argparser.add_argument('--only-starting-jobs', action='store_true',
                                    help='If given, only run the starting jobs (don\'t follow on finish next jobs)')
        self.argparser.add_argument('--comment', help='Can be used for differentiate command line and avoid scheduling '
                                    'fail when a graph manager job is scheduled when another one with same option '
                                    'signature is running. Doesn\'t do anything else.')
        self.argparser.add_argument('--skip-job', default=[], action='append',
                                    help='Skip given job. Can be given multiple times. Also next jobs for the skipped'
                                         'one will be skipped.')
        self.argparser.add_argument('--resume-from-jobid', help='Resume from the given graph manager jobid')

    def parse_args(self):
        args = super(GraphManager, self).parse_args()
        self.jobs_graph = yaml.load(args.jobs_graph) or deepcopy(self.jobs_graph)

        if not self.name:
            self.argparser.error('Manager name not set.')
        return args

    def workflow_loop(self):
        logger.debug("Pending jobs: %r", self.__pending_jobs)
        logger.debug("Running jobs: %r", self.__running_jobs)
        logger.debug("Available resources: %r", self._available_resources)
        logger.debug("Acquired resources: %r", self._acquired_resources)
        self.check_running_jobs()
        if self.__pending_jobs:
            self.run_pending_jobs()
        elif not self.__running_jobs:
            return False
        return True

    def run_job(self, job, is_retry=False):
        task = self.__tasks.get(job)
        if task is not None:
            return task.run(self, is_retry)

        jobconf = self.get_job(job)
        task = self.__tasks.get(jobconf['origin'])
        if task is not None:
            idx = jobconf['index']
            return task.run(self, is_retry, index=idx)

    def _must_wait_time(self, job):
        status = self.__pending_jobs[job]
        if status['wait_time'] is not None:
            wait_time = status['wait_time'] - time() + self.__start_time[job]
            if wait_time > 0:
                logger.info("Job %s must wait %d seconds for running", job, wait_time)
                return True
        return False

    def run_pending_jobs(self):
        """Try running pending jobs.

        Normally, only jobs that have no outstanding dependencies are started.

        If all pending jobs have outstanding dependencies, try to start one job
        ignoring unknown tasks, i.e. those that are not currently pending.

        If none of the pending jobs cannot be started either way, it means
        there's a dependency cycle, in this case an error is raised.

        """

        # Normal mode: start jobs without dependencies.
        for job in sorted(self.__pending_jobs.keys()):
            if len(self.__running_jobs) >= self.max_running_jobs:
                break
            status = self.__pending_jobs[job]

            job_can_run = not status['wait_for'] and not self._must_wait_time(job) and self._try_acquire_resources(job)
            if job_can_run:
                try:
                    jobid = self.run_job(job, status['is_retry'])
                except:
                    self._release_resources(job)
                    raise
                self.__pending_jobs.pop(job)
                self.__running_jobs[job] = jobid

        if not self.__pending_jobs or self.__running_jobs or \
                any(status['wait_time'] is not None for status in self.__pending_jobs.values()):
            return

        # At this point, there are pending jobs, but none were started because
        # of dependencies, try "skip unknown deps" mode: start one job that
        # only has "unseen" dependencies to try to break the "stalemate."
        origin_job = None
        for job in sorted(self.__pending_jobs.keys()):
            if len(self.__running_jobs) >= self.max_running_jobs:
                break
            status = self.__pending_jobs[job]
            job_can_run = (
                all(w not in self.__pending_jobs for w in status['wait_for']) and
                (not origin_job or status.get('origin') == origin_job) and
                self._try_acquire_resources(job))
            origin_job = status.get('origin')
            if job_can_run:
                try:
                    jobid = self.run_job(job, status['is_retry'])
                except:
                    self._release_resources(job)
                    raise
                self.__pending_jobs.pop(job)
                self.__running_jobs[job] = jobid
            if not origin_job and self.__running_jobs:
                return

        if self.__running_jobs:
            return

        # Nothing helped, all pending jobs wait for each other somehow.
        raise RuntimeError("Job dependency cycle detected: %s" % ', '.join(
            '%s waits for %s' % (
                job, sorted(self.__pending_jobs[job]['wait_for']))
            for job in sorted(self.__pending_jobs.keys())))

    def check_running_jobs(self):
        for job, jobid in list(self.__running_jobs.items()):
            outcome = self.is_finished(jobid)
            if outcome is not None:
                logger.info('Job "%s/%s" (%s) finished', self.name, job, jobid)
                for st in self.__pending_jobs.values():
                    st['wait_for'].discard(job)
                for conf in self.jobs_graph.values():
                    if job in conf.get('wait_for', []):
                        conf['wait_for'].remove(job)
                for nextjob in self._get_next_jobs(job, outcome):
                    if nextjob == 'retry':
                        jobconf = self.get_job(job)
                        retries = jobconf.get('retries', 0)
                        if retries > 0:
                            self._add_pending_job(job, is_retry=True)
                            jobconf['retries'] -= 1
                            logger.warning('Will retry job %s (outcome: %s, number of retries left: %s)',
                                           job, outcome, jobconf['retries'])
                    elif nextjob in self.__pending_jobs:
                        logger.error('Job %s already pending', nextjob)
                    else:
                        wait_for = self.get_job(nextjob).get('wait_for', [])
                        self._add_pending_job(nextjob, wait_for)
                self._release_resources(job)
                self.__running_jobs.pop(job)
            else:
                logger.info("Job %s (%s) still running", job, jobid)

    def _try_acquire_resources(self, job):
        result = True
        task_id = self.get_job(job).get('origin', job)
        for required_resources in self.__tasks[task_id].get_required_resources(partial=True):
            for resource, req_amount in required_resources.items():
                if self._available_resources[resource] < req_amount:
                    result = False
                    break
            else:
                for resource, req_amount in required_resources.items():
                    self._available_resources[resource] -= req_amount
                    self._acquired_resources[resource].append((job, req_amount))
                return True
        return result

    def _release_resources(self, job):
        for res, acquired in self._acquired_resources.items():
            for rjob, res_amount in acquired:
                if rjob == job:
                    self._available_resources[res] += res_amount
                    self._acquired_resources[res].remove((rjob, res_amount))

    def _get_next_jobs(self, job, outcome):
        if self.args.only_starting_jobs:
            return []
        on_finish = self.get_job(job).get('on_finish', {})
        if outcome in on_finish:
            nextjobs = on_finish[outcome]
        elif outcome in self.__failed_outcomes:
            nextjobs = on_finish.get('failed', [])
        else:
            nextjobs = on_finish.get('default', [])
        return nextjobs

    @property
    def pending_jobs(self):
        return self.__pending_jobs


from pyfiglet import Figlet
import os
from shutil import copyfile
import shutil
import sqlite3
import subprocess
import winreg
import base64
import subprocess
import datetime
import socket
import ctypes

def init():
	f = Figlet(font='slant')
	print(f.renderText('Reboot'))
	print("This program is Artifact collecter\n")

def showOptions():
	print("\n==========Options==========")
	print("1) Memory Dump")
	print("2) Registry Hive")
	print("3) System Info")
	print("4) System Audit Policy")
	print("5) Group Policy")
	print("6) Event Viewer Log")
	print("7) Services Log")
	print("8) Hosts Data")
	print("9) SRUM (System Resource Utilization Monitor)")
	print("10) Environment Variables")
	print("11) Patch List")
	print("12) Process List")
	print("13) Opened Port")
	print("14) IP Config Info")
	print("15) ARP Info")
	print("16) Net BIOS")
	print("17) Opened Handle")
	print("18) Task Schedule Info")
	print("19) System Logon Info")
	print("20) UserAssist")
	print("21) AutoRun")
	print("22) Registry User")
	print("23) Internet Browser History")
	print("24) Recycle Bin")
	print("25) LNK File")
	print("26) PowerShell Log File")
	print("27) Registerd Service Info")
	print("28) Recent Activity Info")
	print("29) Prefetch")
	print("30) NTFS Artifact")
	print("777) ALL")
	print("0) exit program")

def registryHive():
	print("\n==========Registry Hive File==========")
	current_directory = os.getcwd()

	export_registry_hive(current_directory, winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE")
	export_registry_hive(current_directory, winreg.HKEY_CURRENT_USER, r"Software")

def export_registry_hive(output_directory, hive_name, hive_path):
    registry_folder = os.path.join(output_directory, "RegistryHive")
    if not os.path.exists(registry_folder):
        os.makedirs(registry_folder)

    try:
        with winreg.ConnectRegistry(None, hive_name) as hive:
            hive_output_file = os.path.join(registry_folder, f"{hive_name}_hive.txt")

            with open(hive_output_file, 'w', encoding='utf-16') as file:
                def export_subkeys(key, indent=""):
                    for i in range(winreg.QueryInfoKey(key)[0]):
                        subkey_name = winreg.EnumKey(key, i)
                        subkey_path = os.path.join(hive_path, subkey_name)

                        file.write(f"{indent}[{subkey_path}]\n")

                        with winreg.OpenKey(key, subkey_name) as subkey:
                            export_subkeys(subkey, indent + "  ")

                export_subkeys(hive)

            print(f"{hive_name} hive exported to {hive_output_file}")
    except Exception as e:
        print(f"Error: {e}")

def systemInfo():
	print("\n==========System Info File==========")
	subprocess.run(['systeminfo'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, shell=True, encoding='utf-8')

	output_directory = os.getcwd()
	system_info_folder = os.path.join(output_directory, "SystemInfo")
	if not os.path.exists(system_info_folder):
		os.makedirs(system_info_folder)

	try:
		system_info_output_file = os.path.join(system_info_folder, "system_info.txt")
		with open(system_info_output_file, 'w', encoding='utf-8') as file:
			file.write(subprocess.getoutput('systeminfo'))

		print(f"System info exported to {system_info_output_file}")
	except Exception as e:
		print(f"Error: {e}")

def systemAudit():
	print("\n==========System Audit Policy File==========")
	subprocess.run(['auditpol', '/get', '/category:*'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, shell=True, encoding='utf-8')

	output_directory = os.getcwd()
	audit_policy_folder = os.path.join(output_directory, "AuditPolicy")
	if not os.path.exists(audit_policy_folder):
		os.makedirs(audit_policy_folder)

	try:
		audit_policy_output_file = os.path.join(audit_policy_folder, "audit_policy_info.txt")
		with open(audit_policy_output_file, 'w', encoding='utf-8') as file:
			file.write(subprocess.getoutput('auditpol /get /category:*'))

		print(f"Audit policy info exported to {audit_policy_output_file}")
	except Exception as e:
		print(f"Error: {e}")

def groupPolicy():
	print("\n==========Group Policy File==========")
	subprocess.run(['gpresult', '/R'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, shell=True, encoding='utf-8')
	output_directory = os.getcwd()
	group_policy_folder = os.path.join(output_directory, "GroupPolicy")
	if not os.path.exists(group_policy_folder):
		os.makedirs(group_policy_folder)

	try:
		group_policy_output_file = os.path.join(group_policy_folder, "group_policy_info.txt")
		with open(group_policy_output_file, 'w', encoding='utf-8') as file:
			file.write(subprocess.getoutput('gpresult /R'))

		print(f"Group policy info exported to {group_policy_output_file}")
	except Exception as e:
		print(f"Error: {e}")

def eventLog():
	print("\n==========Event Viewer Log File==========")
	output_directory = os.getcwd()
	event_logs_folder = os.path.join(output_directory, "EventLogs")
	if not os.path.exists(event_logs_folder):
		os.makedirs(event_logs_folder)

	try:
        # 시스템 이벤트 로그 내용 가져오기
		system_log_output_file = os.path.join(event_logs_folder, "system_event_log.txt")
		subprocess.run(['wevtutil', 'qe', 'System', '/f:text', '/c:1'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, shell=True, encoding='utf-8')
		with open(system_log_output_file, 'w', encoding='utf-8') as file:
			file.write(subprocess.getoutput(f'wevtutil qe System /f:text /c:1'))

		print(f"System event log exported to {system_log_output_file}")

        # 응용 프로그램 이벤트 로그 내용 가져오기
		application_log_output_file = os.path.join(event_logs_folder, "application_event_log.txt")
		subprocess.run(['wevtutil', 'qe', 'Application', '/f:text', '/c:1'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, shell=True, encoding='utf-8')
		with open(application_log_output_file, 'w', encoding='utf-8') as file:
			file.write(subprocess.getoutput(f'wevtutil qe Application /f:text /c:1'))

		print(f"Application event log exported to {application_log_output_file}")
	except Exception as e:
		print(f"Error: {e}")

def serviceLog():
	print("\n==========Service Log File==========")
	output_directory = os.getcwd()
	service_log_folder = os.path.join(output_directory, "ServiceLog")
	if not os.path.exists(service_log_folder):
		os.makedirs(service_log_folder)

	try:
        # 서비스 로그를 가져와서 파일로 저장
		service_log_output_file = os.path.join(service_log_folder, "service_log.txt")
		subprocess.run(['net', 'start'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, shell=True, encoding='utf-8')
		with open(service_log_output_file, 'w', encoding='utf-8') as file:
			file.write(subprocess.getoutput('net start'))

		print(f"Service log exported to {service_log_output_file}")
	except Exception as e:
		print(f"Error: {e}")

def hostsData():
	print("\n==========Hosts Data File==========")
	output_directory = os.getcwd()
	hosts_folder = os.path.join(output_directory, "Hosts")
	if not os.path.exists(hosts_folder):
		os.makedirs(hosts_folder)

	try:
		hosts_file_path = r"C:\Windows\System32\drivers\etc\hosts"  # Hosts 파일 경로
		hosts_output_file = os.path.join(hosts_folder, "hosts.txt")

		with open(hosts_file_path, 'r', encoding='utf-8') as input_file, open(hosts_output_file, 'w', encoding='utf-8') as output_file:
			hosts_content = input_file.read()
			output_file.write(hosts_content)

		print(f"Hosts file exported to {hosts_output_file}")
	except Exception as e:
		print(f"Error: {e}")

def srum():
	print("\n==========SRUM File==========")
	output_directory = os.getcwd()
	srum_folder = os.path.join(output_directory, "SRUM")
	if not os.path.exists(srum_folder):
		os.makedirs(srum_folder)

	try:
		with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\Windows NT\CurrentVersion\SRUM\Extensions", 0, winreg.KEY_READ) as key:
			i = 0
			while True:
				try:
					subkey_name = winreg.EnumKey(key, i)
					subkey_path = os.path.join(r"SOFTWARE\Microsoft\Windows NT\CurrentVersion\SRUM\Extensions", subkey_name)

					with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, subkey_path, 0, winreg.KEY_READ) as subkey:
						j = 0
						while True:
							try:
								value_name, srum_data, _ = winreg.EnumValue(subkey, j)

                                # 각 데이터를 파일로 저장 (바이너리 모드 또는 문자열 모드로 열기)
								srum_output_file = os.path.join(srum_folder, f"{subkey_name}_{value_name}.txt")
								with open(srum_output_file, 'wb' if isinstance(srum_data, bytes) else 'w') as file:
									if isinstance(srum_data, bytes):
										file.write(srum_data)
									else:
										file.write(str(srum_data))

								print(f"SRUM data from {subkey_name}/{value_name} exported to {srum_output_file}")

								j += 1
							except OSError as e:
								break

					i += 1
				except OSError as e:
					break

			if i == 0:
				print("No SRUM data found.")
	except Exception as e:
		print(f"Error: {e}")

def environmentVar():
	print("\n==========Envionment Variable File==========")
	output_directory = os.getcwd()
	environ_var_folder = os.path.join(output_directory, "EnvironmentVar")
	if not os.path.exists(environ_var_folder):
		os.makedirs(environ_var_folder)

	output_file = os.path.join(environ_var_folder, "environment_variables.txt")
	try:
		with open(output_file, 'w') as file:
			for key, value in os.environ.items():
				file.write(f"{key} = {value}\n")

		print(f"Environment variables exported to {output_file}")
	except Exception as e:
		print(f"Error: {e}")

def patchList():
	print("\n==========Patch List File==========")
	output_directory = os.getcwd()
	patch_list_folder = os.path.join(output_directory, "PatchList")
	if not os.path.exists(patch_list_folder):
		os.makedirs(patch_list_folder)

	try:
		with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\Windows\CurrentVersion\Component Based Servicing\Packages", 0, winreg.KEY_READ | winreg.KEY_WOW64_64KEY) as key:
			i = 0
			while True:
				try:
					package_name = winreg.EnumKey(key, i)
					patch_list_file = os.path.join(patch_list_folder, f"{package_name}.txt")

					with open(patch_list_file, 'w') as file:
						file.write(package_name)

					i += 1
				except OSError as e:
					break

			if i > 0:
				print(f"Patch list exported to {patch_list_folder}")
			else:
				print("No patch list information found.")
	except Exception as e:
		print(f"Error: {e}")

def processList():
	print("\n==========Process List File==========")
	file_path = os.getcwd()
	file_path = os.path.join(file_path, "ProcessList")
	if not os.path.exists(file_path):
		os.makedirs(file_path)
	file_path = os.path.join(file_path, "process_list.txt")

	process_info_list = []

	for process in psutil.process_iter(attrs=['pid', 'name', 'username', 'memory_info']):
		info = process.info
		pid = info['pid']
		name = info['name']
		username = info['username']
		memory = info['memory_info'].rss  # Resident Set Size: 메모리 사용량

		process_info = f"PID: {pid}, Process Name: {name}, User: {username}, Memory: {memory} bytes"
		process_info_list.append(process_info)

	with open(file_path, 'w') as f:
		for process in process_info_list:
			f.write(process + '\n')

	print(f"실행 중인 프로세스 정보가 {file_path} 에 저장되었습니다.")

def openPort():
	print("\n==========Open Port File==========")
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "OpenPort")
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)
	file_path = os.path.join(destination_folder, "open_ports.txt")

	open_ports = []

	for conn in psutil.net_connections(kind='inet'):
		laddr, raddr, status, pid = conn.laddr, conn.raddr, conn.status, conn.pid
		if raddr:
			open_ports.append(f"{laddr} <--> {raddr} {status} {pid}")
		else:
			open_ports.append(f"{laddr} {status} {pid}")

	with open(file_path, 'w') as f:
		for port in open_ports:
			f.write(f"{port}\n")

	print(f"Open port information has been saved to {file_path}")

def IPConfigInfo():
	print("\n==========IP Config File==========")
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "IPConfig")
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)

	ip_info = {}
	hostname = socket.gethostname()
	local_ip = socket.gethostbyname(hostname)

	ip_info["Hostname"] = hostname
	ip_info["Local IP"] = local_ip

	interfaces = psutil.net_if_addrs()
	for interface, addrs in interfaces.items():
		ip_info[interface] = []
		for addr in addrs:
			ip_info[interface].append(str(addr))

	filename = os.path.join(destination_folder, f"ip_info_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")

	with open(filename, 'w') as f:
		for key, value in ip_info.items():
			f.write(f"{key}: {value}\n")

	print(f"IP 설정 정보가 {filename} 에 저장되었습니다.")

def arpInfo():
	print("\n==========ARP Info File==========")
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "ARPInfo")
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)

	try:
        # 'arp -a' 명령을 실행하고 결과를 반환받습니다.
		arp_output = subprocess.check_output("arp -a", shell=True, stderr=subprocess.STDOUT, text=True)
	except subprocess.CalledProcessError as e:
		arp_output = f"An error occurred while trying to fetch ARP info: {str(e)}"
    # 파일 이름에 현재 시간을 추가하여 고유하게 만듭니다.
	filename = os.path.join(destination_folder, f"arp_info_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")

    # ARP 정보를 .txt 파일에 저장합니다.
	with open(filename, 'w') as f:
		f.write(arp_output)

	print(f"ARP 정보가 {filename} 에 저장되었습니다.")

def netBIOS():
	print("\n==========Net BIOS File==========")
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "NetBIOS")
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)

	try:
        # 'nbtstat -n' 명령을 실행하고 그 결과를 반환합니다.
		netbios_output = subprocess.check_output("nbtstat -n", shell=True, stderr=subprocess.STDOUT, text=True)
	except subprocess.CalledProcessError as e:
		netbios_output = f"An error occurred while trying to fetch NetBIOS info: {str(e)}"

    # 파일 이름에 현재 시간을 추가해 고유한 파일을 생성합니다.
	filename = os.path.join(destination_folder, f"netbios_info_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")

    # NetBIOS 정보를 .txt 파일에 저장합니다.
	with open(filename, 'w') as f:
		f.write(netbios_output)

	print(f"NetBIOS 정보가 {filename} 에 저장되었습니다.")

def openedHandle():
	print("\n==========Opened Handle File==========")
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "OpenHandle")
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)

	processes = []

	for proc in psutil.process_iter(['pid', 'name', 'open_files']):
		processes.append(proc.info)

    # 파일 이름에 현재 시간을 추가해 고유한 파일을 생성합니다.
	filename = os.path.join(destination_folder, f"handle_info_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")

    # 핸들 정보를 .txt 파일에 저장합니다.
	with open(filename, 'w') as f:
		for proc in processes:
			f.write(f"PID: {proc['pid']}, Name: {proc['name']}\n")
			if proc['open_files']:
				for open_file in proc['open_files']:
					f.write(f"\t{open_file}\n")

	print(f"열려있는 핸들 정보가 {filename} 에 저장되었습니다.")

def taskSchedule():
	print("\n==========Task Schedule File==========")
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "JobSchedule")
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)

	try:
        # 'schtasks' 명령을 실행하고 그 결과를 반환합니다.
		output = subprocess.check_output("schtasks /query /fo LIST", shell=True, stderr=subprocess.STDOUT, text=True)
	except subprocess.CalledProcessError as e:
		output = f"An error occurred while trying to fetch task scheduler info: {str(e)}"

    # 파일 이름에 현재 시간을 추가해 고유한 파일을 생성합니다.
	filename = os.path.join(destination_folder, f"task_scheduler_info_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")

    # 작업 스케줄러 정보를 .txt 파일에 저장합니다.
	with open(filename, 'w') as f:
		f.write(output)

	print(f"작업 스케줄 정보가 {filename} 에 저장되었습니다.")

def systemLogon():
	print("\n==========System Logon File==========")
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "SystemLogon")
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)
	file_path = os.path.join(destination_folder, "logon_history.txt")

	query = "wevtutil qe Security /q:\"*[System[Provider[@Name='Microsoft-Windows-Security-Auditing'] and (EventID=4624)]]\" /c:1 /rd:true /f:text"
	result = subprocess.run(query, capture_output=True, text=True, shell=True)

	with open(file_path, 'w') as f:
		f.write(result.stdout)

	print(f"시스템 로그온 정보가 {file_path} 에 저장되었습니다.")

def memoryDump():
	print("\n==========MemoryDump File==========")
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "MemoryDump")

	# "MemoryDump" 폴더 생성 후 저장
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)

	# 메모리 덤프 파일을 저장할 디렉토리 경로 지정
	dump_directory = destination_folder

	# psexec.exe 경로 (psexec을 다운로드한 경로로 설정)
	psexec_path = os.getcwd()
	psexec_path = os.path.join(psexec_path, "PSTools")
	psexec_path = os.path.join(psexec_path, "PsExec.exe")

	# procdump.exe 경로 (procdump를 설치한 경로로 설정)
	procdump_path = os.getcwd()
	procdump_path = os.path.join(procdump_path, "Procdump")
	procdump_path = os.path.join(procdump_path, "procdump.exe")

	# 현재 실행 중인 모든 프로세스 가져오기
	running_processes = list(psutil.process_iter(attrs=['pid', 'name']))

	# 실행 중인 모든 프로세스에 대해 메모리 덤프 생성
	for process in running_processes:
		process_name = process.info['name']
		process_pid = process.info['pid']
		dump_file_path = os.path.join(dump_directory, f"{process_name}_{process_pid}_dump.dmp")

		# procdump을 사용하여 메모리 덤프 실행
		cmd = [procdump_path, f"-ma", process_name, dump_file_path]

		try:
			subprocess.run(cmd, check=True)
			print(f"{process_name} 프로세스의 메모리 덤프 생성 완료:", dump_file_path)
		except subprocess.CalledProcessError as e:
			print(f"{process_name} 프로세스의 메모리 덤프 생성 중 오류 발생:", e)

	print("모든 프로세스의 메모리 덤프 생성 완료")

def userAssist():
	print("\n==========UserAssist File==========")

	keys_to_copy = [
		# 실행파일 기록 subkey
        r"Software\Microsoft\Windows\CurrentVersion\Explorer\UserAssist\{CEBFF5CD-ACE2-4F4F-9178-9926F41749EA}\Count",
		# 바로가기 실행 기록 subkey
        r"Software\Microsoft\Windows\CurrentVersion\Explorer\UserAssist\{F4E57C4B-2036-45F0-A9AB-443BCFE33D9F}\Count"
    ]

	current_directory = os.getcwd()
	destination_folder = os.path.join(current_directory, "UserAssist")

	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)

	try:
		with winreg.OpenKey(winreg.HKEY_CURRENT_USER, keys_to_copy[0], 0, winreg.KEY_READ | winreg.KEY_WOW64_64KEY) as key:
			userAssistRegistry(key, destination_folder, 1)
		with winreg.OpenKey(winreg.HKEY_CURRENT_USER, keys_to_copy[1], 0, winreg.KEY_READ | winreg.KEY_WOW64_64KEY) as key:
			userAssistRegistry(key, destination_folder, 2)
	except Exception as e:
		print(f"Error: {e}")

def sanitize_filename(filename):
    # 파일 이름에 유효하지 않은 문자를 대체
    invalid_chars = ['\\', '/', ':', '*', '?', '"', '<', '>', '|', '.', '&', '!', '@', '#', '$', '%', '^', '(', ')', '[', ']', '{', '}', '+', '=', ',', ';', '`', "'", '~', ' ']
    for char in invalid_chars:
        filename = filename.replace(char, '_')

    # 파일 이름을 최대 100자로 제한
    return filename[:100]

def userAssistRegistry(subkey, output_directory, num):
	try:
		if not os.path.exists(output_directory):
			os.makedirs(output_directory)

		if num == 1:
			subkey_name = str(winreg.QueryInfoKey(subkey)[0])
			subkey_path = os.path.join(output_directory, "CEBFF5CD-ACE2-4F4F-9178-9926F41749EA")
		else:
			subkey_name = str(winreg.QueryInfoKey(subkey)[0])
			subkey_path = os.path.join(output_directory, "F4E57C4B-2036-45F0-A9AB-443BCFE33D9F")

		if not os.path.exists(subkey_path):
			os.makedirs(subkey_path)

		i = 0
		while True:
			try:
				value_name, value_data, _ = winreg.EnumValue(subkey, i)
				value_output_file = os.path.join(subkey_path, f"{sanitize_filename(value_name)}.txt")

                # 중간 디렉터리가 없는 경우 생성
				value_output_dir = os.path.dirname(value_output_file)
				if not os.path.exists(value_output_dir):
					os.makedirs(value_output_dir)

				with open(value_output_file, 'w') as file:
					file.write(str(value_data))  # value_data를 문자열로 변환하여 쓰기
				i += 1
			except OSError as e:
				print(f"Error while processing subkey {subkey_name}: {e}")
				break

		print(f"Data from subkey {subkey_name} has been saved to {subkey_path}")
	except Exception as e:
		print(f"Error: {e}")

def autoRun():
	print("\n==========AutoRun File==========")
	output_directory = os.getcwd()
	autorun_folder = os.path.join(output_directory, "AutoRun")
	if not os.path.exists(autorun_folder):
		os.makedirs(autorun_folder)

	try:
		with winreg.OpenKey(winreg.HKEY_CURRENT_USER, r"Software\Microsoft\Windows\CurrentVersion\Run", 0, winreg.KEY_READ | winreg.KEY_WOW64_64KEY) as key:
			values = {}
			i = 0
			while True:
				try:
					value_name, value_data, _ = winreg.EnumValue(key, i)
					values[value_name] = value_data
					i += 1
				except OSError as e:
					break

			if values:
				output_file = os.path.join(autorun_folder, "autorun_values.txt")
				with open(output_file, 'w') as file:
					for name, data in values.items():
						file.write(f"{name} = {data}\n")

				print(f"AutoRun values exported to {output_file}")
			else:
				print("No AutoRun values found.")
	except Exception as e:
		print(f"Error: {e}")

def registryUser():
	print("\n==========Registry User File==========")
	output_directory = os.getcwd()
	registry_folder = os.path.join(output_directory, "AllUserRegistry")
	if not os.path.exists(registry_folder):
		os.makedirs(registry_folder)

	try:
		with winreg.ConnectRegistry(None, winreg.HKEY_USERS) as users_hive:
			with winreg.OpenKey(users_hive, None) as users_key:
				num_subkeys = winreg.QueryInfoKey(users_key)[0]

				for i in range(num_subkeys):
					user_sid = winreg.EnumKey(users_key, i)
					user_hive_path = f"{user_sid}\\Software"  # 사용자 하이브 경로

					with winreg.ConnectRegistry(None, winreg.HKEY_USERS) as user_hive:
						with winreg.CreateKey(user_hive, user_hive_path) as user_key:
                            # 사용자 레지스트리 정보를 파일로 저장
							user_output_file = os.path.join(registry_folder, f"{user_sid}_registry.txt")
							with open(user_output_file, 'w', encoding='utf-16') as file:
								def export_subkeys(key, indent=""):
									for j in range(winreg.QueryInfoKey(key)[0]):
										subkey_name = winreg.EnumKey(key, j)
										subkey_path = os.path.join(user_hive_path, subkey_name)

										file.write(f"{indent}[{subkey_path}]\n")

										with winreg.OpenKey(key, subkey_name) as subkey:
											export_subkeys(subkey, indent + "  ")

								export_subkeys(user_key)

							print(f"{user_sid} registry exported to {user_output_file}")

	except Exception as e:
		print(f"Error: {e}")

def export_registry_key(reg_file, hkey, key_path):
    try:
        key = winreg.OpenKey(hkey, key_path)
        for i in range(winreg.QueryInfoKey(key)[0]):
            sub_key_name = winreg.EnumKey(key, i)
            sub_key_path = os.path.join(key_path, sub_key_name)
            reg_file.write(f'\n[{sub_key_path}]\n')

            export_registry_key(reg_file, hkey, sub_key_path)

        for i in range(winreg.QueryInfoKey(key)[1]):
            value_name, value_data, value_type = winreg.EnumValue(key, i)
            if value_type == winreg.REG_SZ:
                reg_file.write(f'"{value_name}"="{value_data}"\n')
            elif value_type == winreg.REG_DWORD:
                reg_file.write(f'"{value_name}"=dword:{value_data:08X}\n')
            # 여러 다른 레지스트리 값 유형에 대한 처리 추가 가능

    except Exception as e:
        pass

def save_browser_history():
	print("\n==========Browser History File==========")
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "BrowserHistory")

	# "MemoryDump" 폴더 생성 후 저장
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)

	# 크롬 브라우저의 기록 데이터베이스 경로
	chrome_history_path = os.path.expanduser('~') + '\\AppData\\Local\\Google\\Chrome\\User Data\\Default\\History'
	# 복사할 임시 파일 경로
	temp_db_path = os.path.join(destination_folder, 'temp_history')

	# 기록 데이터베이스 파일 복사
	copyfile(chrome_history_path, temp_db_path)

	# 데이터베이스 연결
	connection = sqlite3.connect(temp_db_path)
	cursor = connection.cursor()

	# 방문 기록 가져오기
	cursor.execute("SELECT title, url, last_visit_time FROM urls")
	history = cursor.fetchall()

	# 파일로 저장
	history_file_path = os.path.join(destination_folder, 'browser_history.txt')
	with open(history_file_path, 'w', encoding='utf-8') as file:
		for item in history:
			title, url, timestamp = item
			time_formatted = str(timestamp)
			file.write(f"Title: {title}\nURL: {url}\nTimestamp: {time_formatted}\n\n")

	# 연결 종료 및 임시 데이터베이스 파일 삭제
	cursor.close()
	connection.close()
	os.remove(temp_db_path)

	print("브라우저 기록이 'Browser' 폴더에 저장되었습니다.")

def recycleBin():
	print("\n==========Recycle Bin File==========")

	# 'RecycleBin' 폴더 생성
	destination_folder = os.path.join(os.getcwd(), 'RecycleBin')
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)

	try:
	    # Powershell 스크립트 실행
		script_path = os.path.join(os.path.dirname(__file__), 'copy_recycle.ps1')
		command = ["powershell", "-ExecutionPolicy", "Bypass", "-File", script_path]
		subprocess.run(command, shell=True, check=True)

		print("휴지통 내용이 'RecycleBin' 폴더에 복사되었습니다.")
	except Exception as e:
	    print(f"오류가 발생했습니다: {e}")

def lnkFile():
	print("\n==========LNK File==========")
	# LNK 폴더 생성 후 저장
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "LNK")
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)

	# .lnk파일 위치
	source_folder = os.path.expanduser("~") + "\\AppData\\Roaming\\Microsoft\\Windows\\Recent"

	# file copy
	for root, dirs, files in os.walk(source_folder):
		for file in files:
			source_file_path = os.path.join(root, file)
			destination_file_path = os.path.join(destination_folder, file)

			try:
				shutil.copy2(source_file_path, destination_file_path)
				print(f"복사 완료: {file}")
			except Exception as e:
				print(f"복사 실패: {file}, 오류: {e}")

def PowerShellLogFile():
	print("\n==========PowerShell Log File==========")
	# PowerShellLog 폴더 생성 후 저장
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "PowerShellLog")
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)

	# .lnk파일 위치
	source_folder = os.path.expanduser("~") + "\\AppData\\Roaming\\Microsoft\\Windows\\PowerShell\\PSReadLine"

	# file copy
	for root, dirs, files in os.walk(source_folder):
		for file in files:
			source_file_path = os.path.join(root, file)
			destination_file_path = os.path.join(destination_folder, file)
			try:
				shutil.copy2(source_file_path, destination_file_path)
				print(f"복사 완료: {file}")
			except Exception as e:
				print(f"복사 실패: {file}, 오류: {e}")

def registeredService():
	print("\n==========Registered Service File==========")
	# PowerShellLog 폴더 생성 후 저장
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "RegisteredService")
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)
	filename = os.path.join(destination_folder, f"service_info_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")

	try:
        # 'sc query' 명령을 실행하여 서비스 정보를 가져옵니다.
		output = subprocess.check_output("sc query", shell=True, text=True, stderr=subprocess.STDOUT)
	except subprocess.CalledProcessError as e:
		output = f"An error occurred while trying to fetch service info: {str(e)}"

	# 서비스 정보를 .txt 파일에 저장합니다.
	with open(filename, 'w') as f:
		f.write(output)

	print(f"등록된 서비스 정보가 {filename} 에 저장되었습니다.")

def recentActivity():
	print("\n==========Recent Activity File==========")
	# PowerShellLog 폴더 생성 후 저장
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "RecentActivity")
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)
    # Recent Items 폴더의 경로를 가져옵니다.
	recent_folder = os.path.join(os.environ['USERPROFILE'], r'AppData\Roaming\Microsoft\Windows\Recent')

	recent_items = []

    # 폴더 내의 모든 파일과 폴더를 나열합니다.
	for item in os.listdir(recent_folder):
		item_path = os.path.join(recent_folder, item)
		item_stat = os.stat(item_path)

        # 파일의 마지막 액세스 시간을 가져옵니다.
		last_access_time = datetime.datetime.fromtimestamp(item_stat.st_atime).strftime('%Y-%m-%d %H:%M:%S')
		recent_items.append(f"{item}: Last accessed at {last_access_time}")

	# 파일 이름에 현재 시간을 추가해 고유한 파일을 생성합니다.
	filename = os.path.join(destination_folder, f"recent_activity_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")

    # 최근 활동 정보를 .txt 파일에 저장합니다.
	with open(filename, 'w') as f:
		for info in recent_items:
			f.write(f"{info}\n")
	print(f"최근 활동 정보가 {filename} 에 저장되었습니다.")

def prefetch():
	print("\n==========Prefetch File==========")
	# PowerShellLog 폴더 생성 후 저장
	current_directory = os.path.dirname(os.path.abspath(__file__))
	destination_folder = os.path.join(current_directory, "Prefetch")
	if not os.path.exists(destination_folder):
		os.makedirs(destination_folder)

	# Prefetch 폴더 경로
	prefetch_path = r"C:\Windows\Prefetch"
	# Prefetch 폴더에서 .pf 파일 목록 수집
	try:
		prefetch_files = [f for f in os.listdir(prefetch_path) if f.endswith('.pf')]
	except PermissionError:
		print("관리자 권한이 필요합니다.")
		exit()

	for prefetch_file in prefetch_files:
		source_file = os.path.join(prefetch_path, prefetch_file)
		destination_file = os.path.join(destination_folder, prefetch_file)
		shutil.copy(source_file, destination_file)

	print(f"prefetch 파일이 {destination_file} 에 저장되었습니다.")

def NTFS():
	print("\n==========NTFS Artifact File==========")
	script_dir = os.path.dirname(os.path.abspath(__file__))
	ntfs_folder = os.path.join(script_dir, 'NTFS')
	os.makedirs(ntfs_folder, exist_ok=True)  # 'NTFS' 폴더 생성 (이미 존재하면 무시)

    # 복사할 NTFS 시스템 파일
	ntfs_files = ["$MFT", "$LogFile", "$Extend\\$UsnJrnl:$J"]

	for ntfs_file in ntfs_files:
		source_file = "C:\\" + ntfs_file
		destination_file_name = ntfs_file.replace('$', '').replace(':', '').replace('\\', '_') + ".txt"
		destination_file = os.path.join(ntfs_folder, destination_file_name)

		command = "fsutil file queryextents {} > {}".format(source_file, destination_file)

		try:
			subprocess.run(command, shell=True, check=True, stderr=subprocess.PIPE)
			print("{}를 성공적으로 복사했습니다.".format(source_file))
		except subprocess.CalledProcessError as e:
			print("{}를 복사하는 데 실패했습니다: {}".format(source_file, e.stderr.decode('utf-8')))

	print(f"NTFS 파일이 {destination_file} 에 저장되었습니다.")

def main():
	init()
	while(True):
		showOptions()
		options = input("[Select Option] : ")

		if options == "0":
			print("Good Bye!")
			exit()
		elif options == "1":
			memoryDump()
		elif options == "2":
			registryHive()
		elif options == "3":
			systemInfo()
		elif options == "4":
			systemAudit()
		elif options == "5":
			groupPolicy()
		elif options == "6":
			eventLog()
		elif options == "7":
			serviceLog()
		elif options == "8":
			hostsData()
		elif options == "9":
			srum()
		elif options == "10":
			environmentVar()
		elif options == "11":
			patchList()
		elif options == "12":
			processList()
		elif options == "13":
			openPort()
		elif options == "14":
			IPConfigInfo()
		elif options == "15":
			arpInfo()
		elif options == "16":
			netBIOS()
		elif options == "17":
			openedHandle()
		elif options == "18":
			taskSchedule()
		elif options == "19":
			systemLogon()
		elif options == "20":
			userAssist()
		elif options == "21":
			autoRun()
		elif options == "22":
			registryUser()
		elif options == "23":
			save_browser_history()
		elif options == "24":
			recycleBin()
		elif options == "25":
			lnkFile()
		elif options == "26":
			PowerShellLogFile()
		elif options == "27":
			registeredService()
		elif options == "28":
			recentActivity()
		elif options == "29":
			prefetch()
		elif options == "30":
			NTFS()
		elif options == "777":
			memoryDump()
			registryHive()
			systemInfo()
			systemAudit()
			groupPolicy()
			eventLog()
			serviceLog()
			hostsData()
			srum()
			environmentVar()
			patchList()
			processList()
			openPort()
			IPConfigInfo()
			arpInfo()
			netBIOS()
			openedHandle()
			taskSchedule()
			systemLogon()
			userAssist()
			autoRun()
			registryUser()
			save_browser_history()
			recycleBin()
			lnkFile()
			PowerShellLogFile()
			registeredService()
			recentActivity()
			prefetch()
			NTFS()
		else :
			print("\nPlease input correct options!")
			pass

if __name__ == "__main__":
	main()


"""Test the concurrency module."""

from typing import Any

import pytest

from rtasr.concurrency import ConcurrencyHandler, ConcurrencyToken

class TestConcurrencyToken:
    """Test the ConcurrencyToken class."""

    @pytest.mark.parametrize("value", [None, "string", True, False, [], {}])
    def test_concurrency_token_invalid(self, value: Any) -> None:
        """Test the concurrency token with invalid values."""
        with pytest.raises(TypeError):
            ConcurrencyToken(value)

    @pytest.mark.parametrize("value", [0, 1, 2, 1.0, 2.0])
    def test_concurrency_token_valid(self, value: int) -> None:
        """Test the concurrency token."""
        token = ConcurrencyToken(value=value)

        assert token.value == value

class TestConcurrencyHandler:
    """Test the ConcurrencyHandler class."""

    @pytest.mark.asyncio
    async def test_concurrency_handler_with_limit(self) -> None:
        """Test the concurrency handler with a limit."""
        limit = 5
        concurrency_handler = ConcurrencyHandler(limit)

        assert hasattr(concurrency_handler, "limit")
        assert hasattr(concurrency_handler, "queue")
        assert concurrency_handler.limit == limit
        assert concurrency_handler.queue.maxsize == limit

        tokens = []
        for _ in range(limit):
            token = await concurrency_handler.get()
            assert isinstance(token, ConcurrencyToken)
            tokens.append(token)

        # Queue should be empty now
        assert concurrency_handler.queue.empty()

        for token in tokens:
            concurrency_handler.put(token)

        # Queue should be refilled
        for _ in range(limit):
            token = await concurrency_handler.get()
            assert isinstance(token, ConcurrencyToken)

    @pytest.mark.asyncio
    async def test_concurrency_handler_without_limit(self) -> None:
        """Test the concurrency handler without a limit."""
        concurrency_handler = ConcurrencyHandler(limit=None)

        token = await concurrency_handler.get()
        assert token is None


import csv
import sqlite3
import numpy as np
import pandas as pd
from nltk.stem.porter import *
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
import string
from PIL import Image
from os import path
import matplotlib.pyplot as plt
import matplotlib as mpl
import pickle

from sklearn.svm import LinearSVC
from sklearn.feature_extraction.text import TfidfVectorizer

from wordcloud import WordCloud, STOPWORDS
from collections import Counter

from nltk.stem import SnowballStemmer
from nltk.stem import WordNetLemmatizer
stemmer = SnowballStemmer('english')
wordnet_lemmatizer = WordNetLemmatizer()

#############################In/out############################
def open_file(csvfile):
    reader = pd.read_csv(csvfile)
    return reader

def output_file(df,string):
    df.to_csv(string, index = False)

#############################Word Cloud & Feature Extraction############################

def text_process(data):
    '''
    Takes in a df in format of [text,stars] performs the following:
    1. Lower capital letters
    2. Remove all punctuation
    3. Remove all stopwords
    4. Reduce words to their word stem
    5. Return a list of words

    '''
    for i in range(len(data)):
        line = data[i]
        line = line.lower()  # lower case
        translation = str.maketrans("", "", string.punctuation);
        line = line.translate(translation)
        split = word_tokenize(line)
        # filter out any tokens not containing letters (e.g., numeric tokens, raw punctuation)
        filtered = []
        for token in split:
            if re.search('[a-zA-Z]', token):
                filtered.append(token)
        word = [i for i in filtered if i not in stopwords.words('english')]

        d = [stemmer.stem(word) for word in word]
        d = [wordnet_lemmatizer.lemmatize(word) for word in d]
        data[i] = d
    return data

def top_words(business_id,review_ml ):
    train = review_ml[review_ml['business_id'] == business_id][review_ml['True(1)/Deceptive(0)'] == 'True']
    text = list(train['Review'])  # text
    text = text_process(text)
    text = sum(text, [])

    counts = Counter(text)
    wordcloud = WordCloud(
        background_color='white',
        max_words=100,
        max_font_size=50,
        min_font_size=10,
        random_state=40,

    ).fit_words(counts)
    fig = plt.figure(1)
    plt.imshow(wordcloud)
    plt.axis('off')  # remove axis
    plt.show()

def change_label(x):
    for i in range(len(x)):
        if x[i] >= 3.0:  # good review: stars >=3.0
            x[i] = 1
        else:  # bad review: stars 3.0
            x[i] = 0
    return x

def bigram(business_id, review_ml):
    # only use true review

    train0 = review_ml[review_ml['business_id'] == business_id]
    train = train0[train0['True(1)/Deceptive(0)'] == 'True']
    #print(train.head())
    #train_data = list(train['Review'])  # text
    label = list(train['Stars'])  # ratings
    #print(label)
    train_label = change_label(label)
    #print(train_label)

    # TfidfVectorizer Transform
    transformer = TfidfVectorizer(stop_words='english',
                                  ngram_range=(2, 2))  # "ignore terms that appear in less than 1% of the documents".
    #print(transformer)

    cvectorizer = transformer.fit(train['Review'])
    #print(cvectorizer)
    transformed = cvectorizer.transform(train['Review'])
    #print(transformed)

    # SVM regression
    clf = LinearSVC()
    clf.fit(transformed, train_label)
    coefficients = clf.coef_.ravel()
    #print(coefficients)
    pos_coefficients = np.argsort(coefficients)[-10:]
    neg_coefficients = np.argsort(coefficients)[:10]
    combine = np.hstack([neg_coefficients, pos_coefficients])
    #print("combine:, ",combine)
    #print("coefficients[combine]: ", coefficients[combine])

    plt.figure(figsize=(7, 4))
    #print("fisnish 1")

    colors = ['red' if i < 0 else 'blue' for i in coefficients[combine]]
    #print("finish 2")

    plt.bar(np.arange(len(coefficients[combine])), coefficients[combine], color=colors)
    #print("finish 3")

    feature_names = np.array(cvectorizer.get_feature_names())
    #print("finish 4")
    plt.title('why the restaurant is rated as bad or good ', fontsize=15)
    #print("finish 5")

    plt.xticks(np.arange(0, 2 * 10), feature_names[combine], rotation=40, ha='right')
    #print("finish 6")

    plt.show()
    #print("finish 7")

#############################helper function#############################
def load_database_data(c, zipcode, business_name_input):
    c.execute('''
            SELECT b_id,r.review, r.r_stars
            FROM business, review_fact_table r
            WHERE postal_code = ? AND name = ? AND r.business_id = b_id''', (zipcode, business_name_input,))
    dataframe = pd.DataFrame(data=c.fetchall(), columns=['business_id', 'review', 'rating'])
    return dataframe
def select_data(c, zipcode, business_name):
    c.execute('''
                SELECT DISTINCT(b_id)
                FROM business, review_fact_table r
                WHERE postal_code = ? AND name = ? AND r.business_id = b_id''', (zipcode, business_name,))
    single_df = pd.DataFrame(data=c.fetchall(), columns=['business_id'])
    return single_df['business_id'][0]

def fake_ratio(predict, single):
    # Load fake results
    predicted_fake = predict
    # reviews that has only that business id
    reviews = predicted_fake[predicted_fake['business_id'] == single]
    n = reviews.shape[0]
    # print(n)
    fake = reviews.groupby('True(1)/Deceptive(0)').count()['Review'][0]
    # print(fake)
    fake_percentage = fake / n
    # print(fake_percentage)
    return fake_percentage

##############################main######################################
def main():
    #open states and income raw data
    zipcode = input("zipcode:")
    business_name = input("restaurant name:")
    print(zipcode,business_name)

    conn = sqlite3.connect('yelp.db')
    c = conn.cursor()
    predicted_fake = open_file('data/predicted_review.csv')
    # find the business id
    single = select_data(c, zipcode, business_name)
    # print(single)
    fake_review_ratio = fake_ratio(predicted_fake,single)
    print(fake_review_ratio)
    #top_words(single, predicted_fake)
    bigram(single, predicted_fake)

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


#!/usr/bin/env python3

# Get sequences from NCBI.
# To be called from Snakefile.
# Usage: python windows.py <infile> <outfile> <email> <window_size>

import os
import sys

from Bio import Entrez
from Bio import SeqIO
import pandas as pd

def main():
    snpfile = sys.argv[1]
    outfile = sys.argv[2]
    email = sys.argv[3]
    window_size = sys.argv[4]

    interest = pd.read_csv(
        snpfile,
        header=0,
        dtype={"chrom": str, "center": np.int64},
    )
    interest.columns = interest.columns.str.lower().str.replace(" ", "_")
    interest[["chrom", "center"]] = (
        interest.iloc[:, 0]
        .str.replace("_[A-Z]{3}?", "")
        .str.replace(" ", "")
        .str.split("_", expand=True)
    )
    interest.assign(
        start=lambda center: center.to_numeric() - window_size,
        end=lambda center: center.to_numeric() + window_size
    )
    interest.index.rename("Index", inplace=True)
    summary = pd.read_csv("Data/dmelSummary.csv")
    summary.index.rename("Index", inplace=True)
    summary.columns = summary.columns.str.lower().str.replace(" ", "_")

    seqs = []
    Entrez.email = email
    for index, row in interest.iterrows():
        with Entrez.efetch(
            db="nucleotide",
            id=summary.loc[summary["name"] == row["chrom"], "refseq"].iat[0],
            rettype="fasta",
            strand=1,
            seq_start=row["start"],
            seq_stop=row["end"],
        ) as handle:
            seqs.append(SeqIO.read(handle, "fasta"))
    SeqIO.write(seqs, outfile, "fasta")

if __name__ == "__main__":
    main()