import os, shutil,math,warnings,time
from glob import glob
from tqdm import tqdm
import json
import pycrs
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import fiona
from fiona.crs import from_epsg
from PIL import Image, ImageOps
from shapely import wkt
from shapely.geometry import box,Point
import pandas as pd
import geopandas as gpd
import rasterio
import rasterio.mask
from rasterio.plot import show
from rasterio.enums import Resampling
from rasterio.warp import calculate_default_transform, reproject, Resampling
from sklearn.decomposition import PCA
from sklearn import metrics
from sklearn.mixture import GaussianMixture
from sklearn.cluster import KMeans
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.neighbors import KNeighborsClassifier

warnings.simplefilter(action='ignore', category=FutureWarning)

class AnalysisTool:
    def __init__(self):
        pass

    def train_point_handler(self, shapefile_path, image_path, array_list=[],raster=None,num_points=1000):

        gdf = gpd.read_file(shapefile_path)

        values = []
        training_points = []

        if len(array_list)==0:
            raster = rasterio.open(image_path)
            array_list=[raster.read(i+1) for i in range(counts)]
        else:
            raster=raster
            array_list=array_list

        counts = raster.meta['count']

        while len(training_points) < num_points:
            idx = np.random.randint(0, gdf['geometry'].count())
            bounds = gdf[gdf['geometry'].index == idx].total_bounds

            xmin = bounds[0]
            xmax = bounds[2]
            ymin = bounds[1]
            ymax = bounds[3]

            x = np.random.uniform(xmin, xmax)
            y = np.random.uniform(ymin, ymax)
            point = Point(x, y)

            if len(gdf[gdf['geometry'].index == idx].loc[gdf[
                    gdf['geometry'].index == idx].intersects(
                        point)].index) == 0:
                pass
            else:
                val_list = []
                for i in range(counts):
                    #array = raster.read(i + 1)
                    array=array_list[i]
                    val = array[raster.index(x, y)]
                    val_list.append(val)
                if np.isnan(val_list).any() == True:
                    pass
                else:
                    values.append(val_list)
                    training_points.append(point)

        del raster

        return values, training_points

    def train_data_process(self,
                           shapefile_dir_path,
                           image_path=None,
                           num_points=1000,
                           show_spectrum=False,
                           band_order=[],
                           write_train_data=False,
                           output_train_data_path=None):

        raster = rasterio.open(image_path)
        band_names=list(raster.descriptions)
        output_train_data_epsg=raster.meta['crs']
        counts=raster.meta['count']
        array_list=[raster.read(i+1) for i in range(counts)]

        train_point_handler = self.train_point_handler
        shapefile_list = glob(shapefile_dir_path + '/*.shp')

        extracted_values = []
        lulc_class = []
        lulc_name = []
        lat = []
        lon = []

        cnt = 1
        for shapefile in tqdm(shapefile_list):
            filename = os.path.basename(shapefile).replace('.shp', '')

            values, training_point = train_point_handler(
                shapefile, image_path, num_points=num_points,array_list=array_list,raster=raster)
            extracted_values.extend(values)

            lat.extend([point.x for point in training_point])
            lon.extend([point.y for point in training_point])

            lulc_name.extend([filename for i in range(len(values))])
            lulc_class.extend([cnt for i in range(len(values))])
            cnt += 1

        output_df = pd.DataFrame({
            'lulc_name': lulc_name,
            'lulc_class': lulc_class,
            'extracted_values': extracted_values,
            'Latitude': lat,
            'Longitude': lon
        })

        output_gdf = gpd.GeoDataFrame(output_df,
                                      geometry=gpd.points_from_xy(
                                          output_df.Latitude,
                                          output_df.Longitude))
        output_gdf.set_crs(output_train_data_epsg, inplace=True)

        if show_spectrum==True:
            fig,ax=plt.subplots(figsize=(20,10))
            if len(band_order)==0:
                spectrum_axis=band_names
            else:
                spectrum_axis=band_order
            classes=set(extracted_data_gdf['lulc_name'])
            for cla in classes:
                extracted_data=extracted_data_gdf[extracted_data_gdf['lulc_name']==cla]
                values=np.array(list(extracted_data['extracted_values']))
                value_list=[]
                for band in spectrum_axis:
                    idx=band_names.index(band)
                    mean_val=np.mean(values[:,idx:idx+1])
                    value_list.append(mean_val)
                ax.plot(spectrum_axis,value_list,label=cla)
            ax.legend()

        if write_train_data == True:
            output_gdf.drop('extracted_values',
                            axis=1).to_file(output_train_data_path)
        else:
            pass

        return output_gdf

    def pca_image(self,
                  image_path,
                  RGBoption=False,
                  return_arrays=False,
                  n_components=10,
                  exception=[],
                  nanvalue=np.nan,
                  image_save_path=None,
                  image_name='pca.png'):
        raster = rasterio.open(image_path)
        width = raster.meta["width"]
        height = raster.meta["height"]
        shape = (height, width)

        counts = raster.meta['count']
        bands_opened = []
        cnt = 1
        for i in range(counts):
            if (cnt in exception) == True:
                pass
            else:
                array = raster.read(cnt).flatten()
                array = np.nan_to_num(array)
                bands_opened.append(array)
            cnt += 1

        bands_stacked = np.dstack(tuple(bands_opened))
        pca = PCA(n_components)
        pca.fit(bands_stacked[0])
        variance = pca.explained_variance_
        totalVariance = np.sum(variance)
        variance = variance / totalVariance * 100

        fig, ax = plt.subplots()

        ax.plot(variance)
        ax.set_xlabel("Components of PCA")
        ax.set_ylabel("Variance explained [%]")

        pca_image = pca.transform(bands_stacked[0])

        first_component = self.nan_mask(image_path,
                                        pca_image[:,
                                                  0].reshape(shape), nanvalue)
        second_component = self.nan_mask(image_path,
                                         pca_image[:,
                                                   1].reshape(shape), nanvalue)
        third_component = self.nan_mask(image_path,
                                        pca_image[:,
                                                  2].reshape(shape), nanvalue)

        if RGBoption == True:
            self.visualize([0, 1, 2],
                           array=np.stack([
                               first_component, second_component,
                               third_component
                           ]),
                           image_save_path=image_save_path,
                           image_name=image_name)
        else:
            first_component = (255 / np.nanmax(first_component) *
                               first_component).astype(np.uint8)
            image = Image.fromarray(first_component)
            image = ImageOps.equalize(image)
            image.show()
            if image_save_path != None:
                image.save(image_save_path + '/' + image_name)
            else:
                pass

        if return_arrays == True:
            return [
                pca_image[:, i].reshape(shape) for i in range(n_components)
            ]
        else:
            pass

    def visualize(self,
                  rgb_band,
                  array=None,
                  image_path=None,
                  nanvalue=np.nan,
                  image_save_path=None,
                  image_name='saved_image.png'):
        if image_path != None:
            raster = rasterio.open(image_path)
            R = raster.read(rgb_band[0])
            G = raster.read(rgb_band[1])
            B = raster.read(rgb_band[2])

            R = self.nan_mask(image_path, R, nanvalue)
            G = self.nan_mask(image_path, G, nanvalue)
            B = self.nan_mask(image_path, B, nanvalue)

            R = (255 / np.nanmax(R) * R).astype(np.uint8)
            G = (255 / np.nanmax(G) * G).astype(np.uint8)
            B = (255 / np.nanmax(B) * B).astype(np.uint8)

            image_stack = Image.fromarray(np.dstack((R, G, B)))
            image = ImageOps.equalize(image_stack)
            image.show()
        else:
            R = array[rgb_band[0]]
            G = array[rgb_band[1]]
            B = array[rgb_band[2]]

            R = (255 / np.nanmax(R) * R).astype(np.uint8)
            G = (255 / np.nanmax(G) * G).astype(np.uint8)
            B = (255 / np.nanmax(B) * B).astype(np.uint8)

            image_stack = Image.fromarray(np.dstack((R, G, B)))
            image = ImageOps.equalize(image_stack)
            image.show()

        if image_save_path != None:
            image.save(image_save_path + '/' + image_name)
        else:
            pass

    def un_supervised_classification(
            self,
            path_to_image,
            output_path,
            input_data=[],
            method='GaussianMixture',
            n=30,
            num_points=1000,
            preview=True,
            output_file_name='un_supervised_classification',
            params=None,
            nanvalue=np.nan,
            cmap='jet'):

        if len(input_data)!=0:
            input_data = input_data
            data = input_data['extracted_values']
        else:
            values = []
            raster = rasterio.open(path_to_image)
            counts = raster.meta['count']
            array_list=[raster.read(i+1) for i in range(counts)]
            bounds = raster.bounds
            xmin = bounds[0]
            xmax = bounds[2]
            ymin = bounds[1]
            ymax = bounds[3]

            cnt=0
            while len(values) < num_points:
                x = np.random.uniform(xmin, xmax)
                y = np.random.uniform(ymin, ymax)
                val_list = []
                for i in range(counts):
                    array=array_list[i]
                    val = array[raster.index(x, y)]
                    val_list.append(val)
                if np.isnan(val_list).any() == True:
                    pass
                else:
                    values.append(val_list)
            del raster
            data=values

        data = [spectrum for spectrum in data]
        X = np.array(data)
        print(X.shape)

        if method == 'GaussianMixture':
            if params != None:
                model = GaussianMixture(**params).fit(X)
            else:
                model = GaussianMixture(n_components=n).fit(X)
        elif method == 'KMeans':
            if params != None:
                model = KMeans(**params).fit(X)
            else:
                model = KMeans(n_clusters=n).fit(X)
        else:
            print(
                'The available methods are either GaussianMixture or KMeans.')

        raster = rasterio.open(path_to_image)
        counts = raster.meta['count']

        image_array = []
        for i in range(counts):
            image_array.append(raster.read(i + 1))

        cluster_output = []

        for i in tqdm(range(image_array[0].shape[0])):
            spectrum_data_at_row_i = []
            spectrum_data_at_row_i_nan_check = []
            for m in range(counts):
                spectrum_data_at_row_i.append(np.nan_to_num(image_array[m][i]))
                spectrum_data_at_row_i_nan_check.append(image_array[m][i])

            spectrum_data_at_row_i_T = np.transpose(
                np.array(spectrum_data_at_row_i))
            spectrum_data_at_row_i_nan_check_T = np.transpose(
                np.array(spectrum_data_at_row_i_nan_check))

            nan_ind = np.where(np.isnan(spectrum_data_at_row_i_nan_check_T))[0]
            classified = np.transpose(
                model.predict(spectrum_data_at_row_i_T)).astype(np.float)
            classified[nan_ind] = nanvalue
            cluster_output.append(classified)

        output_image = np.array(cluster_output)

        output_tif_name = output_file_name + '_' + method + '.tif'

        with rasterio.open(output_path + '/' + output_tif_name,
                           'w',
                           height=output_image.shape[0],
                           width=output_image.shape[1],
                           driver='GTiff',
                           count=1,
                           crs=raster.crs,
                           transform=raster.transform,
                           dtype=output_image.dtype) as output:
            output.write(output_image, 1)
            output.close()

        if preview == True:
            fig, ax = plt.subplots(figsize=(15, 15))
            rasterio.plot.show(output_image, ax=ax, cmap=cmap)
        else:
            pass

        del raster

    def supervised_classification_test(self,
                                      train_data_path=None,
                                      input_data=[],
                                      method='RandomForest',
                                      num_points=1000,
                                      test_size=0.2,
                                      params={}):

        if len(input_data)!=0:
            input_data = input_data
        else:
            input_data= self.train_data_process(train_data_path, path_to_image,num_points)

        data = input_data['extracted_values']
        data = [spectrum for spectrum in data]
        label = input_data['lulc_class']
        label = [l for l in label]

        X = np.array(data)
        y = np.array(label)

        X_train, X_test, y_train, y_test = train_test_split(X,
                                                            y,
                                                            test_size=0.2)

        if method == 'RandomForest':
            model = RandomForestClassifier(**params)
        elif method == 'NeuralNetwork':
            model = MLPClassifier(**params)
        elif method == 'KNeighbors':
            model = KNeighborsClassifier(**params)
        elif method == 'GradientBoosting':
            model = GradientBoostingClassifier(**params)
        else:
            print(
                'The available methods are either RandomForest, NeuralNetwork, KNeighbors, or GradientBoosting.'
            )

        model.fit(X_train, y_train)

        confusion_matrix = metrics.confusion_matrix(y_test,
                                                    model.predict(X_test))
        accuracy_train = metrics.accuracy_score(y_train, model.predict(X_train))
        accuracy_test = metrics.accuracy_score(y_test, model.predict(X_test))
        kappa = metrics.cohen_kappa_score(y_test, model.predict(X_test))

        print(f'Confusion matrix is:')
        print(confusion_matrix)
        print(f'Accuracy on train {accuracy_train:.5f}')
        print(f'Accuracy on test {accuracy_test:.5f}')
        print(f'Kappa coefficient {kappa:.5f}')

    def supervised_classification(self,
                                  path_to_image,
                                  output_path,
                                  train_data_path=None,
                                  input_data=[],
                                  method='RandomForest',
                                  num_points=1000,
                                  preview=True,
                                  output_file_name='supervised_classification',
                                  test_size=0.2,
                                  params={},
                                  nanvalue=np.nan,
                                  cmap='jet'):

        if len(input_data)!=0:
            input_data = input_data
        else:
            input_data= self.train_data_process(train_data_path, path_to_image,num_points)

        data = input_data['extracted_values']
        data = [spectrum for spectrum in data]
        label = input_data['lulc_class']
        label = [l for l in label]

        X = np.array(data)
        y = np.array(label)

        X_train, X_test, y_train, y_test = train_test_split(X,
                                                            y,
                                                            test_size=0.2)

        if method == 'RandomForest':
            model = RandomForestClassifier(**params)
        elif method == 'NeuralNetwork':
            model = MLPClassifier(**params)
        elif method == 'KNeighbors':
            model = KNeighborsClassifier(**params)
        elif method == 'GradientBoosting':
            model = GradientBoostingClassifier(**params)
        else:
            print(
                'The available methods are either RandomForest, NeuralNetwork, Kneighbors, or GradientBoosting.'
            )

        model.fit(X_train, y_train)

        confusion_matrix = metrics.confusion_matrix(y_test,
                                                    model.predict(X_test))
        accuracy = metrics.accuracy_score(y_test, model.predict(X_test))
        kappa = metrics.cohen_kappa_score(y_test, model.predict(X_test))

        print(f'Confusion matrix is:')
        print(confusion_matrix)
        print(f'Accuracy {accuracy:.5f}')
        print(f'Kappa coefficient {kappa:.5f}')

        raster = rasterio.open(path_to_image)
        counts = raster.meta['count']

        image_array = []
        for i in range(counts):
            image_array.append(raster.read(i + 1))

        classified_output = []

        for i in tqdm(range(image_array[0].shape[0])):
            spectrum_data_at_row_i = []
            spectrum_data_at_row_i_nan_check = []
            for m in range(counts):
                spectrum_data_at_row_i.append(np.nan_to_num(image_array[m][i]))
                spectrum_data_at_row_i_nan_check.append(image_array[m][i])

            spectrum_data_at_row_i_T = np.transpose(
                np.array(spectrum_data_at_row_i))
            spectrum_data_at_row_i_nan_check_T = np.transpose(
                np.array(spectrum_data_at_row_i_nan_check))

            nan_ind = np.where(np.isnan(spectrum_data_at_row_i_nan_check_T))[0]

            classified = np.transpose(
                model.predict(spectrum_data_at_row_i_T)).astype(np.float)
            classified[nan_ind] = nanvalue
            classified_output.append(classified)

        output_image = np.array(classified_output)
        output_tif_name = output_file_name + '_' + method + '.tif'

        with rasterio.open(output_path + '/' + output_tif_name,
                           'w',
                           height=output_image.shape[0],
                           width=output_image.shape[1],
                           driver='GTiff',
                           count=1,
                           crs=raster.crs,
                           transform=raster.transform,
                           dtype=output_image.dtype) as output:
            output.write(output_image, 1)
            output.close()

        if preview == True:
            fig, ax = plt.subplots(figsize=(15, 15))
            rasterio.plot.show(output_image, ax=ax, cmap=cmap)
        else:
            pass

        del raster

    def nan_mask(self, path_to_image, input_array, nanvalue):
        raster = rasterio.open(path_to_image)
        counts = raster.meta['count']

        image_array = []
        for i in range(counts):
            image_array.append(raster.read(i + 1))

        nan_mask = []

        for i in range(image_array[0].shape[0]):
            spectrum_data_at_row_i_nan_check = []
            mask = np.zeros(image_array[0].shape[1])

            for m in range(counts):
                spectrum_data_at_row_i_nan_check.append(image_array[m][i])

            spectrum_data_at_row_i_nan_check_T = np.transpose(
                np.array(spectrum_data_at_row_i_nan_check))

            nan_ind = np.where(np.isnan(spectrum_data_at_row_i_nan_check_T))[0]

            mask[nan_ind] = 1
            nan_mask.append(mask)

        nan_mask = np.array(nan_mask)

        idx = np.where(nan_mask == 1)
        input_array[idx] = nanvalue

        return input_array

    def create_composite(self,
                         path_to_image_folder,
                         output_path,
                         singleband=True,
                         merge_method='mean',
                         merge_image_name='merged_by_'):
        file_list = glob(path_to_image_folder + '/*')
        if singleband == True:
            with rasterio.open(file_list[0]) as opened:
                ref = opened
                crs=opened.meta['crs']
                shape = opened.read(1).shape
            opened_array_list = []
            for file in file_list:
                with rasterio.open(file) as opened:
                    opened_array = opened.read(1).flatten()
                    opened_array_list.append(opened_array)
            if merge_method == 'mean':
                output_array = np.nanmean(opened_array_list, axis=0)
            elif merge_method == 'max':
                output_array = np.nanmax(opened_array_list, axis=0)
            elif merge_method == 'min':
                output_array = np.nanmin(opened_array_list, axis=0)
            else:
                print('The available merging methods are mean, max, or min.')
            output_array = output_array.reshape(shape)
            with rasterio.open(output_path +
                               f"/{merge_image_name}{merge_method}.tif",
                               'w',
                               driver='GTiff',
                               width=output_array.shape[1],
                               height=output_array.shape[0],
                               count=1,
                               crs=crs,
                               transform=ref.transform,
                               dtype=output_array.dtype) as output:
                output.write(output_array, 1)
                output.close()
        else:
            with rasterio.open(file_list[0]) as opened:
                ref = opened
                shape = (opened.meta['height'], opened.meta['width'])
                count = opened.meta['count']
                descriptions = opened.descriptions
            band_array_list = []
            for i in range(count):
                opened_array_list = []
                for file in file_list:
                    with rasterio.open(file) as opened:
                        opened_array = opened.read(i + 1).flatten()
                        opened_array_list.append(opened_array)
                if merge_method == 'mean':
                    output_array = np.nanmean(opened_array_list, axis=0)
                elif merge_method == 'max':
                    output_array = np.nanmax(opened_array_list, axis=0)
                elif merge_method == 'min':
                    output_array = np.nanmin(opened_array_list, axis=0)
                else:
                    print(
                        'The available merging methods are mean, max, or min.')
                output_array = output_array.reshape(shape)
                band_array_list.append(output_array)
            with rasterio.open(output_path +
                               f"/{merge_image_name}{merge_method}.tif",
                               'w',
                               driver='GTiff',
                               width=shape[1],
                               height=shape[0],
                               count=count,
                               crs=crs,
                               transform=ref.transform,
                               dtype=output_array.dtype) as output:
                for i in range(count):
                    output.write(band_array_list[i], i + 1)
                    output.set_band_description(i + 1, descriptions[i])
                output.close()

    def clip(self, path_to_image, path_to_shapefile, path_to_output,nodata=np.nan):
        with fiona.open(path_to_shapefile, 'r') as shapefile:
            shapes = [feature['geometry'] for feature in shapefile]

        with rasterio.open(path_to_image) as src:
            out_image, out_transform = rasterio.mask.mask(src,
                                                          shapes,
                                                          crop=True,nodata=np.nan)
            out_meta = src.meta

        out_meta.update({
            'driver': 'GTiff',
            'height': out_image.shape[1],
            'width': out_image.shape[2],
            'transform': out_transform
        })

        filename=os.path.splitext(os.path.basename(path_to_image))[0]
        with rasterio.open(path_to_output+'/'+filename+'_clipped.tif','w',**out_meta) as opened:
            opened.write(out_image)

    def combine_bands(self,path_to_folder,path_to_output,output_filename='combined',remove=True,dtype=None):
        files=glob(path_to_folder+'/*')

        with rasterio.open(files[0]) as ref:
            crs=ref.meta['crs']
            width=ref.meta['width']
            height=ref.meta['height']
            transform=ref.transform
            if dtype==None:
                dtype=ref.read(1).dtype
            else:
                pass

        with rasterio.open(path_to_output+'/'+output_filename+'.tif',
                           'w',
                           driver='GTiff',
                           width=width,
                           height=height,
                           count=len(files),
                           crs=crs,
                           transform=transform,
                           dtype=dtype) as opened:
            cnt=1
            for file in files:
                with rasterio.open(file) as src:
                    output_raster=src.read(1)
                    src.close()
                opened.write(output_raster,cnt)
                band_name=os.path.splitext(os.path.basename(file))[0]
                opened.set_band_description(cnt,band_name)
                cnt+=1
            opened.close()

        if remove==True:
            for file in files:
                try:
                    os.remove(file)
                except:
                    time.sleep(3)
                    os.remove(file)

        else:
            pass

    def reprojection_to_crs(self,path_to_image,dst_crs='EPSG:4326',remove=False):
        epsg='_crs_'+dst_crs.replace('EPSG:','')
        dst_path=path_to_image.replace(os.path.basename(path_to_image),'')
        filename=os.path.splitext(os.path.basename(path_to_image))[0]
        with rasterio.open(path_to_image) as src:
            transform, width, height = calculate_default_transform(src.crs, dst_crs, src.width, src.height, *src.bounds)
            meta = src.meta.copy()
            meta.update({
                'crs': dst_crs,
                'transform': transform,
                'width': width,
                'height': height
            })

            with rasterio.open(dst_path+filename+epsg+'.tif', 'w', **meta) as dst:
                for i in range(1, src.count + 1):
                    reproject(
                        source=rasterio.band(src, i),
                        destination=rasterio.band(dst, i),
                        src_transform=src.transform,
                        src_crs=src.crs,
                        dst_transform=transform,
                        dst_crs=dst_crs,
                        resampling=Resampling.nearest)
        if remove==True:
            os.remove(path_to_image)
        else:
            pass

    def align_raster(self,path_to_ref_image,path_to_target_image,remove=True,nodata=np.nan):
        with rasterio.open(path_to_ref_image) as ref:
            bounds=ref.bounds
            height=ref.meta["height"]
            width=ref.meta["width"]
            crs=ref.meta["crs"]
            transform=ref.meta["transform"]

        bbox = box(bounds[0], bounds[1], bounds[2], bounds[3])
        geo = gpd.GeoDataFrame({'geometry': bbox}, index=[0], crs=crs)
        geo = geo.to_crs(crs=crs)
        coords=[json.loads(geo.to_json())['features'][0]['geometry']]

        dst_path=path_to_target_image.replace(os.path.basename(path_to_target_image),'')
        filename=os.path.splitext(os.path.basename(path_to_target_image))[0]
        os.mkdir(dst_path+'/temp')

        with rasterio.open(path_to_target_image) as src:
            out_img, out_transform = rasterio.mask.mask(src, coords, crop=True,filled=False,nodata=nodata)
            src_crs=src.crs
            src_count=src.count
            dtype=out_img.dtype

        with rasterio.open(dst_path+'/temp'+f'/{filename}_clipped.tif','w',driver='GTiff',
                           width=out_img.shape[2],
                           height=out_img.shape[1],
                           count=src_count,
                           crs=src_crs,
                           transform=out_transform,
                           dtype=dtype) as output:
            for i in range(src_count):
                output.write(out_img[i],i+1)
            output.close()

        with rasterio.open(dst_path+'/temp'+f'/{filename}_clipped.tif') as src_clipped:
            clipped_resampled=src_clipped.read(out_shape=(height,width),resampling=Resampling.nearest)

        with rasterio.open(dst_path+f'/{filename}_aligned.tif',
                           'w',
                           driver='GTiff',
                           width=width,
                           height=height,
                           count=src_count,
                           crs=crs,
                           transform=transform,
                           dtype=dtype) as dst:
            for i in range(src_count):
                dst.write(clipped_resampled[i],i+1)
            dst.close()

        shutil.rmtree(dst_path+'/temp')

        if remove==True:
            os.remove(path_to_target_image)
        else:
            pass


import os
import re
import sys
import logging
import subprocess
from typing import Optional, Dict
from simuleval import options
from simuleval.utils.arguments import cli_argument_list
from simuleval.utils.agent import get_agent_class

logger = logging.getLogger("simuleval.slurm")

def mkdir_output_dir(path: str) -> bool:
    try:
        os.makedirs(path, exist_ok=True)
        return True
    except BaseException as be:
        logger.error(f"Failed to write results to {path}.")
        logger.error(be)
        logger.error("Skip writing predictions.")
        return False

def submit_slurm_job(config_dict: Optional[Dict] = None) -> None:
    if config_dict is not None and "slurm" in config_dict:
        raise RuntimeError("--slurm is only available as a CLI argument")
    parser = options.general_parser()
    cli_arguments = cli_argument_list(config_dict)
    options.add_evaluator_args(parser)
    options.add_scorer_args(parser, cli_arguments)
    options.add_slurm_args(parser)
    options.add_dataloader_args(parser, cli_arguments)
    system_class = get_agent_class(config_dict)
    system_class.add_args(parser)
    args = parser.parse_args(cli_argument_list(config_dict))
    args.output = os.path.abspath(args.output)
    assert mkdir_output_dir(args.output)

    if args.agent is None:
        args.agent = sys.argv[0]

    os.system(f"cp {args.agent} {args.output}/agent.py")
    _args = [sys.argv[0]]
    for arg in sys.argv[1:]:
        if str(arg).isdigit() or str(arg).startswith("--"):
            _args.append(arg)
        else:
            _args.append(f'"{arg}"')
    command = " ".join(_args).strip()
    command = re.sub(r"(--slurm\S*(\s+[^-]\S+)*)", "", command).strip()
    if subprocess.check_output(["which", "simuleval"]).decode().strip() in command:
        command = re.sub(
            r"--agent\s+\S+", f"--agent {args.output}/agent.py", command
        ).strip()
    else:
        # Attention: not fully tested!
        command = re.sub(
            r"[^\"'\s]+\.py", f"{os.path.abspath(args.output)}/agent.py", command
        ).strip()

    if "--output" in command:
        command = re.sub(r"--output\s+\S+", f"--output {args.output}", command).strip()
    else:
        command += f" --output {args.output}"

    command = command.replace("--", "\\\n\t--")
    script = f"""#!/bin/bash
#SBATCH --time={args.slurm_time}
#SBATCH --partition={args.slurm_partition}
#SBATCH --nodes=1
#SBATCH --gpus-per-node=1
#SBATCH --ntasks-per-node=8
#SBATCH --output="{args.output}/slurm-%j.log"
#SBATCH --job-name="{args.slurm_job_name}"

cd {os.path.abspath(args.output)}

GPU_ID=$SLURM_LOCALID

# Change to local a gpu id for debugging, e.g.
# GPU_ID=0

CUDA_VISIBLE_DEVICES=$GPU_ID {command}
    """
    script_file = os.path.join(args.output, "script.sh")
    with open(script_file, "w") as f:
        f.writelines(script)

    process = subprocess.Popen(
        ["sbatch", script_file],
        stderr=subprocess.PIPE,
        stdout=subprocess.PIPE,
    )
    stdout, stderr = process.communicate()
    logger.info("Using slurm.")
    logger.info(f"sbatch stdout: {stdout.decode('utf-8').strip()}")
    stderr = stderr.decode("utf-8").strip()
    if len(stderr) > 0:
        logger.info(f"sbatch stderr: {stderr.decode('utf-8').strip()}")


import argparse
import numpy as np

from torch.utils.tensorboard import SummaryWriter

from basics.agents.base_agent import Agent

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical

class Policy(nn.Module):
    def __init__(self, observation_space, action_space_size):
        super(Policy, self).__init__()

        if len(observation_space) != 1:
            raise RuntimeError("The Policy specified by this agent only supports 1d obs space")

        self.affine1 = nn.Linear(observation_space[0], 128)
        self.dropout = nn.Dropout(p=0.6)
        self.affine2 = nn.Linear(128, action_space_size)

        self.saved_log_probs = []
        self.rewards = []

    def forward(self, x):
        x = self.affine1(x)
        x = self.dropout(x)
        x = F.relu(x)
        action_scores = self.affine2(x)
        return F.softmax(action_scores, dim=1)

class ReinforceAgent(Agent):

    def __init__(self, args, observation_space, action_space_size):
        self.policy = Policy(observation_space, action_space_size)
        self.optimizer = optim.Adam(self.policy.parameters(), lr=1e-2)
        self.eps = np.finfo(np.float32).eps.item()
        self.gamma = args.gamma
        self.episode_reward = 0
        self.running_reward = 10
        self.writer = SummaryWriter()

    def select_action(self, state):
        state = torch.from_numpy(state).float().unsqueeze(0)
        probs = self.policy(state)
        m = Categorical(probs)
        action = m.sample()
        self.policy.saved_log_probs.append(m.log_prob(action))
        return action.item()

    def finish_iteration(self, reward, iteration_count):
        self.policy.rewards.append(reward)
        self.episode_reward += reward

    def finish_episode(self, episode):
        self.running_reward = 0.05 * self.episode_reward + (1 - 0.05) * self.running_reward
        R = 0
        policy_loss = []
        returns = []
        for r in self.policy.rewards[::-1]:
            R = r + self.gamma * R
            returns.insert(0, R)
        returns = torch.tensor(returns)
        returns = (returns - returns.mean()) / (returns.std() + self.eps)
        for log_prob, R in zip(self.policy.saved_log_probs, returns):
            policy_loss.append(-log_prob * R)
        self.optimizer.zero_grad()
        policy_loss = torch.cat(policy_loss).sum()
        policy_loss.backward()
        self.optimizer.step()
        del self.policy.rewards[:]
        del self.policy.saved_log_probs[:]
        self.writer.add_scalar("Last reward", self.episode_reward, episode)
        self.writer.add_scalar("Average reward", self.running_reward, episode)
        self.episode_reward = 0

    def check_solved_env(self, reward_threshold):
        return self.running_reward > reward_threshold

    @staticmethod
    def add_args(parser):
        parser.add_argument('--gamma', type=float, default=0.99, metavar='G',
                    help='discount factor (default: 0.99)')
        return parser


import mock

from buildbot.data import buildslaves
from buildbot.test.fake import fakedb
from buildbot.test.fake import fakemaster
from buildbot.test.util import endpoint
from buildbot.test.util import interfaces
from twisted.internet import defer
from twisted.trial import unittest

testData = [
    fakedb.Builder(id=40, name=u'b1'),
    fakedb.Builder(id=41, name=u'b2'),
    fakedb.Master(id=13),
    fakedb.Master(id=14),
    fakedb.BuilderMaster(id=4013, builderid=40, masterid=13),
    fakedb.BuilderMaster(id=4014, builderid=40, masterid=14),
    fakedb.BuilderMaster(id=4113, builderid=41, masterid=13),

    fakedb.Buildslave(id=1, name=u'linux', info={}),
    fakedb.ConfiguredBuildslave(id=14013,
                                buildslaveid=1, buildermasterid=4013),
    fakedb.ConfiguredBuildslave(id=14014,
                                buildslaveid=1, buildermasterid=4014),
    fakedb.ConnectedBuildslave(id=113, masterid=13, buildslaveid=1),

    fakedb.Buildslave(id=2, name=u'windows', info={"a": "b"}),
    fakedb.ConfiguredBuildslave(id=24013,
                                buildslaveid=2, buildermasterid=4013),
    fakedb.ConfiguredBuildslave(id=24014,
                                buildslaveid=2, buildermasterid=4014),
    fakedb.ConfiguredBuildslave(id=24113,
                                buildslaveid=2, buildermasterid=4113),
    fakedb.ConnectedBuildslave(id=214, masterid=14, buildslaveid=2),
]

def _filt(bs, builderid, masterid):
    bs['connected_to'] = sorted(
        [d for d in bs['connected_to']
         if not masterid or masterid == d['masterid']])
    bs['configured_on'] = sorted(
        [d for d in bs['configured_on']
         if (not masterid or masterid == d['masterid'])
         and (not builderid or builderid == d['builderid'])])
    return bs

def bs1(builderid=None, masterid=None):
    return _filt({
        'buildslaveid': 1,
        'name': 'linux',
        'slaveinfo': {},
        'connected_to': [
            {'masterid': 13},
        ],
        'configured_on': sorted([
            {'builderid': 40, 'masterid': 13},
            {'builderid': 40, 'masterid': 14},
        ]),
    }, builderid, masterid)

def bs2(builderid=None, masterid=None):
    return _filt({
        'buildslaveid': 2,
        'name': 'windows',
        'slaveinfo': {'a': 'b'},
        'connected_to': [
            {'masterid': 14},
        ],
        'configured_on': sorted([
            {'builderid': 40, 'masterid': 13},
            {'builderid': 41, 'masterid': 13},
            {'builderid': 40, 'masterid': 14},
        ]),
    }, builderid, masterid)

class BuildslaveEndpoint(endpoint.EndpointMixin, unittest.TestCase):

    endpointClass = buildslaves.BuildslaveEndpoint
    resourceTypeClass = buildslaves.Buildslave

    def setUp(self):
        self.setUpEndpoint()
        return self.db.insertTestData(testData)

    def tearDown(self):
        self.tearDownEndpoint()

    def test_get_existing(self):
        d = self.callGet(('buildslaves', 2))

        @d.addCallback
        def check(buildslave):
            self.validateData(buildslave)
            buildslave['configured_on'].sort()
            self.assertEqual(buildslave, bs2())
        return d

    def test_get_existing_name(self):
        d = self.callGet(('buildslaves', 'linux'))

        @d.addCallback
        def check(buildslave):
            self.validateData(buildslave)
            buildslave['configured_on'].sort()
            self.assertEqual(buildslave, bs1())
        return d

    def test_get_existing_masterid(self):
        d = self.callGet(('masters', 14, 'buildslaves', 2))

        @d.addCallback
        def check(buildslave):
            self.validateData(buildslave)
            buildslave['configured_on'].sort()
            self.assertEqual(buildslave, bs2(masterid=14))
        return d

    def test_get_existing_builderid(self):
        d = self.callGet(('builders', 40, 'buildslaves', 2))

        @d.addCallback
        def check(buildslave):
            self.validateData(buildslave)
            buildslave['configured_on'].sort()
            self.assertEqual(buildslave, bs2(builderid=40))
        return d

    def test_get_existing_masterid_builderid(self):
        d = self.callGet(('masters', 13, 'builders', 40, 'buildslaves', 2))

        @d.addCallback
        def check(buildslave):
            self.validateData(buildslave)
            buildslave['configured_on'].sort()
            self.assertEqual(buildslave, bs2(masterid=13, builderid=40))
        return d

    def test_get_missing(self):
        d = self.callGet(('buildslaves', 99))

        @d.addCallback
        def check(buildslave):
            self.assertEqual(buildslave, None)
        return d

class BuildslavesEndpoint(endpoint.EndpointMixin, unittest.TestCase):

    endpointClass = buildslaves.BuildslavesEndpoint
    resourceTypeClass = buildslaves.Buildslave

    def setUp(self):
        self.setUpEndpoint()
        return self.db.insertTestData(testData)

    def tearDown(self):
        self.tearDownEndpoint()

    def test_get(self):
        d = self.callGet(('buildslaves',))

        @d.addCallback
        def check(buildslaves):
            [self.validateData(b) for b in buildslaves]
            [b['configured_on'].sort() for b in buildslaves]
            self.assertEqual(sorted(buildslaves), sorted([bs1(), bs2()]))
        return d

    def test_get_masterid(self):
        d = self.callGet(('masters', '13', 'buildslaves',))

        @d.addCallback
        def check(buildslaves):
            [self.validateData(b) for b in buildslaves]
            [b['configured_on'].sort() for b in buildslaves]
            self.assertEqual(sorted(buildslaves),
                             sorted([bs1(masterid=13), bs2(masterid=13)]))
        return d

    def test_get_builderid(self):
        d = self.callGet(('builders', '41', 'buildslaves',))

        @d.addCallback
        def check(buildslaves):
            [self.validateData(b) for b in buildslaves]
            [b['configured_on'].sort() for b in buildslaves]
            self.assertEqual(sorted(buildslaves),
                             sorted([bs2(builderid=41)]))
        return d

    def test_get_masterid_builderid(self):
        d = self.callGet(('masters', '13', 'builders', '41', 'buildslaves',))

        @d.addCallback
        def check(buildslaves):
            [self.validateData(b) for b in buildslaves]
            [b['configured_on'].sort() for b in buildslaves]
            self.assertEqual(sorted(buildslaves),
                             sorted([bs2(masterid=13, builderid=41)]))
        return d

class Buildslave(interfaces.InterfaceTests, unittest.TestCase):

    def setUp(self):
        self.master = fakemaster.make_master(testcase=self,
                                             wantMq=True, wantDb=True, wantData=True)
        self.rtype = buildslaves.Buildslave(self.master)
        return self.master.db.insertTestData([
            fakedb.Master(id=13),
            fakedb.Master(id=14),
        ])

    def test_signature_findBuildslaveId(self):
        @self.assertArgSpecMatches(
            self.master.data.updates.findBuildslaveId,  # fake
            self.rtype.findBuildslaveId)  # real
        def findBuildslaveId(self, name):
            pass

    def test_signature_buildslaveConfigured(self):
        @self.assertArgSpecMatches(
            self.master.data.updates.buildslaveConfigured,  # fake
            self.rtype.buildslaveConfigured)  # real
        def buildslaveConfigured(self, buildslaveid, masterid, builderids):
            pass

    def test_findBuildslaveId(self):
        # this just passes through to the db method, so test that
        rv = defer.succeed(None)
        self.master.db.buildslaves.findBuildslaveId = \
            mock.Mock(return_value=rv)
        self.assertIdentical(self.rtype.findBuildslaveId(u'foo'), rv)

    def test_findBuildslaveId_not_id(self):
        self.assertRaises(ValueError, self.rtype.findBuildslaveId, 'foo')
        self.assertRaises(ValueError, self.rtype.findBuildslaveId, u'123/foo')


import wx

def border(dc, rect, top=True, bottom=True, left=True, right=True,
           rounded=False, size=None):
    """
    Draws a border for the specified rectangle. The current brush and pen are
    used.

    Defaults are to draw the whole rectangle, but if you want to leave sides
    out, specify them as False, like

    >>> draw.border(dc, rect, left=False, right=False)

    If the rounded keyword argument is set to True, then a full rectangle will
    be drawn with rounded edges.
    """
    dc.SetPen(wx.BLACK_PEN)

    if top and bottom and left and right:
        dc.SetBrush(wx.TRANSPARENT_BRUSH)
        if rounded:
            dc.DrawRoundedRectangle(rect.x,rect.y,rect.width,rect.height, rounded)
        else:
            dc.DrawRectangle(rect.x,rect.y,rect.width, rect.height)
    else:
        pen_width = dc.GetPen().GetWidth()
        bottompos = rect.y + rect.height - pen_width
        rightpos  = rect.x + rect.width  - pen_width

        if top:    dc.DrawLine(rect.x, rect.y, rightpos, rect.y)
        if bottom: dc.DrawLine(rect.x, bottompos, rightpos, bottompos)
        if left:   dc.DrawLine(rect.x, rect.y, rect.x, bottompos)
        if right:  dc.DrawLine(rightpos, rect.y, rightpos, bottompos)


import asyncio
import copy
import os
import re
import select
import signal
import sys
import termios
import tty
import unicodedata
from argparse import ArgumentParser

"""Editing widget using the interior of a window object.
 Supports the following Emacs-like key bindings:
Ctrl-A      Go to left edge of window.
Ctrl-B      Cursor left, wrapping to previous line if appropriate.
Ctrl-D      Delete character under cursor.
Ctrl-E      Go to right edge (stripspaces off) or end of line (stripspaces on).
Ctrl-F      Cursor right, wrapping to next line when appropriate.
Ctrl-G      Terminate, returning the window contents.
Ctrl-H      Delete character backward.
Ctrl-J      Terminate if the window is 1 line, otherwise insert newline.
Ctrl-K      If line is blank, delete it, otherwise clear to end of line.
Ctrl-L      Refresh screen.
Ctrl-N      Cursor down; move down one line.
Ctrl-O      Insert a blank line at cursor location.
Ctrl-P      Cursor up; move up one line.
Move operations do nothing if the cursor is at an edge where the movement
is not possible.  The following synonyms are supported where possible:
KEY_LEFT = Ctrl-B, KEY_RIGHT = Ctrl-F, KEY_UP = Ctrl-P, KEY_DOWN = Ctrl-N
KEY_BACKSPACE = Ctrl-h
"""

NUL = 0x00  # ^@
SOH = 0x01  # ^A
STX = 0x02  # ^B
ETX = 0x03  # ^C
EOT = 0x04  # ^D
ENQ = 0x05  # ^E
ACK = 0x06  # ^F
BEL = 0x07  # ^G
BS = 0x08  # ^H
TAB = 0x09  # ^I
HT = 0x09  # ^I
LF = 0x0a  # ^J
NL = 0x0a  # ^J
VT = 0x0b  # ^K
FF = 0x0c  # ^L
CR = 0x0d  # ^M
SO = 0x0e  # ^N
SI = 0x0f  # ^O
DLE = 0x10  # ^P
DC1 = 0x11  # ^Q
DC2 = 0x12  # ^R
DC3 = 0x13  # ^S
DC4 = 0x14  # ^T
NAK = 0x15  # ^U
SYN = 0x16  # ^V
ETB = 0x17  # ^W
CAN = 0x18  # ^X
EM = 0x19  # ^Y
SUB = 0x1a  # ^Z
ESC = 0x1b  # ^[
FS = 0x1c  # ^\
GS = 0x1d  # ^]
RS = 0x1e  # ^^
US = 0x1f  # ^_
SP = 0x20  # space
DEL = 0x7f  # delete

DEFAULT_WIDTH = 80

class Terminal:
    """Input line editor"""

    def __init__(self, history_size=10, max_line_length=100, term_type="vt100", width=80):
        # self._orig_chan = orig_chan
        # self._orig_session = orig_session
        self._history_size = history_size if history_size > 0 else 0
        self._max_line_length = max_line_length
        self._wrap = term_type in (
            "ansi", "cygwin", "linux", "putty", "screen", "teraterm", "cit80", "vt100", "vt102", "vt220", "vt320",
            "xterm", "xterm-color", "xterm-16color", "xterm-256color", "rxvt", "rxvt-color"
        )
        self._width = width or DEFAULT_WIDTH
        self._line_mode = True
        self._echo = True
        self._start_column = 0
        self._end_column = 0
        self._cursor = 0
        self._left_pos = 0
        self._right_pos = 0
        self._pos = 0
        self._line = ""
        self._bell_rung = False
        self._early_wrap = set()
        self._outbuf = list()
        self._keymap = dict()
        self._key_state = self._keymap
        self._erased = ""
        self._history = list()
        self._history_index = 0

        for func, keys in self._keylist:
            for key in keys:
                self._add_key(key, func)

        self._build_printable()

    def _add_key(self, key, func):
        """Add a key to the keymap"""

        keymap = self._keymap

        for ch in key[:-1]:
            if ch not in keymap:
                keymap[ch] = {}

            keymap = keymap[ch]

        keymap[key[-1]] = func

    def _del_key(self, key):
        """Delete a key from the keymap"""

        keymap = self._keymap

        for ch in key[:-1]:
            if ch not in keymap:
                return

            keymap = keymap[ch]

        keymap.pop(key[-1], None)

    def _build_printable(self):
        """Build a regex of printable ASCII non-registered keys"""

        def _escape(c):
            """Backslash escape special characters in regex character range"""

            ch = chr(c)
            return ("\\" if (ch in "-&|[]\\^~") else "") + ch

        def _is_printable(ch):
            """Return if character is printable and has no handler"""

            return ch.isprintable() and ch not in keys

        pat = []
        keys = self._keymap.keys()
        start = ord(" ")
        limit = 0x10000

        while start < limit:
            while start < limit and not _is_printable(chr(start)):
                start += 1

            end = start

            while _is_printable(chr(end)):
                end += 1

            pat.append(_escape(start))

            if start != end - 1:
                pat.append("-" + _escape(end - 1))

            start = end + 1

        self._printable = re.compile("[" + "".join(pat) + "]*")

    def _char_width(self, pos):
        """Return width of character at specified position"""

        return 1 + (unicodedata.east_asian_width(self._line[pos]) in "WF") + ((pos + 1) in self._early_wrap)

    def _determine_column(self, data, column, pos=None):
        """Determine new output column after output occurs"""

        offset = pos
        last_wrap_pos = pos
        wrapped_data = []

        for ch in data:
            if ch == "\b":
                column -= 1
            else:
                if (unicodedata.east_asian_width(ch) in "WF") and (column % self._width) == self._width - 1:
                    column += 1

                    if pos is not None:
                        wrapped_data.append(data[last_wrap_pos - offset:pos - offset])
                        last_wrap_pos = pos

                        self._early_wrap.add(pos)
                else:
                    if pos is not None:
                        self._early_wrap.discard(pos)

                column += 1 + (unicodedata.east_asian_width(ch) in "WF")

            if pos is not None:
                pos += 1

        if pos is not None:
            wrapped_data.append(data[last_wrap_pos - offset:])
            return " ".join(wrapped_data), column
        else:
            return data, column

    def _output(self, data, pos=None):
        """Generate output and calculate new output column"""

        idx = data.rfind("\n")

        if idx >= 0:
            self._outbuf.append(data[:idx + 1])
            tail = data[idx + 1:]
            self._cursor = 0
        else:
            tail = data

        data, self._cursor = self._determine_column(tail, self._cursor, pos)

        self._outbuf.append(data)

        if self._cursor and self._cursor % self._width == 0:
            self._outbuf.append(" \b")

    def _ring_bell(self):
        """Ring the terminal bell"""

        if not self._bell_rung:
            self._outbuf.append("\a")
            self._bell_rung = True

    def _update_input_window(self, new_pos):
        """Update visible input window when not wrapping onto multiple lines"""

        line_len = len(self._line)

        if new_pos < self._left_pos:
            self._left_pos = new_pos
        else:
            if new_pos < line_len:
                new_pos += 1

            pos = self._pos
            column = self._cursor

            while pos < new_pos:
                column += self._char_width(pos)
                pos += 1

            if column >= self._width:
                while column >= self._width:
                    column -= self._char_width(self._left_pos)
                    self._left_pos += 1
            else:
                while self._left_pos > 0:
                    column += self._char_width(self._left_pos)

                    if column < self._width:
                        self._left_pos -= 1
                    else:
                        break

        column = self._start_column
        self._right_pos = self._left_pos

        while self._right_pos < line_len:
            ch_width = self._char_width(self._right_pos)

            if column + ch_width < self._width:
                self._right_pos += 1
                column += ch_width
            else:
                break

        return column

    def _move_cursor(self, column):
        """Move the cursor to selected position in input line"""

        start_row = self._cursor // self._width
        start_col = self._cursor % self._width

        end_row = column // self._width
        end_col = column % self._width

        if end_row < start_row:
            self._outbuf.append("\x1b[" + str(start_row - end_row) + "A")
        elif end_row > start_row:
            self._outbuf.append("\x1b[" + str(end_row - start_row) + "B")

        if end_col > start_col:
            self._outbuf.append("\x1b[" + str(end_col - start_col) + "C")
        elif end_col < start_col:
            self._outbuf.append("\x1b[" + str(start_col - end_col) + "D")

        self._cursor = column

    def _move_back(self, column):
        """Move the cursor backward to selected position in input line"""

        if self._wrap:
            self._move_cursor(column)
        else:
            self._outbuf.append("\b" * (self._cursor - column))
            self._cursor = column

    def _clear_to_end(self):
        """Clear any remaining characters from previous input line"""

        column = self._cursor
        remaining = self._end_column - column

        if remaining > 0:
            self._outbuf.append(" " * remaining)
            self._cursor = self._end_column

            if self._cursor % self._width == 0:
                self._outbuf.append(" \b")

        self._move_back(column)
        self._end_column = column

    def _erase_input(self):
        """Erase current input line"""

        self._move_cursor(self._start_column)
        self._clear_to_end()
        self._early_wrap.clear()

    def _draw_input(self):
        """Draw current input line"""

        if self._line and self._echo:
            if self._wrap:
                self._output(self._line[:self._pos], 0)
                column = self._cursor
                self._output(self._line[self._pos:], self._pos)
            else:
                self._update_input_window(self._pos)
                self._output(self._line[self._left_pos:self._pos])
                column = self._cursor
                self._output(self._line[self._pos:self._right_pos])

            self._end_column = self._cursor
            self._move_back(column)

    def _reposition(self, new_pos, new_column):
        """Reposition the cursor to selected position in input"""

        if self._echo:
            if self._wrap:
                self._move_cursor(new_column)
            else:
                self._update_input(self._pos, self._cursor, new_pos)

        self._pos = new_pos

    def _update_input(self, pos, column, new_pos):
        """Update selected portion of current input line"""

        if self._echo:
            if self._wrap:
                if pos in self._early_wrap:
                    column -= 1

                self._move_cursor(column)
                prev_wrap = new_pos in self._early_wrap
                self._output(self._line[pos:new_pos], pos)
                column = self._cursor
                self._output(self._line[new_pos:], new_pos)
                column += (new_pos in self._early_wrap) - prev_wrap
            else:
                self._update_input_window(new_pos)
                self._move_back(self._start_column)
                self._output(self._line[self._left_pos:new_pos])
                column = self._cursor
                self._output(self._line[new_pos:self._right_pos])

            self._clear_to_end()
            self._move_back(column)

        self._pos = new_pos

    def _insert_printable(self, data):
        """Insert data into the input line"""

        line_len = len(self._line)
        data_len = len(data)

        if self._max_line_length:
            if line_len + data_len > self._max_line_length:
                self._ring_bell()
                data_len = self._max_line_length - line_len
                data = data[:data_len]

        if data:
            pos = self._pos
            new_pos = pos + data_len
            self._line = self._line[:pos] + data + self._line[pos:]

            self._update_input(pos, self._cursor, new_pos)

    def _end_line(self):
        """End the current input line and send it to the session"""

        if (self._echo and not self._wrap and
                (self._left_pos > 0 or self._right_pos < len(self._line))):
            self._output("\b" * (self._cursor - self._start_column) + self._line)
        else:
            self._move_to_end()

        self._output("\r\n")

        self._start_column = 0
        self._end_column = 0
        self._cursor = 0
        self._left_pos = 0
        self._right_pos = 0
        self._pos = 0

        if self._echo and self._history_size and self._line:
            self._history.append(self._line)
            self._history = self._history[-self._history_size:]

        self._history_index = len(self._history)

        data = self._line + "\n"
        self._line = ""

        # self._session.data_received(data, None)

    def _eof_or_delete(self):
        """Erase character to the right, or send EOF if input line is empty"""

        if not self._line:
            self._session.soft_eof_received()
        else:
            self._erase_right()

    def _erase_left(self):
        """Erase character to the left"""

        if self._pos > 0:
            pos = self._pos - 1
            column = self._cursor - self._char_width(pos)
            self._line = self._line[:pos] + self._line[pos + 1:]
            self._update_input(pos, column, pos)
        else:
            self._ring_bell()

    def _erase_right(self):
        """Erase character to the right"""

        if self._pos < len(self._line):
            pos = self._pos
            self._line = self._line[:pos] + self._line[pos + 1:]
            self._update_input(pos, self._cursor, pos)
        else:
            self._ring_bell()

    def _erase_line(self):
        """Erase entire input line"""

        self._erased = self._line
        self._line = ""
        self._update_input(0, self._start_column, 0)

    def _erase_to_end(self):
        """Erase to end of input line"""

        pos = self._pos
        self._erased = self._line[pos:]
        self._line = self._line[:pos]
        self._update_input(pos, self._cursor, pos)

    def _handle_key(self, key, handler):
        """Call an external key handler"""

        result = handler(self._line, self._pos)

        if result is True:
            if key.isprintable():
                self._insert_printable(key)
            else:
                self._ring_bell()
        elif result is False:
            self._ring_bell()
        else:
            line, new_pos = result

            if new_pos < 0:
                self._session.signal_received(line)
            else:
                self._line = line
                self._update_input(0, self._start_column, new_pos)

    def _history_prev(self):
        """Replace input with previous line in history"""

        if self._history_index > 0:
            self._history_index -= 1
            self._line = self._history[self._history_index]
            self._update_input(0, self._start_column, len(self._line))
        else:
            self._ring_bell()

    def _history_next(self):
        """Replace input with next line in history"""

        if self._history_index < len(self._history):
            self._history_index += 1

            if self._history_index < len(self._history):
                self._line = self._history[self._history_index]
            else:
                self._line = ""

            self._update_input(0, self._start_column, len(self._line))
        else:
            self._ring_bell()

    def _move_left(self):
        """Move left in input line"""

        if self._pos > 0:
            pos = self._pos - 1
            column = self._cursor - self._char_width(pos)
            self._reposition(pos, column)
        else:
            self._ring_bell()

    def _move_right(self):
        """Move right in input line"""

        if self._pos < len(self._line):
            pos = self._pos
            column = self._cursor + self._char_width(pos)
            self._reposition(pos + 1, column)
        else:
            self._ring_bell()

    def _move_to_start(self):
        """Move to start of input line"""

        self._reposition(0, self._start_column)

    def _move_to_end(self):
        """Move to end of input line"""

        self._reposition(len(self._line), self._end_column)

    def _redraw(self):
        """Redraw input line"""

        self._erase_input()
        self._draw_input()

    def _insert_erased(self):
        """Insert previously erased input"""

        self._insert_printable(self._erased)

    def _send_break(self):
        """Send break to session"""

        self._session.break_received(0)

    # pylint: disable=bad-whitespace

    _keylist = ((_end_line, ("\n", "\r", "\x1bOM")),
                (_eof_or_delete, ("\x04",)),
                (_erase_left, ("\x08", "\x7f")),
                (_erase_right, ("\x1b[3~",)),
                (_erase_line, ("\x15",)),
                (_erase_to_end, ("\x0b",)),
                (_history_prev, ("\x10", "\x1b[A", "\x1bOA")),
                (_history_next, ("\x0e", "\x1b[B", "\x1bOB")),
                (_move_left, ("\x02", "\x1b[D", "\x1bOD")),
                (_move_right, ("\x06", "\x1b[C", "\x1bOC")),
                (_move_to_start, ("\x01", "\x1b[H", "\x1b[1~")),
                (_move_to_end, ("\x05", "\x1b[F", "\x1b[4~")),
                (_redraw, ("\x12",)),
                (_insert_erased, ("\x19",)),
                (_send_break, ("\x03", "\x1b[33~")))

    # pylint: enable=bad-whitespace

    def register_key(self, key, handler):
        """Register a handler to be called when a key is pressed"""

        self._add_key(key, partial(SSHLineEditor._handle_key,
                                   key=key, handler=handler))
        self._build_printable()

    def unregister_key(self, key):
        """Remove the handler associated with a key"""

        self._del_key(key)
        self._build_printable()

    def set_input(self, line, pos):
        """Set input line and cursor position"""

        self._line = line
        self._update_input(0, self._start_column, pos)

    def set_line_mode(self, line_mode):
        """Enable/disable input line editing"""

        if self._line and not line_mode:
            data = self._line
            self._erase_input()
            self._line = ""

            self._session.data_received(data, None)

        self._line_mode = line_mode

    def set_echo(self, echo):
        """Enable/disable echoing of input in line mode"""

        if self._echo and not echo:
            self._erase_input()
            self._echo = False
        elif echo and not self._echo:
            self._echo = True
            self._draw_input()

    def set_width(self, width):
        """Set terminal line width"""

        self._width = width or DEFAULT_WIDTH

        if self._wrap:
            _, self._cursor = self._determine_column(self._line, self._start_column, 0)

        self._redraw()

    def process_input(self, data, datatype):
        """Process input from channel"""

        if self._line_mode:
            data_len = len(data)
            idx = 0

            while idx < data_len:
                ch = data[idx]
                idx += 1

                if ch in self._key_state:
                    self._key_state = self._key_state[ch]
                    if callable(self._key_state):
                        try:
                            self._key_state(self)
                        finally:
                            self._key_state = self._keymap
                elif self._key_state == self._keymap and ch.isprintable():
                    match = self._printable.match(data, idx - 1)[0]

                    if match:
                        self._insert_printable(match)
                        idx += len(match) - 1
                    else:
                        self._insert_printable(ch)
                else:
                    self._key_state = self._keymap
                    self._ring_bell()

            self._bell_rung = False
            self.write("".join(self._outbuf))
            self._outbuf.clear()
        else:
            self._session.data_received(data, datatype)

    def process_output(self, data):
        """Process output to channel"""

        data = data.replace("\n", "\r\n")

        self._erase_input()
        self._output(data)

        if not self._wrap:
            self._cursor %= self._width

        self._start_column = self._cursor
        self._end_column = self._cursor
        self._draw_input()

        self._chan.write("".join(self._outbuf))
        self._outbuf.clear()

    def __getattr__(self, attr):
        """Delegate most channel functions to original channel"""

        return getattr(self._orig_chan, attr)

    def create_editor(self):
        """Create input line editor if encoding and terminal type are set"""
        return self

    def clear_input(self):
        """Clear input line
           This method clears the current input line.
        """

        self.set_input("", 0)

    def write(self, data, datatype=None):
        """Process data written to the channel"""

        # if self._editor and datatype is None:
        #    self._editor.process_output(data)
        # else:
        #    self._orig_chan.write(data, datatype)
        print(data)

    def session_started(self):
        """Start a session for this newly opened server channel"""

        # self._editor = self._chan.create_editor()
        # self._orig_session.session_started()

        self.create_editor()

    def terminal_size_changed(self, width, height, pixwidth, pixheight):
        """The terminal size has changed"""

        if self._editor:
            self._editor.set_width(width)

        self._orig_session.terminal_size_changed(width, height, pixwidth, pixheight)

    def data_received(self, data, datatype):
        """Process data received from the channel"""
        self.process_input(data, datatype)

    def eof_received(self):
        """Process EOF received from the channel"""

        if self._editor:
            self._editor.set_line_mode(False)

        return self._orig_session.eof_received()

from angelos.common.utils import Util

"""
Control sequence = CSI = ESC [
F = A control sequence, where F is from 0o100 to 0o176 inclusive.

Escape sequence = ESC
F = An escape sequence, where F is from 0o60 to 0o176 inclusive

An intermediate character in an escape sequence or a control sequence, where I is from 0o40 to 0o57 inclusive.

All of the following escape and control sequences are transmitted from the host computer to the VT100 unless otherwise noted

CPR – Cursor Position Report – VT100 to Host
ESC [ Pn ; Pn R

CUB – Cursor Backward – Host to VT100 and VT100 to Host
ESC [ Pn D

CUD – Cursor Down – Host to VT100 and VT100 to Host
ESC [ Pn B

CUF – Cursor Forward – Host to VT100 and VT100 to Host
ESC [ Pn C

CUP – Cursor Position
ESC [ Pn ; Pn H

CUU – Cursor Up – Host to VT100 and VT100 to Host
ESC [ Pn A

DA – Device Attributes
ESC [ Pn c

DECALN – Screen Alignment Display (DEC Private)
ESC # 8

DECDHL – Double Height Line (DEC Private)
Top Half: ESC # 3
Bottom Half: ESC # 4

DECDWL – Double-Width Line (DEC Private)
ESC # 6

DECID – Identify Terminal (DEC Private)
ESC Z

DECKPAM – Keypad Application Mode (DEC Private)
ESC =

DECKPNM – Keypad Numeric Mode (DEC Private)
ESC >

DECLL – Load LEDS (DEC Private)
ESC [ Ps q	default value: 0

DECRC – Restore Cursor (DEC Private)
ESC 8

DECREPTPARM – Report Terminal Parameters
ESC [ <sol>; <par>; <nbits>; <xspeed>; <rspeed>; <clkmul>; <flags> x

DECREQTPARM – Request Terminal Parameters
ESC [ <sol> x

DECSC – Save Cursor (DEC Private)
ESC 7

DECSTBM – Set Top and Bottom Margins (DEC Private)
ESC [ Pn; Pn r

DECSWL – Single-width Line (DEC Private)
ESC # 5

DECTST – Invoke Confidence Test
ESC [ 2 ; Ps y

DSR – Device Status Report
ESC [ Ps n

ED – Erase In Display
ESC [ Ps J

EL – Erase In Line
ESC [ Ps K

HTS – Horizontal Tabulation Set
ESC H

HVP – Horizontal and Vertical Position
ESC [ Pn ; Pn f

IND – Index
ESC D

NEL – Next Line
ESC E

RI – Reverse Index
ESC M

RIS – Reset To Initial State
ESC c

RM – Reset Mode
ESC [ Ps ; Ps ; . . . ; Ps l

SCS – Select Character Set
G0 Sets Sequence	G1 Sets Sequence	Meaning
ESC ( A	ESC ) A	United Kingdom Set
ESC ( B	ESC ) B	ASCII Set
ESC ( 0	ESC ) 0	Special Graphics
ESC ( 1	ESC ) 1	Alternate Character ROM Standard Character Set
ESC ( 2	ESC ) 2	Alternate Character ROM Special Graphics

SGR – Select Graphic Rendition
ESC [ Ps ; . . . ; Ps m

SM – Set Mode
ESC [ Ps ; . . . ; Ps h

TBC – Tabulation Clear
ESC [ Ps g

Cursor Up
ESC A

Cursor Down
ESC B

Cursor Right
ESC C

Cursor Left
ESC D

Enter Graphics Mode
ESC F

Exit Graphics Mode
ESC G

Cursor to Home
ESC H

Reverse Line Feed
ESC I

Erase to End of Screen
ESC J

Erase to End of Line
ESC K

Direct Cursor Address
ESC Y line column

Identify
ESC Z
ESC / Z

Enter Alternate Keypad Mode
ESC =

Exit Alternate Keypad Mode
ESC >

Enter ANSI Mode
ESC <

Cursor up	ESC [ Pn A
Cursor down	ESC [ Pn B
Cursor forward (right)	ESC [ Pn C
Cursor backward (left)	ESC [ Pn D
Direct cursor addressing	ESC [ Pl ; Pc H† or
ESC [ Pl ; Pc f†
Index	ESC D
New line	ESC E
Reverse index	ESC M
Save cursor and attributes	ESC 7
Restore cursor and attributes	ESC 8

Change this line to double-height top half	ESC # 3
Change this line to double-height bottom half	ESC # 4
Change this line to single-width single-height	ESC # 5
Change this line to double-width single-height	ESC # 6

Character Attributes
ESC [ Ps;Ps;Ps;...;Ps m

From cursor to end of line	ESC [ K or ESC [ 0 K
From beginning of line to cursor	ESC [ 1 K
Entire line containing cursor	ESC [ 2 K
From cursor to end of screen	ESC [ J or ESC [ 0 J
From beginning of screen to cursor	ESC [ 1 J
Entire screen	ESC [ 2 J

Programmable LEDs
ESC [ Ps;Ps;...Ps q

United Kingdom (UK)	ESC ( A	ESC ) A
United States (USASCII)	ESC ( B	ESC ) B
Special graphics characters and line drawing set	ESC ( 0	ESC ) 0
Alternate character ROM	ESC ( 1	ESC ) 1
Alternate character ROM special graphics characters	ESC ( 2	ESC ) 2

Scrolling Region
ESC [ Pt ; Pb r

Set tab at current column	ESC H
Clear tab at current column	ESC [ g or ESC [ 0 g
Clear all tabs	ESC [ 3 g

Line feed/new line	New line	ESC [20h	Line feed	ESC [20l*
Cursor key mode	Application	ESC [?1h	Cursor	ESC [?1l*
ANSI/VT52 mode	ANSI	N/A	VT52	ESC [?2l*
Column mode	132 Col	ESC [?3h	80 Col	ESC [?3l*
Scrolling mode	Smooth	ESC [?4h	Jump	ESC [?4l*
Screen mode	Reverse	ESC [?5h	Normal	ESC [?5l*
Origin mode	Relative	ESC [?6h	Absolute	ESC [?6l*
Wraparound	On	ESC [?7h	Off	ESC [?7l*
Auto repeat	On	ESC [?8h	Off	ESC [?8l*
Interlace	On	ESC [?9h	Off	ESC [?9l*
Keypad mode	Application	ESC =	Numeric	ESC >

Invoked by	ESC [ 6 n
Response is	ESC [ Pl ; Pc R †

Invoked by	ESC [ 5 n
Response is	ESC [ 0 n (terminal ok)
ESC [ 3 n (terminal not ok)

Invoked by	ESC [ c or ESC [ 0 c
Response is	ESC [ ? 1 ; Ps c

ESC c

Fill Screen with "Es"	ESC # 8
Invoke Test(s)	ESC [ 2 ; Ps y

Cursor Up	ESC A
Cursor Down	ESC B
Cursor Right	ESC C
Cursor Left	ESC D
Select Special Graphics character set	ESC F
Select ASCII character set	ESC G
Cursor to home	ESC H
Reverse line feed	ESC I
Erase to end of screen	ESC J
Erase to end of line	ESC K
Direct cursor address	ESC Y l c	(see note 1)
Identify	ESC Z	(see note 2)
Enter alternate keypad mode	ESC =
Exit alternate keypad mode	ESC >
Enter ANSI mode	ESC <
"""

REGEX = r"""(\x1B\[[\x40-\x7E]*[\x20-\x2F]|\x1B[\x30-\x7E]*[\x20-\x2F]|[\x00-\x1F])"""  # (CSI I F|ESC I F)

class TTYCodes:
    class SeqCodes:
        """Control and escape sequence codes for the VT100."""

        CTRL_SP = " "  # Space
        CTRL_NUL = "\x00"  # Null
        CTRL_BEL = "\x07"  # Beep
        CTRL_BS = "\x08"  # Backspace
        CTRL_HT = "\x09"  # Tab
        CTRL_LF = "\n"  # New line
        CTRL_VT = "\x0b"  # Vertical tab
        CTRL_FF = "\x0c"  # Form feed
        CTRL_CR = "\r"  # Carriage return
        CTRL_SO = "\x0e"  # Shift out
        CTRL_SI = "\x0f"  # Shift in
        CTRL_CAN = "\x18"  # Cancel
        CTRL_SUB = "\x1a"  # Substitute
        CTRL_ESC = "\x1b"  # Escape
        CTRL_DEL = "\x7f"  # Delete
        CTRL_CSI_C0 = "\x1b["  # Control sequence introducer
        CTRL_CSI_C1 = "\x9b"  # Control sequence introducer
        CTRL_CSI = "\x1b["  # Control sequence introducer
        CTRL_ST_C0 = "\x1b\\"  # String terminator
        CTRL_ST_C1 = "\x9c"  # String terminator
        CTRL_ST = "\x1b\\"  # String terminator
        CTRL_OSC_C0 = "\x1b]"  # OS command
        CTRL_OSC_C1 = "\x9d"  # OS command
        CTRL_OSC = "\x1b]"  # OS command

        ESC_RIS = "c"  # Reset
        ESC_IND = "D"  # Index
        ESC_NEL = "E"  # Next line
        ESC_HTS = "H"  # Tabulation set
        ESC_RI = "M"  # Reverse index
        ESC_DECSC = "7"  # Save cursor
        ESC_DECRC = "8"  # Restore cursor
        ESC_DECALN = "8"  # Alignment display
        ESC_ICH = "@"  # Insert character
        ESC_CUU = "A"  # Cursor up
        ESC_CUD = "B"  # Cursor down
        ESC_CUF = "C"  # Cursor forward
        ESC_CUB = "D"  # Cursor back
        ESC_CNL = "E"  # Cursor next line
        ESC_CPL = "F"  # Cursor previous line
        ESC_CHA = "G"  # Cursor horizontal align
        ESC_CUP = "H"  # Cursor position
        ESC_ED = "J"  # Erase data
        ESC_EL = "K"  # Erase in line
        ESC_IL = "L"  # Insert line
        ESC_DL = "M"  # Delete line
        ESC_DCH = "P"  # Delete character
        ESC_ECH = "X"  # Erase character
        ESC_HPR = "a"  # Horizontal position relative
        ESC_DA = "c"  # Device attributes
        ESC_VPA = "d"  # Vertical position adjust
        ESC_VPR = "e"  # Vertical position relative
        ESC_HVP = "f"  # Horizontal / vertical position
        ESC_TBC = "g"  # Tabulation clear
        ESC_SM = "h"  # Set mode
        ESC_RM = "l"  # Reset mode
        ESC_SGR = "m"  # Select graphics rendition
        ESC_DSR = "n"  # Device status report
        ESC_DECSTBM = "r"  # Select top and bottom margins
        ESC_HPA = "'"  # Horizontal adjust

    class GraphicsCodes:
        """Graphic codes for VT100."""

        BOLD = 1
        ITALICS = 3
        UNDERSCORE = 4
        BLINK = 5
        REVERSE = 7
        STRIKE = 9
        BOLD_RM = 22
        ITALICS_RM = 23
        UNDERSCORE_RM = 24
        BLINK_RM = 25
        REVERSE_RM = 27
        STRIKE_RM = 29

        FG_BLACK = 30
        FG_RED = 31
        FG_GREEN = 32
        FG_BROWN = 33
        FG_BLUE = 34
        FG_MAGENTA = 35
        FG_CYAN = 36
        FG_WHITE = 37
        FG_DEFAULT = 39  # White

        FG_BRIGHTBLACK = 90
        FG_BRIGHTRED = 91
        FG_BRIGHTGREEN = 92
        FG_BRIGHTBROWN = 93
        FG_BRIGHTBLUE = 94
        FG_BRIGHTMAGENTA = 95
        FG_BRIGHTCYAN = 96
        FG_BRIGHTWHITE = 97

        BG_BLACK = 40
        BG_RED = 41
        BG_GREEN = 42
        BG_BROWN = 43
        BG_BLUE = 44
        BG_MAGENTA = 45
        BG_CYAN = 46
        BG_WHITE = 47
        BG_DEFAULT = 49  # Black

        BG_BRIGHTBLACK = 100
        BG_BRIGHTRED = 101
        BG_BRIGHTGREEN = 102
        BG_BRIGHTBROWN = 103
        BG_BRIGHTBLUE = 104
        BG_BRIGHTMAGENTA = 105
        BG_BRIGHTCYAN = 106
        BG_BRIGHTWHITE = 107

    class ModeCodes:
        """Line feed modes for VT100"""

        LNM = 20  # Line feed / New line mode
        IRM = 4  # Insert / Replace mode
        DECTCEM = 25 << 5  # Text cursor enable mode
        DECSCNM = 5 << 5  # Screen mode
        DECOM = 6 << 5  # Origin mode
        DECAWM = 7 << 5  # Auto wrap mode
        DECCOLM = 3 << 5  # Column mode

class Application:
    """Test application for terminal emulator."""

    def __init__(self):
        self._quiter = None
        self._task = None

        self._tty = sys.stdin.fileno()
        self._echo = termios.tcgetattr(self._tty)
        self._no_echo = copy.copy(self._echo)
        self._no_echo[3] = self._no_echo[3] & ~termios.ECHO

        self._terminal = None

    def _quit(self):
        self._quiter.set()

    def _sigint_handler(self):
        self.on_quit()
        self._quit()

    def _sigwinch_handler(self):
        size = os.get_terminal_size()
        self.on_resize(size.columns, size.lines)

    def _input_handler(self):
        self.on_input(sys.stdin.buffer.read1())

    def _build_parser(self) -> ArgumentParser:
        """Build argument parser."""
        parser = ArgumentParser("Angelos Admin Utility")
        parser.add_argument("-v", "--verbose", dest="verbose", default=False, help="Verbose output")

        return parser

    def on_quit(self):
        """Override this method to act upon program quit."""
        pass

    def on_resize(self, columns: int, lines: int):
        """Override this method to act upon terminal resize."""
        print(columns, lines)

    def on_input(self, text: str):
        """Override this method to act upon user keypress."""
        print("Seq:", text)
        # self._terminal.data_received(text, str)

    async def _initialize(self):
        # self._terminal = Terminal()
        tty.setraw(self._tty)

        self._quiter = asyncio.Event()
        self._quiter.clear()

        asyncio.get_event_loop().add_signal_handler(
            signal.CTRL_C_EVENT if os.name == "nt" else signal.SIGINT, self._sigint_handler)
        asyncio.get_event_loop().add_signal_handler(signal.SIGWINCH, self._sigwinch_handler)

        asyncio.get_event_loop().add_reader(sys.stdin, self._input_handler)

        args = self._build_parser().parse_args()
        # self._terminal.session_started()
        # termios.tcsetattr(self._tty, termios.TCSADRAIN, self._no_echo)

    async def _finalize(self):
        termios.tcsetattr(self._tty, termios.TCSADRAIN, self._echo)

    async def run(self):
        """Application main loop."""
        await self._initialize()
        await self._quiter.wait()
        await self._finalize()

    def start(self):
        """Start application main loop."""
        try:
            asyncio.run(self.run())
        except KeyboardInterrupt:
            print("Uncaught keyboard interrupt")
        except Exception as exc:
            Util.print_exception(exc)

if __name__ == "__main__":
    Application().start()


from django.shortcuts import render
from .models import Question,Choice

def index(request):
    latest_question_list=Question.objects.order_by('-pub_date')[:5]
    contest={'latest_question_list':latest_question_list}
    return render(request,'polls/index.html',contest)

def detail(request,question_id):
    try:
        question=Question.objects.get(pk=question_id)
    except Question.DoesNotExist:
        raise Http404("Question does not exit")
    return  render(request,'polls/results.html',{'question':question})

def results(request,question_id):
    question=get_object_or_404(Question,pk=question_id)
    return  render(request,'polls/results.html',{'question':question})


from __future__ import division
import numpy as numpy
import cv2
import math as math

def prepareHough(frame, gray, video):

    kernel = numpy.ones((2, 2), numpy.uint8)

    if video == 'videos/video-3.avi' or video=='videos/video-8.avi':
        gray = cv2.erode(gray, kernel)
    #gray = cv2.morphologyEx(frame, cv2.MORPH_CLOSE, kernel)

    edges = cv2.Canny(gray,50,150,apertureSize = 3)
    minLineLength = 400
    maxLineGap = 5

    lines = cv2.HoughLinesP(edges,1,numpy.pi/180,50,minLineLength,maxLineGap)

    #a,b,c = lines.shape
    #for i in range(a):
        #cv2.line(frame, (lines[i][0][0], lines[i][0][1]), (lines[i][0][2], lines[i][0][3]), (0, 0, 255), 3, cv2.LINE_AA)
    #cv2.imshow('result', frame)
    #cv2.waitKey()
    #lineNumber = len(lines)
    #cv2.imshow('edges', edges)

    if video=='videos/video-3.avi':
        return appendLines3(lines,frame)
    return appendLines(lines, frame)

def appendLines(lines, frame):
    lineNumber = len(lines)
    resenFrame = frame

    lowerLinePoints = []
    upperLinePoints = []

    minx=lines[0][0][0]
    miny=lines[0][0][1]
    maxx=lines[0][0][2]
    maxy=lines[0][0][3]

    distance = math.sqrt(math.pow(minx-maxx,2) + math.pow(miny-maxy,2))

    #MIN I MAX OBE LINIJE ZAJEDNO
    for i in range(lineNumber-1):
        for j in range (lineNumber):
            x1 = lines[i][0][0]
            y1 = lines[i][0][1]
            x2 = lines[j][0][2]
            y2 = lines[j][0][3]

            tempd = math.sqrt(math.pow(x2-x1,2) + math.pow(y2-y1,2))
            if tempd > distance:
                distance = tempd
                minx = x1
                miny = y1
                maxx = x2
                maxy = y2

    middle_pointx = (minx+maxx)/2
    middle_pointy = (miny+maxy)/2

    print("MINIMUM IS : ", minx,miny)
    print("MAXIMUM IS : ", maxx, maxy)
    print("MIDDLE IS : ", middle_pointx, middle_pointy)

    middle_pointx = int(middle_pointx)
    middle_pointy = int(middle_pointy)
    cv2.circle(resenFrame, (middle_pointx, middle_pointy), 2, (0, 0, 255), 3)

    points1 = []
    points2 = []

    for i in range(lineNumber):
        x1 = lines[i][0][0]
        y1 = lines[i][0][1]

        list_item = lines[i][0][0], lines[i][0][1], lines[i][0][2], lines[i][0][3]
        if y1<=middle_pointy:
            points1.append(list_item)
        if abs(y1-middle_pointy)<30:
            points1.append(list_item)
        if abs(y1-middle_pointy)<20:
            points2.append(list_item)
        if y1>=middle_pointy:
            points2.append(list_item)

    minxp1 = points1[0][0]
    minyp1 = points1[0][1]
    maxxp1 = points1[0][2]
    maxyp1 = points1[0][3]

    for i in range(len(points1)):
        x1 = points1[i][0]
        y1 = points1[i][1]
        x2 = points1[i][2]
        y2 = points1[i][3]

        if x1<minxp1:
            minxp1 = x1
            minyp1 = y1
        if x2>maxxp1:
            maxxp1 = x2
            maxyp1 = y2
    cv2.line(resenFrame,(minxp1,minyp1),(maxxp1, maxyp1), (0, 255, 255),3)

    appendItem = minxp1, minyp1, maxxp1, maxyp1
    upperLinePoints.append(appendItem)

    minxp2 = points2[0][0]
    minyp2 = points2[0][1]
    maxxp2 = points2[0][2]
    maxyp2 = points2[0][3]

    for i in range(len(points2)):
        x1 = points2[i][0]
        y1 = points2[i][1]
        x2 = points2[i][2]
        y2 = points2[i][3]

        if x1 < minxp2:
            minxp2 = x1
            minyp2 = y1
        if x2 > maxxp2:
            maxxp2 = x2
            maxyp2 = y2
    cv2.line(resenFrame, (minxp2, minyp2), (maxxp2, maxyp2), (255, 0, 255), 3)
    cv2.imshow('Result',resenFrame)
    cv2.imwrite("lineDetected.jpg", resenFrame);

    appendItem = minxp2, minyp2, maxxp2, maxyp2
    lowerLinePoints.append(appendItem)

    cv2.line(frame,(minx, miny), (maxx,maxy), (0,255,255),3)
    cv2.waitKey()
    return lowerLinePoints, upperLinePoints

def appendLines3(lines,frame):
    lineNumber = len(lines)
    resenFrame = frame

    lowerLinePoints = []
    upperLinePoints = []

    x1k = lines[0][0][0]
    y1k = lines[0][0][1]
    x2k = lines[0][0][2]
    y2k = lines[0][0][3]

    k = ((y2k - y1k) / (x2k - x1k))
    n1 = y1k - k * x1k

    x1m = lines[lineNumber - 1][0][0]
    y1m = lines[lineNumber - 1][0][1]
    x2m = lines[lineNumber - 1][0][2]
    y2m = lines[lineNumber - 1][0][3]

    m = ((y2m - y1m) / (x2m - x1m))
    n2 = y1m - k * x1m

    first_line = []
    second_line = []

    minx = lines[0][0][0]
    miny = lines[0][0][1]
    maxx = lines[0][0][2]
    maxy = lines[0][0][3]

    for i in range(lineNumber):
        x1 = lines[i][0][0]
        y1 = lines[i][0][1]
        x2 = lines[i][0][2]
        y2 = lines[i][0][3]

        i_k = (y2 - y1) / (x2 - x1)
        i_n = y1-i_k*x1

        list_item = lines[i][0][0], lines[i][0][1], lines[i][0][2], lines[i][0][3]
        if abs(n1-i_n)<10:
            first_line.append(list_item)
        if abs(n1-i_n)>20:
            second_line.append(list_item)

        #cv2.line(frame, (x1, y1), (x2, y2), (0, 0, 255), 3)
        #cv2.imshow('test',frame);
        #cv2.waitKey()
    sminx = second_line[0][0]
    sminy = second_line[0][1]
    smaxx = second_line[0][2]
    smaxy = second_line[0][3]

    first_line_lenght = len(first_line)
    second_line_lenght = len(second_line)

    for i in range(second_line_lenght):
        x1 = second_line[i][0]
        y1 = second_line[i][1]
        x2 = second_line[i][2]
        y2 = second_line[i][3]
        if x1<sminx:
            sminx = x1
            sminy = y1
        if x2>smaxx:
            smaxx=x2
            smaxy=y2
    cv2.line(frame,(sminx, sminy), (smaxx,smaxy), (0,255,255),3)
    appendItem = sminx, sminy, smaxx, smaxy
    upperLinePoints.append(appendItem)
    #cv2.imshow("linije", frame)
    #cv2.waitKey()

    ssminx = first_line[0][0]
    ssminy = first_line[0][1]
    ssmaxx = first_line[0][2]
    ssmaxy = first_line[0][3]

    for i in range(first_line_lenght):
        x1 = first_line[i][0]
        y1 = first_line[i][1]
        x2 = first_line[i][2]
        y2 = first_line[i][3]
        if x1 < ssminx:
            ssminx = x1
            ssminy = y1
        if x2 > ssmaxx:
            ssmaxx = x2
            ssmaxy = y2
    cv2.line(frame, (ssminx, ssminy), (ssmaxx, ssmaxy), (255, 0, 255), 3)
    appendItem = ssminx, ssminy, ssmaxx, ssmaxy
    lowerLinePoints.append(appendItem)
    cv2.imshow("linije", frame)
    cv2.waitKey()

    return lowerLinePoints, upperLinePoints


import pickle
import os
import random
import numpy as np
from config import BASE_PATH

class DataContainer:
    def __init__(self, pretrained_emb_path,
                 lexicon_path, corpus_vocab_path, down_sampling=False,
                 seed=101, name=None, verbose=True):
        if down_sampling:
            name = '{}_balanced'.format(name)
        self.name = name
        self.seed = seed
        self.verbose = verbose
        self.down_sampling = down_sampling

        data = self.load_from_cache(name=name, pretrained_emb_path=pretrained_emb_path,
                                    lexicon_path=lexicon_path, corpus_vocab_path=corpus_vocab_path)

        self.lexicon_words, self.non_lexicon_words, self.lexicon, self.init_emb = data
        _vec = next(iter(self.init_emb.values()))
        self.emb_shape = _vec.shape

        self.pos_words, self.neg_words = self.init_poles()

        _scores = [self.lexicon[key] for key in self.lexicon]
        if len(set(_scores)) == 2:
            self.binary = True
        else:
            self.binary = False

        if self.verbose:
            print('Num of lexicon words:', len(self.lexicon_words))
            print('Class Distribution of the lexicon words:')
            print('Positive:', len(self.pos_words), 'Negative:', len(self.neg_words),
                  'Neutral:', len(self.lexicon_words) - len(self.pos_words) - len(self.neg_words))

            _len_nl = 0
            for _bth in self.non_lexicon_words:
                _len_nl += len(_bth)
            print('Num of non-lexicon words:', _len_nl)

            print('The shape of the word representations is:', self.emb_shape, '!')

            if self.binary:
                print('***The lexicon is binary!***')
            else:
                print('***The lexicon is not binary!***')

        # self.init_2nd_emb = None
        self.init_1st_emb = dict(self.init_emb)

    def update_emb(self, embeddings):
        for word in embeddings:
            self.init_emb[word] = embeddings[word]
        # self.init_2nd_emb = dict(self.init_emb)

    def rescale(self, num, _min, _max):
        random.seed(self.seed)
        res = (num - _min) * 2 / (_max - _min) - 1
        if res == 0:
            res += random.choice([0.000001, -0.000001])
            self.seed += 1
        return res

    @staticmethod
    def batchify(words, batch_size):
        assert batch_size > 1
        batches = []
        batch = []
        count = 0
        for word in words:
            batch.append(word)
            count += 1
            if count > batch_size:
                batches.append(batch)
                batch = []
                count = 0
        batches.append(batch)
        return batches

    def preprocess(self, pretrained_emb_path, lexicon_path, corpus_vocab_path):
        '''
        cross compare vocabularies and initialize lexicon words, non lexicon words;
        rescale the lexicon scores to [-1, 1]
        cache the preprocessed data if wanted

        :param pretrained_emb_path:
        :param lexicon_path:
        :param corpus_vocab_path:
        :return:
        '''
        pretrained_emb = pickle.load(open(pretrained_emb_path, 'rb'))
        lexicon = pickle.load(open(lexicon_path, 'rb'))
        corpus_vocab = pickle.load(open(corpus_vocab_path, 'rb'))

        inter_corpus_vocab = list(set(pretrained_emb.keys()) & set(corpus_vocab.keys()))

        lexicon_words = list(set(pretrained_emb.keys()) & set(lexicon.keys()))
        non_lexicon_words = list(set(inter_corpus_vocab) - set(lexicon_words))

        random.seed(self.seed)
        random.shuffle(lexicon_words)
        random.shuffle(non_lexicon_words)

        _scores = [lexicon[key] for key in lexicon]
        _min = np.floor(min(_scores))
        _max = np.ceil(max(_scores))

        scaled_lexicon = {}
        for word in lexicon_words:
            scaled_lexicon[word] = self.rescale(lexicon[word], _min, _max)

        # down sampling
        if self.down_sampling:
            print('Down Sampling...')
            _pos_word = []
            _neg_word = []
            _neutral_word = []
            for word in lexicon_words:
                if scaled_lexicon[word] > 0.001:
                    _pos_word.append(word)
                elif scaled_lexicon[word] < -0.001:
                    _neg_word.append(word)
                else:
                    _neutral_word.append(word)
            _maxLen = min(len(_pos_word), len(_neg_word))

            _down_sampled = _pos_word[:_maxLen] + _neg_word[:_maxLen] + _neutral_word
            lexicon_words = _down_sampled

        full_vocab = list(set(lexicon_words) | set(non_lexicon_words))
        init_emb = {}
        for word in full_vocab:
            init_emb[word] = pretrained_emb[word]

        batch_non_lexicon_words = self.batchify(non_lexicon_words, len(lexicon_words) / 9)

        data = (lexicon_words, batch_non_lexicon_words, scaled_lexicon, init_emb)
        return data

    def init_poles(self):
        '''
        initialize the poles and save to the data container

        :return:
        '''
        pos_words = []
        neg_words = []
        _agg_pos_emb = np.zeros(shape=self.emb_shape)
        _agg_neg_emb = np.zeros(shape=self.emb_shape)
        _num_pos = 0
        _num_neg = 0

        for word in self.lexicon_words:
            if self.lexicon[word] > 0.001:
                pos_words.append(word)
                _agg_pos_emb += self.init_emb[word]
                _num_pos += 1
            elif self.lexicon[word] < -0.001:
                neg_words.append(word)
                _agg_neg_emb += self.init_emb[word]
                _num_neg += 1

        assert _num_pos > 0
        assert _num_neg > 0
        _pos_pole = _agg_pos_emb / _num_pos
        _neg_pole = _agg_neg_emb / _num_neg
        self.init_emb['POSPOLE'] = _pos_pole
        self.init_emb['NEGPOLE'] = _neg_pole
        self.lexicon['POSPOLE'] = 1
        self.lexicon['NEGPOLE'] = -1
        self.lexicon_words.append('POSPOLE')
        self.lexicon_words.append('NEGPOLE')
        pos_words.append('POSPOLE')
        neg_words.append('NEGPOLE')

        return pos_words, neg_words

    @staticmethod
    def _check_cache():
        cache_dir = os.path.join(BASE_PATH, "_cache")
        if not os.path.exists(cache_dir):
            os.makedirs(cache_dir)

    def _get_cache_filename(self):
        return os.path.join(BASE_PATH, "_cache",
                            "_cachedData_{}.pkl".format(self.name))

    def _write_cache(self, data):
        self._check_cache()

        cache_file = self._get_cache_filename()

        with open(cache_file, 'wb') as pickle_file:
            pickle.dump(data, pickle_file)

    def load_from_cache(self, name, pretrained_emb_path,
                        lexicon_path, corpus_vocab_path):

        # NOT using cache
        if name is None:
            print("cache deactivated for data loading!")
            return self.preprocess(pretrained_emb_path=pretrained_emb_path,
                                   lexicon_path=lexicon_path,
                                   corpus_vocab_path=corpus_vocab_path)

        # using cache
        cache_file = self._get_cache_filename()

        if os.path.exists(cache_file):
            print("Loading {} from cache ...".format(self.name))
            with open(cache_file, 'rb') as f:
                return pickle.load(f)
        else:
            print("No cache file for {} ...".format(self.name))
            data = self.preprocess(pretrained_emb_path=pretrained_emb_path,
                                   lexicon_path=lexicon_path,
                                   corpus_vocab_path=corpus_vocab_path)
            self._write_cache(data)
            return data


from typing import List
from torch import Tensor
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torch.utils.data import DataLoader, Dataset
from torchvision import models
from config import Config
import time
import copy
from tqdm import tqdm

def train_model(model: models.ResNet, train_dataloader: DataLoader, train_dataset: Dataset):
    """modelオブジェクトとDataLoaderオブジェクトを渡して、モデルの学習を行う.

    Parameters
    ----------
    model : models.ResNet
        _description_
    train_dataloader : DataLoader
        _description_
    """
    print('len of train data is {}'.format(train_dataset.__len__()))
    # 演算を行うデバイスを設定
    # GPUを使えるかどうかに応じてtorch.deviceオブジェクトを生成
    # torch.deviceはテンソルをどのデバイスに割り当てるかを表すクラス。
    device = torch.device(
        'cuda') if torch.cuda.is_available() else torch.device('cpu')
    # modelをGPUに渡す
    model.to(device)

    # 損失関数の設定
    criterion = nn.CrossEntropyLoss()
    # 最適化手法の設定
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer_ft = optim.SGD(params=params,
                             lr=Config.learning_rate_resnet,
                             momentum=Config.momentum_resnet
                             )
    # 7エポックごとにLRを0.1倍ずつ減衰させる。
    exp_lr_scheduler = optim.lr_scheduler.StepLR(
        optimizer=optimizer_ft, step_size=7, gamma=0.1
    )

    # GPUのキャッシュクリア
    torch.cuda.empty_cache()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    num_epochs = 25

    since = time.time()

    # 学習
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch+1, num_epochs))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                exp_lr_scheduler.step()
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in train_dataloader:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer_ft.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    # tensor(max, max_indices)なのでpredは0,1のラベル
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()  # 誤差逆伝搬で、各パラメータの勾配gradの値を計算(実は累算してる!だからzero_gradを使ってる)
                        optimizer_ft.step()  # - grad＊学習率を使って、パラメータを更新

        print(f'epoch {epoch}/{num_epochs} : loss {loss}')

    time_elapsed = time.time() - since

    # save weight
    model_path = 'model.pth'
    # model.state_dict()として保存した方が無駄な情報を削れてファイルサイズを小さくできるらしい.
    torch.save(model.state_dict(), model_path)

    return model