from map_thicket_agc import get_logger
import os
import warnings
from collections import OrderedDict
import numpy as np
from sklearn import linear_model
import rasterio
from rasterio.features import sieve
from rasterio.windows import Window
from rasterio.mask import raster_geometry_mask
from rasterio import fill
import geopandas as gpd, pandas as pd

logger = get_logger(__name__)

def nanentropy(x, axis: tuple=None):
    """
    Find entropy ignoring NaNs

    Parameters
    ----------
    x : numpy.array_like
        data in 3 or fewer dimensions
    axis : int
        specify dimension(s) of x along which to find entropy, default=None in which case all dimensions are used

    Returns
    -------
    Entropy of x long specified dimensions
    """
    if not axis is None:
        if (len(axis) > x.ndim) or (np.any(np.array(axis) > x.ndim)):
            raise Exception('len(axis) > x.ndim) or (np.any(axis > x.ndim))')
        elif len(axis) == x.ndim:
            axis = None

    if axis is None:    # find entropy along all dimensions of x
        nbins = 100
        # quantise x into nbins bins
        x = x[~np.isnan(x)]
        if len(x) < 10:
            return 0.
        x = x - x.min()
        x = np.int32(np.round(np.float32(nbins * x) / x.max()))

        value, counts = np.unique(x, return_counts=True)
        p = np.array([count/float(x.size) for count in counts])
        return -np.sum(p * np.log2(p))
    else:        # slice 3D array
        along_axis = np.setdiff1d(range(0, x.ndim), axis)[0]
        e = np.zeros((x.shape[along_axis]))
        for axis_slice in range(x.shape[along_axis]):
            if along_axis == 0:
                xa = x[axis_slice,:,:]
            elif along_axis == 1:
                xa = x[:,axis_slice,:]
            elif along_axis == 2:
                xa = x[:,:,axis_slice]
            else:
                raise Exception('along_axis > 2 or < 0')
            e[axis_slice] = nanentropy(xa, axis=None)
        return e

class PatchFeatureExtractor:
    def __init__(self, num_bands=1, apply_rolling_window=False, rolling_window_xsize=None, rolling_window_xstep=None):
        """
        Virtual class for extracting features from image patches, that hides extraction details from user.
        Optional application of rolling window (x direction only)

        Parameters
        ----------
        apply_rolling_window : bool
            include rolling window in feature functions
        rolling_window_xsize : int
            x size of rolling window in pixels (optional)
        rolling_window_xstep : int
            Number of pixels to step in x direction (optional)
        """
        self.fn_dict = {}
        if apply_rolling_window:
            if rolling_window_xsize is None or rolling_window_xstep is None:
                raise Exception("rolling_window_xsize and rolling_window_xstep must be specified")
        self._num_bands = num_bands
        self._apply_rolling_window = apply_rolling_window
        self._rolling_window_xsize = rolling_window_xsize
        self._rolling_window_xstep = rolling_window_xstep
        # self._pan_bands = None
        self._generate_fn_dict()

    def _rolling_window_view(self, x):
        """
        Return a 3D strided view of 2D array to allow fast rolling window operations.
        Rolling windows are stacked along the third dimension.  No data copying is involved.

        Parameters
        ----------
        x : numpy.array_like
            array to return view of

        Returns
        -------
        3D rolling window view of x
        """

        shape = x.shape[:-1] + (self._rolling_window_xsize, int(1 + (x.shape[-1] - self._rolling_window_xsize) / self._rolling_window_xstep))
        strides = x.strides[:-1] + (x.strides[-1], self._rolling_window_xstep * x.strides[-1])
        return np.lib.stride_tricks.as_strided(x, shape=shape, strides=strides, writeable=False)

    def _generate_fn_dict(self):
        """
        Generate feature extraction dictionary self.fn_dict with values= pointer to feature extraction functions, and
        keys= descriptive feature strings.  Suitable for use with modelling from patches or mapping a whole image i.e.
        ImageFeatureExtractor and ImageMapper.
        """
        raise NotImplementedError()

    def extract_features(self, im_patch, mask=None, fn_keys=None):
        """
        Virtual method to extract features from image patch with optional mask
        """
        raise NotImplementedError()

class MsPatchFeatureExtractor(PatchFeatureExtractor):
    # TODO: PatchFeatureExtractor and MsPatchFeatureExtractor may benefit from refactoring that gets around the
    #  necessity for apply_rolling_window which is confusing e.g. rolling_window_xsize and rolling_window_xstep are
    #  always passed to extract_features(), and not to __init__. extract_features() then passes to feature functions in
    #  fn_dict.  If there is no rolling window then rolling_window_xsize and rolling_window_xstep are set to the patch x
    #  size.  This would avoid the need for upfront specification of rolling window.

    def __init__(self, num_bands=9, apply_rolling_window=False, rolling_window_xsize=None, rolling_window_xstep=None):
        """
        PatchFeatureExtractor for extracting band ratio, vegetation index and texture features from multi-spectral image patches.

        Parameters
        ----------
        num_bands : int
            number of multi-spectral bands
        apply_rolling_window : bool
            include rolling window in feature functions
        rolling_window_xsize : int
            x size of rolling window in pixels (optional).  If None, no rolling window applied
        rolling_window_xstep : int
            Number of pixels to step in x direction (optional).
        """
        self._pan_bands, self._band_dict = self.get_band_info(num_bands)
        PatchFeatureExtractor.__init__(self, num_bands=num_bands, apply_rolling_window=apply_rolling_window, rolling_window_xsize=rolling_window_xsize,
                                       rolling_window_xstep=rolling_window_xstep)

    @staticmethod
    def get_band_info(num_bands=9):
        """
        Get an array of band numbers to sum to generate panchromatic, and a dict of band labels and numbers for
        assumed multi-spectral images (Worldview3 or NGI aerial)

        Parameters
        ----------
        num_bands : int
            number of multi-spectral bands in the image

        Returns
        -------
        (pan_bands, band_dict), where
            pan_bands: array of band numbers to sum to create pan
            band_dict: dictions of {band name: band number} for multi-spectral bands
        """
        if num_bands == 8:  # assume WV3
            # pan_bands = np.array([1, 2, 3, 4, 5, 6])
            pan_bands = [1, 2, 4, 5]
            band_dict = OrderedDict([('C', 0), ('B', 1), ('G', 2), ('Y', 3), ('R', 4), ('RE', 5), ('NIR', 6), ('NIR2', 7)])
        elif num_bands == 9:  # assume WV3 MS & PAN in band 0
            # pan_bands = np.array([0])     # old version
            pan_bands = [2, 3, 5, 6]
            band_dict = OrderedDict(
                [('C', 1), ('B', 2), ('G', 3), ('Y', 4), ('R', 5), ('RE', 6), ('NIR', 7), ('NIR2', 8)])
        else:  # assume NGI aerial
            pan_bands = [0, 1, 2, 3]
            band_dict = OrderedDict([('R', 0), ('G', 1), ('B', 2), ('NIR', 3)])
        return pan_bands, band_dict

    def _generate_fn_dict(self):
        """
        Generate feature extraction dictionary with values= pointers to feature extraction functions, and
        keys= descriptive feature strings.
        Suitable for use with modelling from patches or mapping a whole image i.e. ImageFeatureExtractor and ImageMapper.
        Generates band ratios, veg. indices, simple texture measures, and non-linear transformations thereof.
        """

        if len(self.fn_dict) > 0:
            return

        # inner band ratios
        self.inner_dict = OrderedDict()
        self.inner_dict['pan/1'] = lambda pan, bands: pan
        for _num_key in list(self._band_dict.keys()):
            self.inner_dict['{0}/pan'.format(_num_key)] = lambda pan, bands, num_key=_num_key: bands[self._band_dict[num_key], :] / pan
            self.inner_dict['pan/{0}'.format(_num_key)] = lambda pan, bands, num_key=_num_key: pan / bands[self._band_dict[num_key], :]
            self.inner_dict['{0}/1'.format(_num_key)] = lambda pan, bands, num_key=_num_key: bands[self._band_dict[num_key], :]
            for _den_key in list(self._band_dict.keys()):
                if not _num_key == _den_key:
                    self.inner_dict['{0}/{1}'.format(_num_key, _den_key)] = lambda pan, bands, num_key=_num_key, den_key=_den_key:\
                        bands[self._band_dict[num_key], :] / bands[self._band_dict[den_key], :]

        # inner veg indices
        SAVI_L = 0.05   # wikipedia
        nir_keys = [key for key in list(self._band_dict.keys()) if ('NIR' in key) or ('RE' in key)]
        for _nir_key in nir_keys:
            post_fix = '' if _nir_key == 'NIR' else '_{0}'.format(_nir_key)
            self.inner_dict['NDVI' + post_fix] = lambda pan, bands, nir_key=_nir_key: 1 + (bands[self._band_dict[nir_key], :] - bands[self._band_dict['R'], :]) / \
                                                                                (bands[self._band_dict[nir_key], :] + bands[self._band_dict['R'], :])
            self.inner_dict['SAVI' + post_fix] = lambda pan, bands, nir_key=_nir_key: 1 + (1 + SAVI_L) * (bands[self._band_dict[nir_key], :] - bands[self._band_dict['R'], :]) / \
                                                                                (SAVI_L + bands[self._band_dict[nir_key], :] + bands[self._band_dict['R'], :])

        # window functions
        if self._apply_rolling_window:
            self.win_dict = OrderedDict({'mean': lambda x: np.nanmean(self._rolling_window_view(x), axis=(0, 1)),
                                         'std': lambda x: np.nanstd(self._rolling_window_view(x), axis=(0, 1)),
                                         'entropy': lambda x: nanentropy(self._rolling_window_view(x), axis=(0, 1))})
        else:
            self.win_dict = OrderedDict({'mean': lambda x: np.nanmean(x, axis=(0,1)), 'std': lambda x: np.nanstd(x, axis=(0,1)), 'entropy': lambda x: nanentropy(x, axis=(0,1))})

        self.scale_dict = OrderedDict({'log': np.log10, 'sqr': lambda x: np.power(x, 2), 'sqrt': np.sqrt})

        # combine inner, window and scaling functions
        for inner_key, _inner_fn in self.inner_dict.items():
            for win_key, _win_fn in self.win_dict.items():
                fn_key = '({0}({1}))'.format(win_key, inner_key)
                fn = lambda pan, bands, win_fn=_win_fn, inner_fn=_inner_fn: win_fn(inner_fn(pan, bands))
                self.fn_dict[fn_key] = fn
                if win_key == 'mean':   # for backward compatibility - TODO: try apply scaling to all windows, not only mean
                    for scale_key, _scale_fn in self.scale_dict.items():
                        fn_key = '{0}({1}({2}))'.format(scale_key, win_key, inner_key)
                        fn = lambda pan, bands, scale_fn=_scale_fn, win_fn=_win_fn, inner_fn=_inner_fn: scale_fn(win_fn(inner_fn(pan, bands)))
                        self.fn_dict[fn_key] = fn

    def extract_features(self, im_patch, mask=None, fn_keys=None):
        """
        Extract dictionary of features for multi-spectral image patch with optional mask

        Parameters
        ----------
        im_patch : numpy.array_like
            multi-spectal image patch - bands along first dimension (as with rasterio)
        mask : numpy.array_like
            mask of pixels to include (optional) - same x-y size as im_patch
        fn_keys : list
            feature keys to extract (optional - default extract all)

        Returns
        -------
        Dictionary of features feat_dict = {'<feature string>': <feature value>, ...}
        """
        if fn_keys is None:
            fn_keys = self.fn_dict.keys()

        if mask is None:
            mask = np.all(im_patch>0, axis=0)

        mask = np.bool8(mask)
        im_patch_mask = np.float64(im_patch)
        im_patch_mask[:, ~mask] = np.nan
        if self._num_bands != im_patch.shape[0]:
            raise Exception("im_patch must have the same number of bands as passed to MsPatchFeatureExtractor(...)")
        pan_mask = im_patch_mask[self._pan_bands, :].mean(axis=0)

        feat_dict = OrderedDict()
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", category=RuntimeWarning)    # suppress mean of empty slice warning
            for fn_key in fn_keys:
                feat_dict[fn_key] = self.fn_dict[fn_key](pan_mask, im_patch_mask)

        return feat_dict

class ImageFeatureExtractor(object):
    # Note: this could be a function rather than a class but have made it a class in line with PatchFeatureExtractor etc above
    def __init__(self, image_filename=None, plot_data_gdf=gpd.GeoDataFrame(), store_thumbnail=True):
        """
        Virtual base class to extract features from patches (e.g. ground truth plots) in an image

        Parameters
        ----------
        image_filename : str
            path to image file from which to extract features
        plot_data_gdf : geopandas.GeoDataFrame
            plot polygons with optional ground truth and index of plot ID strings
        """
        self.im_plot_data_gdf = gpd.GeoDataFrame()  # geodataframe of features in subindex 'feats' and data in 'data'
        self._image_reader = rasterio.open(image_filename, 'r')
        self._plot_data_gdf = plot_data_gdf
        self._patch_feature_extractor = None     # not implemented, to be specified in derived class constructor
        self._store_thumbnail = store_thumbnail

    def __del__(self):
        self._image_reader.close()

    def extract_image_features(self, feat_keys=None):
        """
        Extract features from plot polygons specified by plot_data_gdf in image

        Returns
        -------
        geopandas.GeoDataFrame of features in subindex 'feats' and data in 'data'
        """
        plot_data_gdf = self._plot_data_gdf
        plot_data_gdf = plot_data_gdf.to_crs(self._image_reader.crs)   # convert plot co-ordinates to image projection

        im_plot_data_dict = dict()      # dict to store plot features etc
        im_plot_count = 0
        max_thumbnail_vals = np.zeros(self._image_reader.count)    # max vals for each image band to scale thumbnails
        im_feat_dict = {}
        im_data_dict = {}

        for plot_id, plot in plot_data_gdf.iterrows():     # loop through plot polygons
            # convert polygon to raster mask
            plot_mask, plot_transform, plot_window = raster_geometry_mask(self._image_reader, [plot['geometry']], crop=True,
                                                                          all_touched=False)
            plot_cnrs_pixel =  np.array(plot_window.toranges())
            plot_mask = ~plot_mask  # TODO: can we lose this?

            # check plot window lies inside image
            if not (np.all(plot_cnrs_pixel[1, :] < self._image_reader.width) and np.all(plot_cnrs_pixel[0, :] < self._image_reader.height)
                    and np.all(plot_cnrs_pixel >= 0)):  # and plot.has_key('Yc') and plot['Yc'] > 0.:
                logger.warning(f'Excluding plot {plot["ID"]} - outside image extent')
                continue

            im_buf = self._image_reader.read(window=plot_window)     # read plot ROI from image

            if np.all(im_buf == 0):
                logger.warning(f'Excluding plot {plot["ID"]} - all pixels are zero')
                continue

            plot_mask = plot_mask & np.all(im_buf > 0, axis=0) & np.all(~np.isnan(im_buf), axis=0)  # exclude any nan or -ve pixels

            im_feat_dict = self._patch_feature_extractor.extract_features(im_buf, mask=plot_mask, fn_keys=feat_keys)  # extract image features for this plot

            # create plot thumbnail for visualisation in numpy format
            if self._store_thumbnail:
                thumbnail = np.float32(np.moveaxis(im_buf, 0, 2))
                thumbnail[~plot_mask] = 0.
                im_data_dict = {**im_feat_dict, **plot, 'thumbnail': thumbnail}  # combine features and other plot data

                # calc max thumbnail vals for scaling later
                max_val = np.percentile(thumbnail, 98., axis=(0, 1))
                max_thumbnail_vals[max_val > max_thumbnail_vals] = max_val[max_val > max_thumbnail_vals]
            else:
                im_data_dict = {**im_feat_dict, **plot}                             # combine features and other plot data

            im_plot_data_dict[plot_id] = im_data_dict       # add to dict of all plots
            im_plot_count += 1

            log_dict = {'ABC': 'Abc' in plot, 'Num zero pixels': (im_buf == 0).sum(), 'Num -ve pixels': (im_buf < 0).sum(),
                'Num nan pixels': np.isnan(im_buf).sum()}
            logger.info(', '.join([f'Plot {plot_id}'] + ['{}: {}'.format(k, v) for k, v in log_dict.items()]))

        logger.info('Processed {0} plots'.format(im_plot_count))

        # scale thumbnails for display
        if self._store_thumbnail:
            for im_data_dict in im_plot_data_dict.values():
                thumbnail = im_data_dict['thumbnail']
                for b in range(0, self._image_reader.count):
                    thumbnail[:, :, b] /= max_thumbnail_vals[b]
                    thumbnail[:, :, b][thumbnail[:, :, b] > 1.] = 1.
                im_data_dict['thumbnail'] = thumbnail

        # create MultiIndex column labels that separate features from other data
        data_labels = ['feats']*len(im_feat_dict) + ['data']*(len(im_data_dict) - len(im_feat_dict))
        columns = pd.MultiIndex.from_arrays([data_labels, list(im_data_dict.keys())], names=['high','low'])

        # create geodataframe of results
        self.im_plot_data_gdf = gpd.GeoDataFrame.from_dict(im_plot_data_dict, orient='index')
        self.im_plot_data_gdf = self.im_plot_data_gdf.set_crs(self._image_reader.crs)
        self.im_plot_data_gdf.columns = columns
        self.im_plot_data_gdf[('data','ID')] = self.im_plot_data_gdf.index

        return self.im_plot_data_gdf

class MsImageFeatureExtractor(ImageFeatureExtractor):
    def __init__(self, image_filename=None, plot_data_gdf=gpd.GeoDataFrame(), store_thumbnail=True):
        """
        Class to extract multi-spectral features from patches (e.g. ground truth plots) in an image

        Parameters
        ----------
        image_filename : str
            path to image file from which to extract features
        plot_data_gdf : geopandas.GeoDataFrame
            plot polygons with optional ground truth and index of plot ID strings
        """
        ImageFeatureExtractor.__init__(self, image_filename=image_filename, plot_data_gdf=plot_data_gdf, store_thumbnail=store_thumbnail)
        self._patch_feature_extractor = MsPatchFeatureExtractor(num_bands=self._image_reader.count)

class ImageMapper(object):
    def __init__(self, image_file_name='', map_file_name='', model=linear_model.LinearRegression(), model_feat_keys: list=None,
                 save_feats=False):
        """
        Virtual base class to generate raster map by applying a fitted model to an image

        Parameters
        ----------
        image_file_name : str
            path to input image file
        map_file_name: str
            output map file path to create
        model : sklearn.BaseEstimator
            trained/ fitted model to apply to image_file_name
        model_feat_keys : list
            list of descriptive strings specifying features to extract from image_file_name and to feed to model,
            (as understood by feature_extractor_type)
        save_feats : bool
            include extracted features as bands in map_file_name (default = False)
        """
        self.image_file_name = image_file_name
        self.map_file_name = map_file_name
        self.model = model
        self.model_keys = model_feat_keys
        self.save_feats = save_feats
        self.nodata = np.nan        # nodata value to use for map raster
        self._feature_extractor_type = PatchFeatureExtractor    # not implemented, to be specified in derived class

    def map(self, win_size=(1, 1), step_size=(1, 1)):
        """
        Apply model to image features and write to map raster file

        Parameters
        ----------
        win_size : numpy.array_like
            (x, y) rolling window size in pixels (default = (1,1))
        step_size : numpy.array_like
            (x, y) rolling window steps in pixels (default = (1,1))
        """
        if np.any(np.array(step_size) > np.array(win_size)):
            logger.warning('step_size <= win_size')

        with rasterio.Env():
            with rasterio.open(self.image_file_name, 'r') as image:
                # create the feature extraction class
                feature_extractor = self._feature_extractor_type(num_bands=image.count, apply_rolling_window=True,
                                                                          rolling_window_xsize=win_size[0], rolling_window_xstep=step_size[0])
                map_win_off = np.array([0, 0])  # test
                prog_update = 10

                if self.save_feats:
                    out_bands = len(self.model_keys) + 1
                else:
                    out_bands = 1

                # setup the output raster metadata based on image metadata and specified win_size and step_size
                map_profile = image.profile
                map_size = np.floor([1 + (image.width - win_size[0])/step_size[0], 1 + (image.height - win_size[1])/step_size[1]]).astype('int32')
                # map_size = np.floor([image.width / step_size[0], image.height / step_size[1]]).astype('int32')
                map_profile.update(dtype=rasterio.float32, count=out_bands, compress='deflate', driver='GTiff', width=map_size[0],
                                   height=map_size[1], nodata=self.nodata)

                map_profile['transform'] = map_profile['transform'] * rasterio.Affine.scale(*step_size)
                map_profile['transform'] = map_profile['transform'] * rasterio.Affine.translation(*np.floor(np.array(win_size)/(2 * np.array(step_size))))
                if (map_size[0] / map_profile['blockxsize'] < 10) | (map_size[1] / map_profile['blockysize'] < 10):
                    map_profile['tiled'] = False
                    map_profile['blockxsize'] = map_size[0]
                    map_profile['blockysize'] = 1

                with rasterio.open(self.map_file_name, 'w', **map_profile) as map_im:
                    # read image by groups of N rows where N = win_size[1], then slide window in x direction
                    # Note: rasterio index is x,y, numpy index is row, col (i.e. y,x)
                    for cy in range(0, image.height - win_size[1] + 1, step_size[1]):
                        image_win = Window(0, cy, (map_size[0] - 1)*step_size[0] + win_size[0], win_size[1])
                        bands = list(range(1, image.count + 1))
                        image_buf = image.read(bands, window=image_win).astype(rasterio.float32)  # NB bands along first dim

                        # pan = image_buf[feature_extractor._pan_bands, :, :].mean(axis=0)
                        map_win = Window(map_win_off[0], int(cy/step_size[0]) + map_win_off[1], map_size[0], 1)

                        # mask out negative and zero values to prevent NaN outputs - overly conservative but neater & faster than other options
                        image_nan_mask = np.all(image_buf > 0, axis=0)  # & np.all(pan != 0, axis=0)
                        # extract features (model features only - not entire library) from image_buf using rolling window
                        feat_dict = feature_extractor.extract_features(image_buf, mask=image_nan_mask, fn_keys=self.model_keys)

                        feat_buf = np.array([fv for fv in feat_dict.values()]).transpose()
                        feat_buf_nan_mask = np.any(np.isnan(feat_buf) | np.isinf(feat_buf), axis=1)    # work around for nan's which raise an error in sklearn
                        feat_buf[feat_buf_nan_mask, :] = 0
                        map_buf = self.model.predict(feat_buf).astype(map_im.dtypes[0])
                        map_buf[feat_buf_nan_mask] = self.nodata

                        map_im.write(map_buf.reshape(1, -1), indexes=1, window=map_win)
                        if self.save_feats:
                            feat_buf[feat_buf_nan_mask, :] = self.nodata
                            feat_buf = feat_buf.astype(map_im.dtypes[0]).transpose()[:, np.newaxis, :]   # order axes correctly for rasterio write
                            map_im.write(feat_buf, indexes=np.arange(2, 2 + len(self.model_keys)), window=map_win)

                        # TODO: proper progress bar
                        if np.ceil(100 * cy / (image.height - win_size[1] + 1)) >= prog_update:
                            logger.info(f'Progress {prog_update}%')
                            prog_update += 10

class MsImageMapper(ImageMapper):
    def __init__(self, image_file_name='', map_file_name='', model=linear_model.LinearRegression, model_feat_keys: list=None,
                 save_feats=False):
        """
        Class to generate raster map by applying fitted model to a multispectral image

        Parameters
        ----------
        image_file_name : str
            path to input image file
        map_file_name: str
            output map file path to create
        model : sklearn.BaseEstimator
            trained/ fitted model to apply to image_file_name
        model_feat_keys : list
            list of descriptive strings specifying features to extract from image_file_name and to feed to model,
            (as understood by feature_extractor_type)
        save_feats : bool
            include extracted features as bands in map_file_name (default = False)
        """
        ImageMapper.__init__(self, image_file_name=image_file_name, map_file_name=map_file_name, model=model, model_feat_keys=model_feat_keys,
                 save_feats=save_feats)
        self._feature_extractor_type = MsPatchFeatureExtractor

def thicket_agc_post_proc(image_mapper = ImageMapper()):
    """
    Post process thicket AGC map - helper function for MsImageMapper
    Writes out a cleaned version of map generated by ImageMapper.map(...), ills nodata and places sensible limits on
    AGC values.

    Parameters
    ----------
    image_mapper : ImageMapper
        instance of ImageMapper that has generated map

    Returns
    -------
    post processed raster file name
    """
    with rasterio.Env():
        with rasterio.open(image_mapper.map_file_name, 'r') as in_ds:
            out_profile = in_ds.profile
            out_profile.update(count=1)
            split_ext = os.path.splitext(image_mapper.map_file_name)
            out_file_name = '{0}_postproc{1}'.format(split_ext[0], split_ext[1])
            with rasterio.open(out_file_name, 'w', **out_profile) as out_ds:
                if (not out_profile['tiled']) or (np.prod(in_ds.shape) < 10e6):
                    in_windows = enumerate([Window(0,0,in_ds.width, in_ds.height)])    # read whole raster at once
                else:
                    in_windows = in_ds.block_windows(1)                                # read in blocks

                for ji, block_win in in_windows:
                    in_block = in_ds.read(1, window=block_win, masked=True)

                    in_block[in_block < 0] = 0
                    in_block.mask = (in_block.mask.astype(np.bool) | (in_block > 95) | (in_block < 0)).astype(rasterio.uint8)

                    in_mask = in_block.mask.copy()
                    sieved_msk = sieve(in_mask.astype(rasterio.uint8), size=2000)

                    out_block = fill.fillnodata(in_block, mask=None, max_search_distance=20, smoothing_iterations=1)
                    out_block[sieved_msk.astype(np.bool)] = image_mapper.nodata
                    out_ds.write(out_block, indexes=1, window=block_win)
        return out_file_name


# -*- encoding: utf-8 -*-

from django.test import TestCase

from caffe.models import Caffe
from employees.models import Employee

from .forms import CashReportForm, CompanyForm, ExpenseForm, FullExpenseForm
from .models import Company, Expense

class CompanyFormTest(TestCase):
    """Tests Company form."""

    def setUp(self):
        """Test data setup."""

        self.caffe = Caffe.objects.create(
            name='kafo',
            city='Gliwice',
            street='Wieczorka',
            house_number='14',
            postal_code='44-100'
        )
        self.filtry = Caffe.objects.create(
            name='filtry',
            city='Warszawa',
            street='Filry',
            house_number='14',
            postal_code='44-100'
        )

    def test_company(self):
        """Check validation."""

        bad_company = CompanyForm(
            {'name': ''},
            caffe=self.caffe
        )

        self.assertFalse(bad_company.is_valid())

        good_company = CompanyForm(
            {'name': 'very good'},
            caffe=self.caffe
        )

        self.assertTrue(good_company.is_valid())
        good_company.save()

        copycat_company = CompanyForm(
            {'name': 'very good'},
            caffe=self.caffe
        )
        self.assertFalse(copycat_company.is_valid())

        with self.assertRaises(Exception):
            CompanyForm({'bakery'})

    def test_name_validation(self):
        """Check name validation."""

        # same name, different caffes
        Company.objects.create(name='Bakery', caffe=self.filtry)

        form_correct = CompanyForm(
            {'name': 'Bakery'},
            caffe=self.caffe
        )
        self.assertTrue(form_correct.is_valid())

        # same name, same caffes
        Company.objects.create(name='Bakery', caffe=self.caffe)

        form_incorrect = CompanyForm(
            {'name': 'Bakery'},
            caffe=self.caffe
        )
        self.assertFalse(form_incorrect.is_valid())

class ExpenseFormTest(TestCase):
    """Tests Expense form."""

    def setUp(self):
        """Prepare objects for tests."""

        self.caffe = Caffe.objects.create(
            name='kafo',
            city='Gliwice',
            street='Wieczorka',
            house_number='14',
            postal_code='44-100'
        )
        self.filtry = Caffe.objects.create(
            name='filtry',
            city='Warszawa',
            street='Filry',
            house_number='14',
            postal_code='44-100'
        )

        self.company = Company.objects.create(
            name='GoodCake',
            caffe=self.caffe
        )

        self.company_f = Company.objects.create(
            name='GoodCake',
            caffe=self.filtry
        )

    def test_validation(self):
        """Check validation."""

        empty_name_expense = ExpenseForm(
            {'name': '', 'company': self.company},
            caffe=self.caffe
        )
        self.assertFalse(empty_name_expense.is_valid())

        empty_company_expense1 = ExpenseForm(
            {'name': 'newspapers'},
            caffe=self.caffe
        )
        self.assertTrue(empty_company_expense1.is_valid())

        empty_company_expense2 = ExpenseForm(
            {'name': 'newspapers', 'company': None},
            caffe=self.caffe
        )
        self.assertTrue(empty_company_expense2.is_valid())

        empty_company_expense3 = ExpenseForm(
            {'name': 'newspapers', 'company': ''},
            caffe=self.caffe
        )
        self.assertTrue(empty_company_expense3.is_valid())

        complete_expense = ExpenseForm(
            {'name': 'cakes', 'company': self.company.pk},
            caffe=self.caffe
        )
        self.assertTrue(complete_expense.is_valid())

        different_caffe = ExpenseForm(
            {'name': 'cakes', 'company': self.company.pk},
            caffe=self.filtry
        )
        with self.assertRaises(Exception):
            different_caffe.save()

        with self.assertRaises(Exception):
            ExpenseForm({'name': 'cakes', 'company': self.company.pk})

    def test_name_validation(self):
        """Check name validation."""

        Expense.objects.create(
            name='cakes',
            company=self.company_f,
            caffe=self.filtry
        )

        form_correct = ExpenseForm(
            {'name': 'cakes', 'company': self.company.pk},
            caffe=self.caffe
        )
        self.assertTrue(form_correct.is_valid())

        # invalid name
        Expense.objects.create(
            name='cakes',
            company=self.company,
            caffe=self.caffe
        )

        form_incorrect = ExpenseForm(
            {'name': 'cakes', 'company': self.company.pk},
            caffe=self.caffe
        )
        self.assertFalse(form_incorrect.is_valid())

class FullExpenseFormTest(TestCase):
    """Tests FullExpense form."""

    def setUp(self):
        """Prepare objects for tests."""

        self.caffe = Caffe.objects.create(
            name='kafo',
            city='Gliwice',
            street='Wieczorka',
            house_number='14',
            postal_code='44-100'
        )
        self.filtry = Caffe.objects.create(
            name='filtry',
            city='Warszawa',
            street='Filry',
            house_number='14',
            postal_code='44-100'
        )

        self.company = Company.objects.create(
            name='GoodCake',
            caffe=self.caffe
        )

        self.expense = Expense.objects.create(
            name='cakes',
            company=self.company,
            caffe=self.caffe
        )

    def test_validation(self):
        """Test validation."""

        cakes_for_10 = FullExpenseForm(
            {'expense': self.expense.pk, 'amount': 10},
            caffe=self.caffe
        )
        self.assertTrue(cakes_for_10.is_valid())

        cakes_for_nothing = FullExpenseForm(
            {'expense': self.expense.pk, 'amount': None},
            caffe=self.caffe
        )
        self.assertFalse(cakes_for_nothing.is_valid())

        money_for_nothing = FullExpenseForm(
            {'expense': None, 'amount': 10},
            caffe=self.caffe
        )
        self.assertFalse(money_for_nothing.is_valid())

        different_caffe = FullExpenseForm(
            {'expense': self.expense.pk, 'amount': 1},
            caffe=self.filtry
        )

        with self.assertRaises(Exception):
            different_caffe.save()

class CashReportFormTest(TestCase):
    """Tests CashReport form."""

    def setUp(self):
        """Test data setup."""

        self.caffe = Caffe.objects.create(
            name='kafo',
            city='Gliwice',
            street='Wieczorka',
            house_number='14',
            postal_code='44-100'
        )
        self.filtry = Caffe.objects.create(
            name='filtry',
            city='Warszawa',
            street='Filry',
            house_number='14',
            postal_code='44-100'
        )

        self.kate = Employee.objects.create(
            username='KateT',
            first_name='Kate',
            last_name='Tempest',
            telephone_number='12345678',
            email='kate@tempest.com',
            favorite_coffee='flat white',
            caffe=self.caffe
        )

    def test_validation(self):
        """Test validation of CashReport form."""

        no_cash_report = CashReportForm(
            {'card_payments': 500, 'amount_due': 1700},
            creator=self.kate,
            caffe=self.caffe,
        )
        self.assertFalse(no_cash_report.is_valid())

        no_cards_report = CashReportForm(
            {
                'cash_before_shift': 1000,
                'cash_after_shift': 2000,
                'amount_due': 1700
            },
            creator=self.kate,
            caffe=self.caffe
        )
        self.assertFalse(no_cards_report.is_valid())

        no_due_report = CashReportForm(
            {
                'card_payments': 500,
                'cash_before_shift': 1000,
                'cash_after_shift': 2000,
            },
            creator=self.kate,
            caffe=self.caffe
        )
        self.assertFalse(no_due_report.is_valid())

        perfectly_fine_report = CashReportForm(
            {
                'card_payments': 500,
                'cash_before_shift': 1000,
                'cash_after_shift': 2000,
                'amount_due': 1700
            },
            creator=self.kate,
            caffe=self.caffe
        )
        self.assertTrue(perfectly_fine_report.is_valid())

        different_caffe = CashReportForm(
            {
                'card_payments': 500,
                'cash_before_shift': 1000,
                'cash_after_shift': 2000,
                'amount_due': 1700
            },
            creator=self.kate,
            caffe=self.filtry
        )
        with self.assertRaises(Exception):
            different_caffe.save()

        with self.assertRaises(Exception):
            CashReportForm(
                {
                    'card_payments': 500,
                    'cash_before_shift': 1000,
                    'cash_after_shift': 2000,
                    'amount_due': 1700
                },
                creator=self.kate
            )

        with self.assertRaises(Exception):
            CashReportForm(
                {
                    'card_payments': 500,
                    'cash_before_shift': 1000,
                    'cash_after_shift': 2000,
                    'amount_due': 1700
                },
                caffe=self.caffe
            )


from io import BytesIO
import logging
import os
from pathlib import Path
import threading
from typing import Callable, Dict, Any

from flask import Flask, jsonify, request, send_file

from werkzeug.serving import make_server, run_simple, WSGIRequestHandler

from lansync.models import NodeChunk
from lansync.market import Market
from lansync.session import instance as session

class WSGIRequestHandlerHTTP11(WSGIRequestHandler):
    protocol_version = "HTTP/1.0"

certs_dir = Path.cwd() / "certs"

app = Flask(__name__)

@app.route("/market/<namespace_name>/<key>", methods=["POST"])
def exchange(namespace_name, key):
    market = Market.load_from_file(request.stream)
    own_market_exists = Market.load_from_db(namespace_name, key) is not None
    market.exchange_with_db()
    if not own_market_exists:
        session.sync_worker.schedule_event("scheduled_sync")
    fd = BytesIO()
    market.dump_to_file(fd)
    fd.seek(os.SEEK_SET)
    return send_file(fd, mimetype="application/octet-stream")

@app.route("/chunk/<namespace_name>/<content_hash>", methods=["GET", "HEAD"])
def chunk(namespace_name, content_hash):
    node_chunk_pair = NodeChunk.find(namespace_name, content_hash)
    if node_chunk_pair is None:
        return jsonify({"ok": False, "error": "Not found"}), 404

    chunk, read_chunk = node_chunk_pair

    if request.method == "HEAD":
        return "", 200

    fd = BytesIO(read_chunk())
    return send_file(fd, mimetype="application/octet-stream")

def run(app, debug=False, on_start: Callable[[int], None] = None):
    options: Dict[str, Any] = {}
    options.setdefault("threaded", True)
    options.setdefault(
        "ssl_context", (os.fspath(certs_dir / "alpha.crt"), os.fspath(certs_dir / "alpha.key"),)
    )
    options.setdefault("request_handler", WSGIRequestHandlerHTTP11)

    host = "0.0.0.0"
    port = 0

    if debug:
        options.setdefault("use_reloader", debug)
        options.setdefault("use_debugger", debug)
        run_simple(host, port, app, **options)  # ???
    else:
        server = make_server(host, port, app, **options)
        _, port = server.server_address

        logging.info("Serving on port: %d", port)
        if on_start is not None:
            on_start(port)

        server.serve_forever()

def run_in_thread(debug=False, on_start: Callable[[int], None] = None):
    threading.Thread(target=run, args=(app, debug, on_start), daemon=True).start()


# proc1: subscribes to the book-100 and trade data to get the snapshot message followed by 100 book-100 or trade
# messages in a new file each time, disconnects, sleeps for 5 min, awakens, repeats until an hour has passed at which point it is dead

import asyncio
import time
from multiprocessing import Process, Queue
from main import main
import random

def send_kill(q, pid):
    print("putting", pid)
    q.put(pid)
    #print("waiting")
    #response = q.get(block=True) # should be pretty fast
    #return response

if __name__ == "__main__":

    one_hour_start = time.time()
    ctr = 0
    procs = []
    killQ = []
    WAIT_SEC = 2
    TOTAL_RUNTIME_MIN = 5
    forced_sleep_interval_sec = 30

    for i in range(3):
        killQ.append(Queue())
        p  = Process(target=main, args=(time.time(), killQ[i], 'trade',f'{i}'), name = ('process_' + str(i+1)))
        procs.append(p)
        p.start()
        print( 'starting', p.name)

    while (time.time() - one_hour_start) < TOTAL_RUNTIME_MIN * 60:
        print(time.time() - one_hour_start)

        dead = random.randint(0,10)
        if dead % 2 == 0:
            procsecuted = random.randint(0,2)
            send_kill(killQ[procsecuted], procs[procsecuted].pid)
            # procs[procsecuted].terminate()
            print("killed proc!", procs[procsecuted])
            procs.remove(procs[procsecuted])

            asyncio.sleep(forced_sleep_interval_sec - 10)

            newproc = Process(target=main, args=(time.time(), killQ[procsecuted], 'trade',f'{procsecuted}'), name = ('process_' + str(procsecuted)))
            procs.append(newproc)
            newproc.start()
            print("started new proc")
        print("sleeping")
        asyncio.sleep(forced_sleep_interval_sec)

    print("that's time")
    for idx,p in enumerate(procs):
        send_kill(killQ[idx], p.pid)
        p.terminate()
        p.join()


import numpy as np

def check_crossroad(robot, point1, point2, point3, point4) -> bool:
    r, p1, p3 = map(np.array, [robot, point1, point3])
    vec1, vec2 = (r - p1) >= 0, (p3 - r) >= 0
    return bool(np.prod(vec1) and np.prod(vec2))

def check_collision(coefficients) -> list:
    ans = []
    for i in range(len(coefficients)):
        for j in range(i + 1, len(coefficients)):
            eq1 = np.array(coefficients[i][:2])
            eq2 = np.array(coefficients[j][:2])
            if np.linalg.det(np.vstack((eq1, eq2))):
                ans.append((i, j))
                ans.append((j, i))
    return ans

def check_path(points_list) -> float:
    path = 0
    for i in range(1, len(points_list)):
        vec1, vec2 = np.array(points_list[i - 1]), np.array(points_list[i])
        path += np.linalg.norm(vec2 - vec1, ord=2)
    return round(path, 2)


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

from alipay.aop.api.constant.ParamConstants import *
from alipay.aop.api.domain.ActivityOrderDTO import ActivityOrderDTO
from alipay.aop.api.domain.BudgetInfo import BudgetInfo
from alipay.aop.api.domain.ConstraintInfo import ConstraintInfo
from alipay.aop.api.domain.PromoTool import PromoTool
from alipay.aop.api.domain.PublishChannel import PublishChannel
from alipay.aop.api.domain.RecruitInfo import RecruitInfo

class CampDetail(object):

    def __init__(self):
        self._activity_orders = None
        self._audit_status = None
        self._auto_delay_flag = None
        self._budget_info = None
        self._constraint_info = None
        self._desc = None
        self._end_time = None
        self._ext_info = None
        self._id = None
        self._name = None
        self._promo_tools = None
        self._publish_channels = None
        self._recruit_info = None
        self._start_time = None
        self._status = None
        self._type = None

    @property
    def activity_orders(self):
        return self._activity_orders

    @activity_orders.setter
    def activity_orders(self, value):
        if isinstance(value, list):
            self._activity_orders = list()
            for i in value:
                if isinstance(i, ActivityOrderDTO):
                    self._activity_orders.append(i)
                else:
                    self._activity_orders.append(ActivityOrderDTO.from_alipay_dict(i))
    @property
    def audit_status(self):
        return self._audit_status

    @audit_status.setter
    def audit_status(self, value):
        self._audit_status = value
    @property
    def auto_delay_flag(self):
        return self._auto_delay_flag

    @auto_delay_flag.setter
    def auto_delay_flag(self, value):
        self._auto_delay_flag = value
    @property
    def budget_info(self):
        return self._budget_info

    @budget_info.setter
    def budget_info(self, value):
        if isinstance(value, BudgetInfo):
            self._budget_info = value
        else:
            self._budget_info = BudgetInfo.from_alipay_dict(value)
    @property
    def constraint_info(self):
        return self._constraint_info

    @constraint_info.setter
    def constraint_info(self, value):
        if isinstance(value, ConstraintInfo):
            self._constraint_info = value
        else:
            self._constraint_info = ConstraintInfo.from_alipay_dict(value)
    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        self._desc = value
    @property
    def end_time(self):
        return self._end_time

    @end_time.setter
    def end_time(self, value):
        self._end_time = value
    @property
    def ext_info(self):
        return self._ext_info

    @ext_info.setter
    def ext_info(self, value):
        self._ext_info = value
    @property
    def id(self):
        return self._id

    @id.setter
    def id(self, value):
        self._id = value
    @property
    def name(self):
        return self._name

    @name.setter
    def name(self, value):
        self._name = value
    @property
    def promo_tools(self):
        return self._promo_tools

    @promo_tools.setter
    def promo_tools(self, value):
        if isinstance(value, list):
            self._promo_tools = list()
            for i in value:
                if isinstance(i, PromoTool):
                    self._promo_tools.append(i)
                else:
                    self._promo_tools.append(PromoTool.from_alipay_dict(i))
    @property
    def publish_channels(self):
        return self._publish_channels

    @publish_channels.setter
    def publish_channels(self, value):
        if isinstance(value, list):
            self._publish_channels = list()
            for i in value:
                if isinstance(i, PublishChannel):
                    self._publish_channels.append(i)
                else:
                    self._publish_channels.append(PublishChannel.from_alipay_dict(i))
    @property
    def recruit_info(self):
        return self._recruit_info

    @recruit_info.setter
    def recruit_info(self, value):
        if isinstance(value, RecruitInfo):
            self._recruit_info = value
        else:
            self._recruit_info = RecruitInfo.from_alipay_dict(value)
    @property
    def start_time(self):
        return self._start_time

    @start_time.setter
    def start_time(self, value):
        self._start_time = value
    @property
    def status(self):
        return self._status

    @status.setter
    def status(self, value):
        self._status = value
    @property
    def type(self):
        return self._type

    @type.setter
    def type(self, value):
        self._type = value

    def to_alipay_dict(self):
        params = dict()
        if self.activity_orders:
            if isinstance(self.activity_orders, list):
                for i in range(0, len(self.activity_orders)):
                    element = self.activity_orders[i]
                    if hasattr(element, 'to_alipay_dict'):
                        self.activity_orders[i] = element.to_alipay_dict()
            if hasattr(self.activity_orders, 'to_alipay_dict'):
                params['activity_orders'] = self.activity_orders.to_alipay_dict()
            else:
                params['activity_orders'] = self.activity_orders
        if self.audit_status:
            if hasattr(self.audit_status, 'to_alipay_dict'):
                params['audit_status'] = self.audit_status.to_alipay_dict()
            else:
                params['audit_status'] = self.audit_status
        if self.auto_delay_flag:
            if hasattr(self.auto_delay_flag, 'to_alipay_dict'):
                params['auto_delay_flag'] = self.auto_delay_flag.to_alipay_dict()
            else:
                params['auto_delay_flag'] = self.auto_delay_flag
        if self.budget_info:
            if hasattr(self.budget_info, 'to_alipay_dict'):
                params['budget_info'] = self.budget_info.to_alipay_dict()
            else:
                params['budget_info'] = self.budget_info
        if self.constraint_info:
            if hasattr(self.constraint_info, 'to_alipay_dict'):
                params['constraint_info'] = self.constraint_info.to_alipay_dict()
            else:
                params['constraint_info'] = self.constraint_info
        if self.desc:
            if hasattr(self.desc, 'to_alipay_dict'):
                params['desc'] = self.desc.to_alipay_dict()
            else:
                params['desc'] = self.desc
        if self.end_time:
            if hasattr(self.end_time, 'to_alipay_dict'):
                params['end_time'] = self.end_time.to_alipay_dict()
            else:
                params['end_time'] = self.end_time
        if self.ext_info:
            if hasattr(self.ext_info, 'to_alipay_dict'):
                params['ext_info'] = self.ext_info.to_alipay_dict()
            else:
                params['ext_info'] = self.ext_info
        if self.id:
            if hasattr(self.id, 'to_alipay_dict'):
                params['id'] = self.id.to_alipay_dict()
            else:
                params['id'] = self.id
        if self.name:
            if hasattr(self.name, 'to_alipay_dict'):
                params['name'] = self.name.to_alipay_dict()
            else:
                params['name'] = self.name
        if self.promo_tools:
            if isinstance(self.promo_tools, list):
                for i in range(0, len(self.promo_tools)):
                    element = self.promo_tools[i]
                    if hasattr(element, 'to_alipay_dict'):
                        self.promo_tools[i] = element.to_alipay_dict()
            if hasattr(self.promo_tools, 'to_alipay_dict'):
                params['promo_tools'] = self.promo_tools.to_alipay_dict()
            else:
                params['promo_tools'] = self.promo_tools
        if self.publish_channels:
            if isinstance(self.publish_channels, list):
                for i in range(0, len(self.publish_channels)):
                    element = self.publish_channels[i]
                    if hasattr(element, 'to_alipay_dict'):
                        self.publish_channels[i] = element.to_alipay_dict()
            if hasattr(self.publish_channels, 'to_alipay_dict'):
                params['publish_channels'] = self.publish_channels.to_alipay_dict()
            else:
                params['publish_channels'] = self.publish_channels
        if self.recruit_info:
            if hasattr(self.recruit_info, 'to_alipay_dict'):
                params['recruit_info'] = self.recruit_info.to_alipay_dict()
            else:
                params['recruit_info'] = self.recruit_info
        if self.start_time:
            if hasattr(self.start_time, 'to_alipay_dict'):
                params['start_time'] = self.start_time.to_alipay_dict()
            else:
                params['start_time'] = self.start_time
        if self.status:
            if hasattr(self.status, 'to_alipay_dict'):
                params['status'] = self.status.to_alipay_dict()
            else:
                params['status'] = self.status
        if self.type:
            if hasattr(self.type, 'to_alipay_dict'):
                params['type'] = self.type.to_alipay_dict()
            else:
                params['type'] = self.type
        return params

    @staticmethod
    def from_alipay_dict(d):
        if not d:
            return None
        o = CampDetail()
        if 'activity_orders' in d:
            o.activity_orders = d['activity_orders']
        if 'audit_status' in d:
            o.audit_status = d['audit_status']
        if 'auto_delay_flag' in d:
            o.auto_delay_flag = d['auto_delay_flag']
        if 'budget_info' in d:
            o.budget_info = d['budget_info']
        if 'constraint_info' in d:
            o.constraint_info = d['constraint_info']
        if 'desc' in d:
            o.desc = d['desc']
        if 'end_time' in d:
            o.end_time = d['end_time']
        if 'ext_info' in d:
            o.ext_info = d['ext_info']
        if 'id' in d:
            o.id = d['id']
        if 'name' in d:
            o.name = d['name']
        if 'promo_tools' in d:
            o.promo_tools = d['promo_tools']
        if 'publish_channels' in d:
            o.publish_channels = d['publish_channels']
        if 'recruit_info' in d:
            o.recruit_info = d['recruit_info']
        if 'start_time' in d:
            o.start_time = d['start_time']
        if 'status' in d:
            o.status = d['status']
        if 'type' in d:
            o.type = d['type']
        return o


"""Tests for the jupyter notebooks."""
import os
import subprocess
from sty import fg, ef, rs

def test_notebooks(notebook):
    """Function to run a notebook.

    Parameters
    ----------
    notebook : str
        Path and name of notebook to be tested.
    """
    # Break up the path.
    psplit = notebook.split('/')
    name = psplit.pop(-1)
    print(fg.red + '\ntesting notebook: ' + fg.rs + fg.green +
          '{}\n'.format(name) + fg.rs)
    dirpath = '/'.join(psplit)

    # Execute the jupyter notebook.
    cmd = 'jupyter nbconvert {} --to notebook --execute '.format(notebook)
    cmd += '--ExecutePreprocessor.timeout=6000 --output out_file'
    subprocess.call(cmd, shell=True)

    # Try and clean up output. Otherwise fail and exit.
    try:
        os.remove('{}/out_file.ipynb'.format(dirpath))
    except FileNotFoundError:
        print(ef.bold + fg(201) + '\n{} failed\n'.format(name) + fg.rs +
              rs.bold)
        exit()

if __name__ == '__main__':
    path = os.getcwd()

    flist = []

    # Find all notebooks in the tutorials directory.
    for root, dirs, files in os.walk(path):
        for file in files:
            if 'checkpoint' not in file:
                if file.endswith(".ipynb"):
                    flist.append(os.path.join(root, file))

    # Run all the tests.
    for f in flist:
        test_notebooks(f)


name = input("Enter File :")
try :
    fhand = open(name)
except :
    print("Enter Correct file name.")
    exit()

counts = dict()

for line in fhand :
    line.strip()
    if line.startswith('From '):
        words = line.split()
        s = words[5]
        f = s[:2]
        counts[f] = counts.get(f, 0) + 1

sortc = sorted( counts.items() )

for k, v in sortc :
    print(k, v)


import sys
sys.stdin=open("../input.txt", "rt")

n = int(input())

for i in range(1, n+1):
    word = input()
    word = word.upper()
    c = len(word) // 2
    cnt = 0
    for j in range(c):
        if word[j] != word[-1-j]:
            print(f"#{i} NO")
            break
    else:
        print(f"#{i} YES")

        # print(word[j], word[len(word)-(j+1)])
        # print(j, len(word) - (j+1))


import matplotlib
matplotlib.use('TkAgg')
from matplotlib import pyplot as plt
from math import *

xf = []
yf = []
xF = []
yFr = []
yFi = []

min_xf = -2.0
max_xf = 4.0

min_xF = -9.0
max_xF = 9.0

tf = min_xf
tF = min_xF
for i in range(0, 1000):
    xf.append(tf)
    xF.append(tF)

    yf.append(exp(-tf))
    yFr.append((complex(1,0)/ complex(1,tF)).real)
    yFi.append((complex(1,0)/ complex(1,tF)).imag)

    tf = tf + (max_xf - min_xf)/1000
    tF = tF + (max_xF - min_xF)/1000

plt.subplot(121)
plt.plot(xf, yf, label = 'f(t) = exp(-t)')
plt.xlabel = 't'
plt.ylabel = 'f(t)'
plt.legend()

plt.subplot(122)
plt.plot(xF, yFr, 'b', label='f(t) = Re(1/(1+jt))')
plt.plot(xF, yFi, 'r', label='f(t) = Im(1/(1+jt))')
plt.xlabel = 't'
plt.ylabel = 'f(t)'
plt.legend()

#plt.savefig('plottedPic.png')
plt.savefig('plottedPic.png', bbox_inches='tight')

plt.show()


import os
import random
import torch
import math
import glob

from PIL import Image
import numpy as np

from .dataloader_base import DLBase
import datasets.daugm_video as tf

class DataVideo(DLBase):

    def __init__(self, cfg, split, val=False):
        super(DataVideo, self).__init__()

        self.cfg = cfg
        self.split = split
        self.val = val

        self.cfg_frame_gap = cfg.DATASET.VAL_FRAME_GAP if val else cfg.DATASET.FRAME_GAP

        self._init_palette(cfg.TRAIN.BATCH_SIZE * cfg.MODEL.GRID_SIZE**2)

        # train/val/test splits are pre-cut
        split_fn = os.path.join(cfg.DATASET.ROOT, "filelists", self.split + ".txt")
        assert os.path.isfile(split_fn)

        self.images = []

        token = None # sequence token (new video when it changes)

        subsequence = []
        ignored = [0]
        num_frames = [0]

        def add_sequence():

            if cfg.DATASET.VIDEO_LEN > len(subsequence):
                # found a very short sequence
                ignored[0] += 1
            else:
                # adding the subsequence
                self.images.append(tuple(subsequence))
                num_frames[0] += len(subsequence)

            del subsequence[:]

        with open(split_fn, "r") as lines:
            for line in lines:
                _line = line.strip("\n").split(' ')

                assert len(_line) > 0, "Expected at least one path"
                _image = _line[0]

                # each sequence may have a different length
                # do some book-keeping e.g. to ensure we have
                # sequences long enough for subsequent sampling
                _token = _image.split("/")[-2] # parent directory

                # sequence ID is in the filename
                #_token = os.path.basename(_image).split("_")[0]
                if token != _token:
                    if not token is None:
                        add_sequence()
                    token = _token

                # image 1
                _image = os.path.join(cfg.DATASET.ROOT, _image.lstrip('/'))
                #assert os.path.isfile(_image), '%s not found' % _image
                subsequence.append(_image)

        # update the last sequence
        # returns the total amount of frames
        add_sequence()
        print("Dataloader: {}".format(split), " / Frame Gap: ", self.cfg_frame_gap)
        print("Loaded {} sequences / {} ignored / {} frames".format(len(self.images), \
                                                                    ignored[0], \
                                                                    num_frames[0]))

        self._num_samples = num_frames[0]
        self._init_augm(cfg)

    def _init_augm(self, cfg):

        # general (unguided) affine transformations
        tfs_pre = [tf.CreateMask()]
        self.tf_pre = tf.Compose(tfs_pre)

        # photometric noise
        tfs_affine = []

        # guided affine transformations
        tfs_augm = []

        # 1.
        # general affine transformations
        #
        tfs_pre.append(tf.MaskScaleSmallest(cfg.DATASET.SMALLEST_RANGE))

        if cfg.DATASET.RND_CROP:
            tfs_pre.append(tf.MaskRandCrop(cfg.DATASET.CROP_SIZE, pad_if_needed=True))
        else:
            tfs_pre.append(tf.MaskCenterCrop(cfg.DATASET.CROP_SIZE))

        if cfg.DATASET.RND_HFLIP:
            tfs_pre.append(tf.MaskRandHFlip())

        # 2.
        # Guided affine transformation
        #
        if cfg.DATASET.GUIDED_HFLIP:
            tfs_affine.append(tf.GuidedRandHFlip())

        # this will add affine transformation
        if cfg.DATASET.RND_ZOOM:
            tfs_affine.append(tf.MaskRandScaleCrop(*cfg.DATASET.RND_ZOOM_RANGE))

        self.tf_affine = tf.Compose(tfs_affine)
        self.tf_affine2 = tf.Compose([tf.AffineIdentity()])

        tfs_post = [tf.ToTensorMask(),
                    tf.Normalize(mean=self.MEAN, std=self.STD),
                    tf.ApplyMask(-1)]

        # image to the teacher will have no noise
        self.tf_post = tf.Compose(tfs_post)

    def set_num_samples(self, n):
        print("Re-setting # of samples: {:d} -> {:d}".format(self._num_samples, n))
        self._num_samples = n

    def __len__(self):
        return len(self.images) #self._num_samples

    def denorm(self, image):

        if image.dim() == 3:
            assert image.dim() == 3, "Expected image [CxHxW]"
            assert image.size(0) == 3, "Expected RGB image [3xHxW]"

            for t, m, s in zip(image, self.MEAN, self.STD):
                t.mul_(s).add_(m)
        elif image.dim() == 4:
            # batch mode
            assert image.size(1) == 3, "Expected RGB image [3xHxW]"

            for t, m, s in zip((0,1,2), self.MEAN, self.STD):
                image[:, t, :, :].mul_(s).add_(m)

        return image

    def _get_affine(self, params):

        N = len(params)

        # construct affine operator
        affine = torch.zeros(N, 2, 3)

        aspect_ratio = float(self.cfg.DATASET.CROP_SIZE[0]) / \
                            float(self.cfg.DATASET.CROP_SIZE[1])

        for i, (dy,dx,alpha,scale,flip) in enumerate(params):

            # R inverse
            sin = math.sin(alpha * math.pi / 180.)
            cos = math.cos(alpha * math.pi / 180.)

            # inverse, note how flipping is incorporated
            affine[i,0,0], affine[i,0,1] = flip * cos, sin * aspect_ratio
            affine[i,1,0], affine[i,1,1] = -sin / aspect_ratio, cos

            # T inverse Rinv * t == R^T * t
            affine[i,0,2] = -1. * (cos * dx + sin * dy)
            affine[i,1,2] = -1. * (-sin * dx + cos * dy)

            # T
            affine[i,0,2] /= float(self.cfg.DATASET.CROP_SIZE[1] // 2)
            affine[i,1,2] /= float(self.cfg.DATASET.CROP_SIZE[0] // 2)

            # scaling
            affine[i] *= scale

        return affine

    def _get_affine_inv(self, affine, params):

        aspect_ratio = float(self.cfg.DATASET.CROP_SIZE[0]) / \
                            float(self.cfg.DATASET.CROP_SIZE[1])

        affine_inv = affine.clone()
        affine_inv[:,0,1] = affine[:,1,0] * aspect_ratio**2
        affine_inv[:,1,0] = affine[:,0,1] / aspect_ratio**2
        affine_inv[:,0,2] = -1 * (affine_inv[:,0,0] * affine[:,0,2] + affine_inv[:,0,1] * affine[:,1,2])
        affine_inv[:,1,2] = -1 * (affine_inv[:,1,0] * affine[:,0,2] + affine_inv[:,1,1] * affine[:,1,2])

        # scaling
        affine_inv /= torch.Tensor(params)[:,3].view(-1,1,1)**2

        return affine_inv

    def __getitem__(self, index):

        # searching for the video clip ID
        sequence = self.images[index] # % len(self.images)]
        seqlen = len(sequence)

        assert self.cfg_frame_gap > 0, "Frame gap should be positive"
        t_window = self.cfg_frame_gap * self.cfg.DATASET.VIDEO_LEN

        # reduce sampling gap for short clips
        t_window = min(seqlen, t_window)
        frame_gap = t_window // self.cfg.DATASET.VIDEO_LEN

        # strided slice
        frame_ids = torch.arange(t_window)[::frame_gap]
        frame_ids = frame_ids[:self.cfg.DATASET.VIDEO_LEN]
        assert len(frame_ids) == self.cfg.DATASET.VIDEO_LEN

        # selecting a random start
        index_start = random.randint(0, seqlen - frame_ids[-1] - 1)
        # permuting the frames in the batch
        random_ids = torch.randperm(self.cfg.DATASET.VIDEO_LEN)
        # adding the offset
        frame_ids = frame_ids[random_ids] + index_start

        # forward sequence
        images = []
        for frame_id in frame_ids:
            fn = sequence[frame_id]
            images.append(Image.open(fn).convert('RGB'))

        # 1. general transforms
        frames, valid = self.tf_pre(images)

        # 1.1 creating two sequences in forward/backward order
        frames1, valid1 = frames[:], valid[:]

        # second copy
        frames2 = [f.copy() for f in frames]
        valid2 = [v.copy() for v in valid]

        # 2. guided affine transforms
        frames1, valid1, affine_params1 = self.tf_affine(frames1, valid1)
        frames2, valid2, affine_params2 = self.tf_affine2(frames2, valid2)

        # convert to tensor, zero out the values
        frames1 = self.tf_post(frames1, valid1)
        frames2 = self.tf_post(frames2, valid2)

        # converting the affine transforms
        aff_reg = self._get_affine(affine_params1)
        aff_main = self._get_affine(affine_params2)

        aff_reg_inv = self._get_affine_inv(aff_reg, affine_params1)

        aff_reg = aff_main # identity affine2_inv
        aff_main = aff_reg_inv

        frames1 = torch.stack(frames1, 0)
        frames2 = torch.stack(frames2, 0)

        assert frames1.shape == frames2.shape

        return frames2, frames1, aff_main, aff_reg

class DataVideoKinetics(DataVideo):

    def __init__(self, cfg, split):
        super(DataVideo, self).__init__()

        self.cfg = cfg
        self.split = split

        self._init_palette(cfg.TRAIN.BATCH_SIZE * cfg.MODEL.GRID_SIZE**2)

        # train/val/test splits are pre-cut
        split_fn = os.path.join(cfg.DATASET.ROOT, "filelists", self.split + ".txt")
        assert os.path.isfile(split_fn)

        self.videos = []

        with open(split_fn, "r") as lines:
            for line in lines:
                _line = line.strip("\n").split(' ')
                assert len(_line) > 0, "Expected at least one path"

                _vid = _line[0]

                # image 1
                _vid = os.path.join(cfg.DATASET.ROOT, _vid.lstrip('/'))
                #assert os.path.isdir(_vid), "{} does not exist".format(_vid)
                self.videos.append(_vid)

        # update the last sequence
        # returns the total amount of frames
        print("DataloaderKinetics: {}".format(split))
        print("Loaded {} sequences".format(len(self.videos)))
        self._init_augm(cfg)

    def __len__(self):
        return len(self.videos)

    def __getitem__(self, index):

        C = self.cfg.DATASET
        path = self.videos[index]

        # filenames
        fns = sorted(glob.glob(path + "/*.jpg"))
        total_len = len(fns)

        # temporal window to consider
        temp_window = min(total_len, C.VIDEO_LEN * C.FRAME_GAP)
        gap = temp_window / C.VIDEO_LEN
        start_frame = random.randint(0, total_len - temp_window)

        images = []
        for idx in range(C.VIDEO_LEN):
            frame_id = start_frame + int(idx * gap)
            fn = fns[frame_id]
            images.append(Image.open(fn).convert('RGB'))

        # 1. general transforms
        frames, valid = self.tf_pre(images)

        # 1.1 creating two sequences in forward/backward order
        frames1, valid1 = frames[:], valid[:]

        # second copy
        frames2 = [f.copy() for f in frames]
        valid2 = [v.copy() for v in valid]

        # 2. guided affine transforms
        frames1, valid1, affine_params1 = self.tf_affine(frames1, valid1)
        frames2, valid2, affine_params2 = self.tf_affine2(frames2, valid2)

        # convert to tensor, zero out the values
        frames1 = self.tf_post(frames1, valid1)
        frames2 = self.tf_post(frames2, valid2)

        # converting the affine transforms
        affine1 = self._get_affine(affine_params1)
        affine2 = self._get_affine(affine_params2)

        affine1_inv = self._get_affine_inv(affine1, affine_params1)

        affine1 = affine2 # identity affine2_inv
        affine2 = affine1_inv

        frames1 = torch.stack(frames1, 0)
        frames2 = torch.stack(frames2, 0)

        assert frames1.shape == frames2.shape

        return frames2, frames1, affine2, affine1


from .sr_view import SRView
import nextcord
from nextcord.ext import commands

ownerRole = 942457163614400613

class Button(commands.Cog, name="SR"):
    def __init__(self, bot: commands.Bot):
        self.bot = bot
        self.sr_persistent_view = False
    @commands.Cog.listener()
    async def on_ready(self):
        if not self.sr_persistent_view:
            self.bot.add_view(SRView())
            self.sr_persistent_view = True
            print('loading persistent view')

    @commands.command()
    async def sr(self, ctx: commands.Context):
        if ownerRole in[r.id for r in ctx.author.roles]:
            embed = nextcord.Embed(
                title = "Self Roles",
                description="Get your very own self roles here, Look at the list below for some options on what to pick.",
                color=nextcord.Color.fuchsia()
            )
            await ctx.send(embed=embed, view=SRView())
        else:
            await ctx.send("This command is for the owner only!")

def setup(bot: commands.Bot):
    bot.add_cog(Button(bot))


# genetic algorithm search for continuous function optimization
from numpy.random import randint
from numpy.random import rand
from algorithms.functions import beale_function
import methods.mutation as mutation
import methods.cross as cross
import methods.selection as select
from ui.UserInputs import UserInputs
from statistics import stdev, mean

def inversion(bitstring: list, r_inve):
    c1 = bitstring.copy()
    if rand() < r_inve:
        pt1 = randint(1, (len(bitstring) - 2) / 2)
        pt2 = randint((len(bitstring) - 2) / 2, len(bitstring) - 2)
        x = bitstring[pt1:pt2]
        x.reverse()
        c1 = bitstring[:pt1] + x + bitstring[pt2:]
    return c1

# decode bitstring to numbers
def decode(bounds, n_bits, bitstring):
    decoded = list()
    largest = 2 ** n_bits
    for i in range(len(bounds)):
        start, end = i * n_bits, (i * n_bits) + n_bits
        substring = bitstring[start:end]
        chars = ''.join([str(s) for s in substring])
        integer = int(chars, 2)
        value = bounds[i][0] + (integer / largest) * (bounds[i][1] - bounds[i][0])
        decoded.append(value)
    return decoded

def genetic_algorithm(user_input):
    bounds = [[user_input.begin_range_a, user_input.end_range_b], [user_input.begin_range_a, user_input.end_range_b]]
    pop = [randint(0, 2, user_input.number_of_bits * len(bounds)).tolist() for _ in range(user_input.population_amount)]
    best, best_eval = 0, beale_function(decode(bounds, user_input.number_of_bits, pop[0]))
    gen_b_rows = []
    gen_avg_rows = []
    gen_std_dev_rows = []
    for gen in range(user_input.epochs_amount):
        decoded = [decode(bounds, user_input.number_of_bits, p) for p in pop]
        scores = [beale_function(d) for d in decoded]
        best_iter = scores[0]
        for i in range(user_input.population_amount):
            if user_input.maximum:
                if scores[i] > best_eval:
                    best, best_eval = pop[i], scores[i]
                    print(">%d, new best f(%s) = %f" % (gen, decoded[i], scores[i]))
                if scores[i] > best_iter:
                    best_iter = scores[i]
            else:
                if scores[i] < best_eval:
                    best, best_eval = pop[i], scores[i]
                    print(">%d, new best f(%s) = %f" % (gen, decoded[i], scores[i]))
                if scores[i] < best_iter:
                    best_iter = scores[i]
        gen_b_rows.append([str(gen), str(best_iter)])
        gen_avg_rows.append([str(gen), mean(scores)])
        gen_std_dev_rows.append([str(gen), str(stdev(scores))])
        if user_input.selection_method == "Tournament":
            selected = [select.tournament(pop, scores, user_input.best_and_tournament_chromosome_amount) for _ in
                        range(user_input.population_amount)]
        if user_input.selection_method == "Roulette":
            selected = [select.roulette(pop, scores) for _ in range(user_input.population_amount)]
        if user_input.selection_method == "Best":
            selected = [select.best(pop, scores, _) for _ in range(user_input.population_amount)]
        children = list()
        if user_input.elite_strategy:
            children.append(best)
        for i in range(0, user_input.population_amount, 2):
            p1, p2 = selected[i], selected[i + 1]
            if user_input.cross_method == "One Point Cross":
                cross_result = cross.crossover(p1, p2, user_input.cross_probability)
            if user_input.cross_method == "Two Point Cross":
                cross_result = cross.crossover2(p1, p2, user_input.cross_probability)
            if user_input.cross_method == "Three Point Cross":
                cross_result = cross.crossover3(p1, p2, user_input.cross_probability)
            if user_input.cross_method == "Uniform Cross":
                cross_result = cross.uniformCrossover(p1, p2)
            for c in cross_result:
                if user_input.mutation_method == "Edge":
                    c = mutation.edge_mutation(c, user_input.mutation_probability)
                if user_input.mutation_method == "One Point":
                    c = mutation.op_mutation(c, user_input.mutation_probability)
                if user_input.mutation_method == "Two Point":
                    c = mutation.tp_mutation(c, user_input.mutation_probability)
                c = inversion(c, user_input.inversion_probability)
                children.append(c)
        pop = children

    print('Done!')
    decoded = decode(bounds, user_input.number_of_bits, best)
    print('f(%s) = %f' % (decoded, best_eval))
    return [decoded, best_eval, gen_b_rows, gen_avg_rows, gen_std_dev_rows]

defaultUser = UserInputs(
    begin_range_a=-4.5,
    end_range_b=4.5,
    population_amount=100,
    number_of_bits=16,
    epochs_amount=1000,
    best_and_tournament_chromosome_amount=3,
    cross_probability=0.5,
    mutation_probability=0.1,
    inversion_probability=0.1,
    selection_method="Best",
    cross_method="cross",
    mutation_method="mutate",
    maximum=False,
    elite_strategy=True,
)


import os # Buat bersihin layar
import math # yang dipake math.sqrt buat ngitung nilai akar
import copy # Buat bikin duplikat data objek
from random import randint # Buat cari nilai random kelompok
import pandas as pd # Buat baca data dari excel

# Clear layar terminal
os.system('cls' if os.name == 'nt' else 'clear')

# Inisialisasi jumlah kelompok, f, dan nilai threshold
k = int(input("Masukkan jumlah kelompok: "))
f = 0
t = float(input("Masukkan nilai threshold: "))

# Clear layar terminal
os.system('cls' if os.name == 'nt' else 'clear')

initialData = [] # Tempat menyimpan data iterasi 0 / data awal
featureNames = [] # Tempat menyimpan daftar nama fitur untuk ditampilkan

# Buat fungsi untuk baca data dari excel
def readData(filename):
    data = pd.read_excel(filename)
    data.drop(data.columns[0], axis=1, inplace=True) # Data kolom pertama tidak dibaca karena isinya data nomor urutan data

    # Inisialisasi isi data awal
    for i in range(len(data)):
        initialData.append({
            "features": [],
            "cluster": randint(1, k) # Memasukkan setiap data ke kelompok secara acak
        })

    # Mengisi data awal dari excel
    for col in data.columns:
        featureNames.append(col)
        for i, val in enumerate(data[col]):
            initialData[i]["features"].append(val)

# readData("random/kmeans.xlsx")
# readData("random/doc1.xlsx")
# readData("random/doc2.xlsx")
# readData("random/doc3.xlsx")
# readData("random/doc4.xlsx")
# readData("random/doc5.xlsx")
# readData("random/doc6.xlsx")
# readData("random/doc7.xlsx")
# readData("random/doc8.xlsx")
readData("random/doc9.xlsx")

# Menghitung centroid setiap kelompok
def calculateCentroid(data):
    # Mencari nilai kfx
    kfxs = []
    for i in range(k):
        # Inisialisasi data kfx
        kfxs.append({
            "total": 0,
            # Tempat menyimpan total data di setiap kelompok
            "values": []
            # Tempat menyimpan jumlah nilai fitur di setiap kelompok
        })

        for _ in range(len(featureNames)):
            kfxs[i]["values"].append(0)

    # Menghitung total data di setiap kelompok
    for d in data:
        kfxs[d["cluster"] - 1]["total"] += 1

    # Menghitung jumlah nilai fitur di setiap kelompok
    for fIndex in range(len(featureNames)):
        for d in data:
            kfxs[d["cluster"] - 1]["values"][fIndex] += d["features"][fIndex]

    # Menghitung centroid
    centroids = []
    for kIndex, kfx in enumerate(kfxs):
        centroids.append([])
        for fIndex in range(len(featureNames)):
            if(kfx["total"] > 0):
                centroids[kIndex].append(kfx["values"][fIndex] / kfx["total"])
            else:
                centroids[kIndex].append(0)

    print("kfxs:\n", kfxs)
    print("centroids:\n", centroids)

    return centroids

# Menghitung jarak data ke centroid
def measureDistance(centroids, data):
    distances = []
    for dIndex, d in enumerate(data):
        distances.append([])
        # Menghitung jarak menggunakan rumus euclidean
        for c in centroids:
            distance = 0
            for xIndex, x in enumerate(d["features"]):
                distance += pow(abs(x - c[xIndex]), 2)
            distance = math.sqrt(distance)
            distances[dIndex].append(distance)
    print("Jarak:\n", distances)
    return distances

# Memindahkan setiap data ke kelompok dengan centroid terdekat
def clusterCheck(distances, data):
    newData = copy.deepcopy(data)
    isChanged = False
    for dataIndex, d in enumerate(newData):
        for distanceIndex, di in enumerate(distances[dataIndex]):
            if di < distances[dataIndex][d["cluster"] - 1]:
                d["cluster"] = distanceIndex + 1
                isChanged = True
    return [isChanged, newData]

# Menghitung fungsi f (objektif)
def calculateDeltaF(distances, data):
    oldF = f
    newF = 0
    for dIndex, d in enumerate(data):
        newF += distances[dIndex][d["cluster"] - 1]
    deltaf = abs(newF - oldF)
    return [deltaf, newF]

# Membuat fungsi untuk menampilkan data
def displayResult(data):
    dash = "---------"
    for x in range(len(featureNames) + 1):
        dash += "----------------"
    print(dash)

    displayHeader = "| Data\t|"
    for x in featureNames:
        displayHeader = displayHeader + " " + x + "\t|"
    displayHeader = displayHeader + " Kelompok\t|"
    print(displayHeader)
    print(dash)

    for index, i in enumerate(data):
        displayData = "| " + str(index + 1) + "\t"
        for x in range(len(featureNames)):
            displayData = displayData + "| " + str(i["features"][x]) + ("\t\t" if len(str(i["features"][x])) <= 5 else "\t")
        displayData += "| " + str(i["cluster"]) + "\t\t|"
        print(displayData)
    print(dash)

print("k = ", k)
print("f = ", f)
print("t = ", t)

print("\nData awal")
displayResult(initialData)
print("\n")

stop = False
oldData = initialData
newData = initialData
iterasi = 1
while (not stop):
    centroids = calculateCentroid(oldData)
    distances = measureDistance(centroids, oldData)
    newData = clusterCheck(distances, oldData)
    fdelta = calculateDeltaF(distances, oldData)
    oldData = newData[1]

    print("Iterasi " + str(iterasi))
    displayResult(oldData)
    print("Delta F\t: " + str(fdelta[0]))
    print("\n")

    if fdelta[0] < t or not newData[0]:
        stop = True

    f = fdelta[1]
    iterasi += 1

print("Hasil")
displayResult(newData[1])


from copy import deepcopy
import logging
from pathlib import Path
from typing import Dict, List, Tuple
import json_tricks

import numpy as np
from torch import Tensor
import torch
from torch.nn import Module

from ........nni.algorithms.compression.v2.pytorch.base import Task, TaskResult
from ........nni.algorithms.compression.v2.pytorch.utils import (
    config_list_canonical,
    compute_sparsity,
    get_model_weights_numel
)
from .base import TaskGenerator

_logger = logging.getLogger(__name__)

class FunctionBasedTaskGenerator(TaskGenerator):
    def __init__(self, total_iteration: int, origin_model: Module, origin_config_list: List[Dict],
                 origin_masks: Dict[str, Dict[str, Tensor]] = {}, log_dir: str = '.', keep_intermediate_result: bool = False):
        """
        Parameters
        ----------
        total_iteration
            The total iteration number.
        origin_model
            The origin unwrapped pytorch model to be pruned.
        origin_config_list
            The origin config list provided by the user. Note that this config_list is directly config the origin model.
            This means the sparsity provided by the origin_masks should also be recorded in the origin_config_list.
        origin_masks
            The pre masks on the origin model. This mask maybe user-defined or maybe generate by previous pruning.
        log_dir
            The log directory use to saving the task generator log.
        keep_intermediate_result
            If keeping the intermediate result, including intermediate model and masks during each iteration.
        """
        self.total_iteration = total_iteration
        super().__init__(origin_model, origin_config_list=origin_config_list, origin_masks=origin_masks,
                         log_dir=log_dir, keep_intermediate_result=keep_intermediate_result)

    def reset(self, model: Module, config_list: List[Dict] = [], masks: Dict[str, Dict[str, Tensor]] = {}):
        self.current_iteration = 0
        self.target_sparsity = config_list_canonical(model, config_list)
        super().reset(model, config_list=config_list, masks=masks)

    def init_pending_tasks(self) -> List[Task]:
        origin_model = torch.load(self._origin_model_path)
        origin_masks = torch.load(self._origin_masks_path)

        task_result = TaskResult('origin', origin_model, origin_masks, origin_masks, None)

        return self.generate_tasks(task_result)

    def generate_tasks(self, task_result: TaskResult) -> List[Task]:
        compact_model = task_result.compact_model
        compact_model_masks = task_result.compact_model_masks

        # save intermediate result
        model_path = Path(self._intermediate_result_dir, '{}_compact_model.pth'.format(task_result.task_id))
        masks_path = Path(self._intermediate_result_dir, '{}_compact_model_masks.pth'.format(task_result.task_id))
        torch.save(compact_model, model_path)
        torch.save(compact_model_masks, masks_path)

        # get current2origin_sparsity and compact2origin_sparsity
        origin_model = torch.load(self._origin_model_path)
        current2origin_sparsity, compact2origin_sparsity, _ = compute_sparsity(origin_model, compact_model, compact_model_masks, self.target_sparsity)
        _logger.info('\nTask %s total real sparsity compared with original model is:\n%s', str(task_result.task_id), json_tricks.dumps(current2origin_sparsity, indent=4))
        if task_result.task_id != 'origin':
            self._tasks[task_result.task_id].state['current2origin_sparsity'] = current2origin_sparsity

        # if reach the total_iteration, no more task will be generated
        if self.current_iteration > self.total_iteration:
            return []

        task_id = self._task_id_candidate
        new_config_list = self.generate_config_list(self.target_sparsity, self.current_iteration, compact2origin_sparsity)
        new_config_list = self.allocate_sparsity(new_config_list, compact_model, compact_model_masks)
        config_list_path = Path(self._intermediate_result_dir, '{}_config_list.json'.format(task_id))

        with Path(config_list_path).open('w') as f:
            json_tricks.dump(new_config_list, f, indent=4)
        task = Task(task_id, model_path, masks_path, config_list_path)

        self._tasks[task_id] = task

        self._task_id_candidate += 1
        self.current_iteration += 1

        return [task]

    def generate_config_list(self, target_sparsity: List[Dict], iteration: int, compact2origin_sparsity: List[Dict]) -> List[Dict]:
        raise NotImplementedError()

    def allocate_sparsity(self, new_config_list: List[Dict], model: Module, masks: Dict[str, Dict[str, Tensor]]):
        return new_config_list

class AGPTaskGenerator(FunctionBasedTaskGenerator):
    def generate_config_list(self, target_sparsity: List[Dict], iteration: int, compact2origin_sparsity: List[Dict]) -> List[Dict]:
        config_list = []
        for target, mo in zip(target_sparsity, compact2origin_sparsity):
            ori_sparsity = (1 - (1 - iteration / self.total_iteration) ** 3) * target['total_sparsity']
            sparsity = max(0.0, (ori_sparsity - mo['total_sparsity']) / (1 - mo['total_sparsity']))
            assert 0 <= sparsity <= 1, 'sparsity: {}, ori_sparsity: {}, model_sparsity: {}'.format(sparsity, ori_sparsity, mo['total_sparsity'])
            config_list.append(deepcopy(target))
            config_list[-1]['total_sparsity'] = sparsity
        return config_list

class LinearTaskGenerator(FunctionBasedTaskGenerator):
    def generate_config_list(self, target_sparsity: List[Dict], iteration: int, compact2origin_sparsity: List[Dict]) -> List[Dict]:
        config_list = []
        for target, mo in zip(target_sparsity, compact2origin_sparsity):
            ori_sparsity = iteration / self.total_iteration * target['total_sparsity']
            sparsity = max(0.0, (ori_sparsity - mo['total_sparsity']) / (1 - mo['total_sparsity']))
            assert 0 <= sparsity <= 1, 'sparsity: {}, ori_sparsity: {}, model_sparsity: {}'.format(sparsity, ori_sparsity, mo['total_sparsity'])
            config_list.append(deepcopy(target))
            config_list[-1]['total_sparsity'] = sparsity
        return config_list

class LotteryTicketTaskGenerator(FunctionBasedTaskGenerator):
    def reset(self, model: Module, config_list: List[Dict] = [], masks: Dict[str, Dict[str, Tensor]] = {}):
        self.current_iteration = 1
        self.target_sparsity = config_list_canonical(model, config_list)
        super(FunctionBasedTaskGenerator, self).reset(model, config_list=config_list, masks=masks)

    def generate_config_list(self, target_sparsity: List[Dict], iteration: int, compact2origin_sparsity: List[Dict]) -> List[Dict]:
        config_list = []
        for target, mo in zip(target_sparsity, compact2origin_sparsity):
            # NOTE: The ori_sparsity calculation formula in compression v1 is as follow, it is different from the paper.
            # But the formula in paper will cause numerical problems, so keep the formula in compression v1.
            ori_sparsity = 1 - (1 - target['total_sparsity']) ** (iteration / self.total_iteration)
            # The following is the formula in paper.
            # ori_sparsity = (target['total_sparsity'] * 100) ** (iteration / self.total_iteration) / 100
            sparsity = max(0.0, (ori_sparsity - mo['total_sparsity']) / (1 - mo['total_sparsity']))
            assert 0 <= sparsity <= 1, 'sparsity: {}, ori_sparsity: {}, model_sparsity: {}'.format(sparsity, ori_sparsity, mo['total_sparsity'])
            config_list.append(deepcopy(target))
            config_list[-1]['total_sparsity'] = sparsity
        return config_list

class SimulatedAnnealingTaskGenerator(TaskGenerator):
    def __init__(self, origin_model: Module, origin_config_list: List[Dict], origin_masks: Dict[str, Dict[str, Tensor]] = {},
                 start_temperature: float = 100, stop_temperature: float = 20, cool_down_rate: float = 0.9,
                 perturbation_magnitude: float = 0.35, log_dir: str = '.', keep_intermediate_result: bool = False):
        """
        Parameters
        ----------
        origin_model
            The origin unwrapped pytorch model to be pruned.
        origin_config_list
            The origin config list provided by the user. Note that this config_list is directly config the origin model.
            This means the sparsity provided by the origin_masks should also be recorded in the origin_config_list.
        origin_masks
            The pre masks on the origin model. This mask maybe user-defined or maybe generate by previous pruning.
        start_temperature
            Start temperature of the simulated annealing process.
        stop_temperature
            Stop temperature of the simulated annealing process.
        cool_down_rate
            Cool down rate of the temperature.
        perturbation_magnitude
            Initial perturbation magnitude to the sparsities. The magnitude decreases with current temperature.
        log_dir
            The log directory use to saving the task generator log.
        keep_intermediate_result
            If keeping the intermediate result, including intermediate model and masks during each iteration.
        """
        self.start_temperature = start_temperature
        self.stop_temperature = stop_temperature
        self.cool_down_rate = cool_down_rate
        self.perturbation_magnitude = perturbation_magnitude

        super().__init__(origin_model, origin_masks=origin_masks, origin_config_list=origin_config_list,
                         log_dir=log_dir, keep_intermediate_result=keep_intermediate_result)

    def reset(self, model: Module, config_list: List[Dict] = [], masks: Dict[str, Dict[str, Tensor]] = {}):
        self.current_temperature = self.start_temperature

        self.weights_numel, self.masked_rate = get_model_weights_numel(model, config_list, masks)
        self.target_sparsity_list = config_list_canonical(model, config_list)
        self._adjust_target_sparsity()

        self._temp_config_list = None
        self._current_sparsity_list = None
        self._current_score = None

        super().reset(model, config_list=config_list, masks=masks)

    def _adjust_target_sparsity(self):
        """
        If origin_masks is not empty, then re-scale the target sparsity.
        """
        if len(self.masked_rate) > 0:
            for config in self.target_sparsity_list:
                sparsity, op_names = config['total_sparsity'], config['op_names']
                remaining_weight_numel = 0
                pruned_weight_numel = 0
                for name in op_names:
                    remaining_weight_numel += self.weights_numel[name]
                    if name in self.masked_rate and self.masked_rate[name] != 0:
                        pruned_weight_numel += 1 / (1 / self.masked_rate[name] - 1) * self.weights_numel[name]
                total_mask_rate = pruned_weight_numel / (pruned_weight_numel + remaining_weight_numel)
                config['total_sparsity'] = max(0, (sparsity - total_mask_rate) / (1 - total_mask_rate))

    def _init_temp_config_list(self):
        self._temp_config_list = []
        self._temp_sparsity_list = []
        for config in self.target_sparsity_list:
            sparsity_config, sparsity = self._init_config_sparsity(config)
            self._temp_config_list.extend(sparsity_config)
            self._temp_sparsity_list.append(sparsity)

    def _init_config_sparsity(self, config: Dict) -> Tuple[List[Dict], List]:
        assert 'total_sparsity' in config, 'Sparsity must be set in config: {}'.format(config)
        target_sparsity = config['total_sparsity']
        op_names = config['op_names']

        if target_sparsity == 0:
            return [], []

        while True:
            random_sparsity = sorted(np.random.uniform(0, 1, len(op_names)))
            rescaled_sparsity = self._rescale_sparsity(random_sparsity, target_sparsity, op_names)
            if rescaled_sparsity is not None and rescaled_sparsity[0] >= 0 and rescaled_sparsity[-1] < 1:
                break

        return self._sparsity_to_config_list(rescaled_sparsity, config), rescaled_sparsity

    def _rescale_sparsity(self, random_sparsity: List, target_sparsity: float, op_names: List) -> List:
        assert len(random_sparsity) == len(op_names)

        num_weights = sorted([self.weights_numel[op_name] for op_name in op_names])
        sparsity = sorted(random_sparsity)

        total_weights = 0
        total_weights_pruned = 0

        # calculate the scale
        for idx, num_weight in enumerate(num_weights):
            total_weights += num_weight
            total_weights_pruned += int(num_weight * sparsity[idx])
        if total_weights_pruned == 0:
            return None

        scale = target_sparsity / (total_weights_pruned / total_weights)

        # rescale the sparsity
        sparsity = np.asarray(sparsity) * scale
        return sparsity

    def _sparsity_to_config_list(self, sparsity: List, config: Dict) -> List[Dict]:
        sparsity = sorted(sparsity)
        op_names = [k for k, _ in sorted(self.weights_numel.items(), key=lambda item: item[1]) if k in config['op_names']]
        assert len(sparsity) == len(op_names)
        return [{'total_sparsity': sparsity, 'op_names': [op_name]} for sparsity, op_name in zip(sparsity, op_names)]

    def _update_with_perturbations(self):
        self._temp_config_list = []
        self._temp_sparsity_list = []
        # decrease magnitude with current temperature
        magnitude = self.current_temperature / self.start_temperature * self.perturbation_magnitude
        for config, current_sparsity in zip(self.target_sparsity_list, self._current_sparsity_list):
            if len(current_sparsity) == 0:
                self._temp_sparsity_list.append([])
                continue
            while True:
                perturbation = np.random.uniform(-magnitude, magnitude, len(current_sparsity))
                temp_sparsity = np.clip(0, current_sparsity + perturbation, None)
                temp_sparsity = self._rescale_sparsity(temp_sparsity, config['total_sparsity'], config['op_names'])
                if temp_sparsity is not None and temp_sparsity[0] >= 0 and temp_sparsity[-1] < 1:
                    self._temp_config_list.extend(self._sparsity_to_config_list(temp_sparsity, config))
                    self._temp_sparsity_list.append(temp_sparsity)
                    break

    def _recover_real_sparsity(self, config_list: List[Dict]) -> List[Dict]:
        """
        If the origin masks is not None, then the sparsity in new generated config_list need to be rescaled.
        """
        for config in config_list:
            assert len(config['op_names']) == 1
            op_name = config['op_names'][0]
            if op_name in self.masked_rate:
                config['total_sparsity'] = self.masked_rate[op_name] + config['total_sparsity'] * (1 - self.masked_rate[op_name])
        return config_list

    def init_pending_tasks(self) -> List[Task]:
        origin_model = torch.load(self._origin_model_path)
        origin_masks = torch.load(self._origin_masks_path)

        self.temp_model_path = Path(self._intermediate_result_dir, 'origin_compact_model.pth')
        self.temp_masks_path = Path(self._intermediate_result_dir, 'origin_compact_model_masks.pth')
        torch.save(origin_model, self.temp_model_path)
        torch.save(origin_masks, self.temp_masks_path)

        task_result = TaskResult('origin', origin_model, origin_masks, origin_masks, None)

        return self.generate_tasks(task_result)

    def generate_tasks(self, task_result: TaskResult) -> List[Task]:
        # initial/update temp config list
        if self._temp_config_list is None:
            self._init_temp_config_list()
        else:
            score = self._tasks[task_result.task_id].score
            if self._current_sparsity_list is None:
                self._current_sparsity_list = deepcopy(self._temp_sparsity_list)
                self._current_score = score
            else:
                delta_E = np.abs(score - self._current_score)
                probability = np.exp(-1 * delta_E / self.current_temperature)
                if self._current_score < score or np.random.uniform(0, 1) < probability:
                    self._current_score = score
                    self._current_sparsity_list = deepcopy(self._temp_sparsity_list)
                    self.current_temperature *= self.cool_down_rate
            if self.current_temperature < self.stop_temperature:
                return []
            self._update_with_perturbations()

        task_id = self._task_id_candidate
        new_config_list = self._recover_real_sparsity(deepcopy(self._temp_config_list))
        config_list_path = Path(self._intermediate_result_dir, '{}_config_list.json'.format(task_id))

        with Path(config_list_path).open('w') as f:
            json_tricks.dump(new_config_list, f, indent=4)

        task = Task(task_id, self.temp_model_path, self.temp_masks_path, config_list_path)

        self._tasks[task_id] = task

        self._task_id_candidate += 1

        return [task]