text:
'''
TODO: Write a function to calculate the initial flux errors (to be used in the spline
weighting) by heavily smoothing the spectrum and calculating the stddev of the points
around the smoothed flux
'''
import os
from collections import namedtuple

from astropy.io import fits
from astropy.io import ascii as asc
from astropy.table import Table
from astropy.modeling import models,fitting
from astropy.convolution import convolve, Box1DKernel
from astropy.time import Time
import numpy as np
from scipy import signal, interpolate
from matplotlib import pyplot as plt
import matplotlib.collections as collections
from matplotlib.backends.backend_pdf import PdfPages

from utilities_az import spectroscopy as spec

endpoint = namedtuple('endpoint', ['wave', 'flux', 'error'])

FIG_DIR = '../figures'

def read_iraf_spectrum(filename, redshift=0.0069):
    ofile = fits.open(filename)
    flux = ofile[0].data[0,0,:]
    err = ofile[0].data[3,0,:]
    wave = spec.calc_wavelength(ofile[0].header, np.arange(len(flux))+1)
    rest_wave = spec.apply_redshift(wave, redshift)
    return(spec.spectrum1d(rest_wave, flux, error=err))

def smooth_signal(flux, width, poly_deg):
    smoothed_flux = signal.savgol_filter(flux, width, poly_deg)
    return smoothed_flux

def find_blue_edge(wave, flux, wcenter, binsize, wmin=None):
    '''
    Calculate the slope in each bin starting at wmin, until the bin changes sign, use center for blue_edge
    '''
    wcenter_indx = np.argmin(np.abs(wave-wcenter))
    ifit = np.polyfit(wave[wcenter_indx-binsize: wcenter_indx+1],flux[wcenter_indx-binsize: wcenter_indx+1], 1)
    slope_product = 1
    #plt.plot(wave, flux)
    #plt.xlim(wmin, wcenter)
    if wmin is None:
        search_indx = np.arange(binsize,wcenter_indx+1)
    else:
        min_indx = np.argmin(np.abs(wave - wmin))
        search_indx = np.arange(min_indx, wcenter_indx)
    for indx in search_indx[::-1]:
        last_slope = ifit[0]
        if indx-binsize < 0:
            break
        ifit = np.polyfit(wave[indx-binsize:indx+1], flux[indx-binsize:indx+1], 1)
        #plt.plot(wave[indx-binsize:indx+1], flux[indx-binsize:indx+1])
        #plt.plot(wave[indx-binsize:indx+1], np.polyval(ifit, wave[indx-binsize:indx+1]))
        slope_product = last_slope*ifit[0] #if this is negative then the slope has changed sign
        if slope_product < 0:
            break
    edge_indx = indx - binsize//2
    return edge_indx, wave[edge_indx]

def find_red_edge(wave, flux, wcenter, binsize, wmax = None):
    '''
    Calculate the slope in each bin starting at wmin, until the bin changes sign, use center for red_edge
    binsize is in pixels
    '''
    wcenter_indx = np.argmin(np.abs(wave-wcenter))
    #fig = plt.figure()
    #ax1 = fig.add_subplot(2,1,1)
    #ax2 = fig.add_subplot(2,1,2, sharex=ax1)
    #ax1.plot(wave, flux)
    #ax2.plot(wave, wave-wcenter)
    ifit = np.polyfit(wave[wcenter_indx:wcenter_indx+binsize+1],flux[wcenter_indx:wcenter_indx+binsize+1], 1)
    slope_product = 1
    #plt.plot(wave, flux)
    #plt.xlim(wmin, wcenter)
    #plt.axvline(wave[wcenter_indx], color='y')
    if wmax is None:
        search_indx = np.arange(wcenter_indx, len(flux))
    else:
        max_indx = np.argmin(np.abs(wave - wmax))
        search_indx = np.arange(wcenter_indx, max_indx+1)
    #plt.plot(wave[search_indx], flux[search_indx])
    for indx in search_indx:
        last_slope = ifit[0]
        ifit = np.polyfit(wave[indx:indx+binsize+1], flux[indx:indx+binsize+1], 1)
        slope_product = last_slope*ifit[0] #if this is negative then the slope has changed sign
        #plt.plot(wave[indx:indx+binsize+1], flux[indx:indx+binsize+1])
        #plt.plot(wave[indx:indx+binsize+1], np.polyval(ifit, wave[indx:indx+binsize+1]))
        slope_product = last_slope*ifit[0] #if this is negative then the slope has changed sign
        if slope_product < 0:
            break
    edge_indx = indx + binsize//2
    return edge_indx, wave[edge_indx]

def check_max(wave, flux, edge_indx, binsize, absorption=True):
    '''
    Fit a quadratic and verify that it is the correct direction
    '''
    wmin_indx = edge_indx - binsize//2
    wmax_indx = edge_indx + binsize//2
    fit = np.polyfit(wave[wmin_indx:wmax_indx+1], flux[wmin_indx:wmax_indx+1], 2)
    if fit[0]>0:
        concavity = 'up'
    if fit[0]<0:
        concavity = 'down'
    if ((absorption is True) and (concavity is 'down')) or ((absorption is False) and (concavity is 'up')):
        good_fit = True
    else:
        good_fit = False
    return good_fit

def calc_rmse(data, model):
    rmse = np.sqrt(np.sum((model-data)**2)/len(data))
    print('rmse calculated over {} points'.format(len(data)))
    return rmse

def find_boundary(wave, flux, wmin, wmax, binsize, visualize=False):
    wmin_indx = np.argmin(np.abs(wave-wmin))
    wmax_indx = np.argmin(np.abs(wave-wmax))
    slope_product = []
    ifit = np.polyfit(wave[wmin_indx-binsize//2:wmin_indx+binsize//2+1], flux[wmin_indx-binsize//2:wmin_indx+binsize//2+1], 1)
    search_indx = np.arange(wmin_indx, wmax_indx+1)
    if visualize:
        fig = plt.figure()
        ax1 = fig.add_subplot(2,1,1)
        ax2 = fig.add_subplot(2,1,2, sharex=ax1)
        ax1.set_xlim(wmin, wmax)
    for indx in search_indx:
        last_slope = ifit[0]
        ifit = np.polyfit(wave[indx-int(binsize//2):indx+int(binsize//2)+1], flux[indx-int(binsize//2):indx+int(binsize//2)+1], 1)
        slope_product.append(last_slope*ifit[0]) #if this is negative then the slope has changed sign
        if visualize:
            ax1.plot(wave[indx-int(binsize//2):indx+int(binsize//2)+1], np.polyval(ifit, wave[indx-int(binsize//2):indx+int(binsize//2)+1]))
    slope_product = np.array(slope_product)
    slope_change_indx = search_indx[slope_product<0]
    if visualize:
        ax1.set_title('Slope Plot binsize={}, fit_wmin={:4.2f}, fit_wmax={:4.2f}'.format(binsize, wave[wmin_indx-binsize//2], wave[indx+int(binsize//2)]))
        ax1.plot(wave, flux)
        ax2.plot(wave[search_indx], slope_product)
        ax2.axhline(0, color='k', linestyle=':')
    if len(slope_change_indx) == 3:
        blue_edge_indx, wcenter_indx, red_edge_indx = slope_change_indx
        return blue_edge_indx, wcenter_indx, red_edge_indx
    else:
        return None, None, None

def determine_error_binsize(wave, wave_bin=100):
    '''
    We should be calculating noise over the same wavelength range
    rather than the same number of pixels as long as one wavelength
    bin includes enough pixels. Set binsize to be 100A. If there are
    fewer than 10 pixels in 100A (dispersion is greater than 10A/pix)
    then issue a warning and make binsize 10 pixels regardless of how
    many angstroms this represents

    wave: array of wavelengths
    wave_bin: size of bin in angstroms

    outputs: binsize in pixels

    Note: right now this calculates the dispersion for the full wavelength range.
    For a grating/grism with a large variation in dispersion, it might make sense to
    just calculate this over the feature wavelength range.
    '''
    dispersion = np.median(wave[1:]-wave[:-1])
    binsize = np.ceil(wave_bin/dispersion)
    if binsize < 10:
        print('WARNING: dispersion = {}, \
        leading to binsize < 10 for {}$\AA$ bins, \
        setting binsize=10, making wave_bin={}'.format(dispersion, wave_bin, 10*dispersion))
        binsize=10
    return binsize

def define_continuum(wave, flux, edge_indx, binsize, err_binsize, absorption=True, visualize=False):
    '''
    Fit a quadratic and verify that it is the correct direction
    '''
    #Silverman says: "Once these two endpoints are determined, a quadratic function is
    # fit to the data in wavelength bins centred on each endpoint."
    #Let's start with fitting over 2 wavelength bins?
    wmin_indx = edge_indx - int(np.floor(1*binsize))
    wmax_indx = edge_indx + int(np.ceil(1*binsize))
    quad_model = models.Polynomial1D(degree=2)
    fitter = fitting.LinearLSQFitter()
    fit = fitter(quad_model, wave[wmin_indx:wmax_indx+1], flux[wmin_indx:wmax_indx+1])
    fit_extreme_wl = -fit.c1.value/(2*fit.c2.value)

    #calc rmse over edge_indx +/- 20 pixels
    wmin_rmse = edge_indx - int(err_binsize//2)
    wmax_rmse = edge_indx + int(err_binsize//2)
    rmse = calc_rmse(flux[wmin_rmse:wmax_rmse], fit(wave[wmin_rmse:wmax_rmse]))

    if fit.c2.value>0:
        concavity = 'up'
    if fit.c2.value<0:
        concavity = 'down'

    if ((absorption is True) and (concavity is 'down')) or ((absorption is False) and (concavity is 'up')):
        good_fit = True
    else:
        good_fit = False
    if visualize:
        fig = plt.figure()
        ax = fig.add_subplot(1,1,1)
        ax.plot(wave, flux)
        ax.set_xlim(wave[edge_indx-4*binsize], wave[edge_indx+4*binsize])
        ax.plot(wave[wmin_rmse:wmax_rmse], fit(wave[wmin_rmse:wmax_rmse]), label='RMSE range')
        ax.plot(wave[wmin_indx:wmax_indx+1], fit(wave[wmin_indx:wmax_indx+1]), label='fit range')
        ax.axvline(wave[edge_indx], label='input continuum', color='k')
        ax.errorbar(fit_extreme_wl, fit(fit_extreme_wl), yerr=rmse, fmt='.', label='Edge w/error', zorder=10 )
        ax.set_ylim(0.9*np.min(flux[wmin_rmse:wmax_rmse]), 1.1*np.max(flux[wmin_rmse:wmax_rmse]))
        ax.legend(loc='best')
        ax.set_title('absorption={}, concavity={}, good_fit={}'.format(absorption, concavity, good_fit))
    endpt = endpoint(fit_extreme_wl, fit(fit_extreme_wl), rmse)
    return good_fit, endpt

def calc_pseudo_ew(wave, flux, continuum_l, continuum_r, absorption=True, visualize=False):
    '''
    wave: array
        array of wavelength (can be whole spectrum)
    flux: array
        array of fluxes (can be whole spectrum)

    * Create a fit to the continuum and define the continuum for each wavelength in wave
    * Use continuum wavelengths to define index location of feature
    * Calc pseudo equivalent width using flux, continuum, and delta wave as calculated from the
    wave array
    '''
    fitter = fitting.LinearLSQFitter()
    lin_mod = models.Linear1D()
    continuum_fit = fitter(lin_mod,[continuum_l.wave, continuum_r.wave], [continuum_l.flux, continuum_r.flux])
    line_indx = np.int_(np.arange(len(wave))[(wave>=continuum_l.wave)&(wave<=continuum_r.wave)])
    continuum = continuum_fit(wave[line_indx])
    delta_lambda = wave[line_indx]-wave[line_indx-1]
    if absorption is True:
        pew = np.sum(delta_lambda*(continuum - flux[line_indx])/continuum)
    else:
        pew = np.sum(delta_lambda*(flux[line_indx]-continuum)/continuum) #Check that this is true
    if visualize is True:
        fig = plt.figure()
        ax1 = fig.add_subplot(1,2,1)
        ax2 = fig.add_subplot(1,2,2)
        ax2.axhline(1, color='k')
        ax1.plot(wave, flux)
        ax1.plot(wave[line_indx], flux[line_indx], label='data')
        ax1.plot(wave[line_indx], continuum, label='continuum')
        ax1.set_xlim(continuum_l.wave-10, continuum_r.wave+10)
        if absorption is True:
            ax2.plot(wave[line_indx], (continuum - flux[line_indx])/continuum, label='sum for pEW')
        else:
            ax2.plot(wave[line_indx], (flux[line_indx]-continuum)/continuum, label='sum for pEW')
    return pew

def calc_continuum(wave, continuum_l, continuum_r):
    fitter = fitting.LinearLSQFitter()
    lin_mod = models.Linear1D()
    continuum_fit = fitter(lin_mod,[continuum_l.wave, continuum_r.wave], [continuum_l.flux, continuum_r.flux])
    continuum = continuum_fit(wave)
    return continuum

def find_velocity(wave, flux, error, wcenter, continuum_l, continuum_r, binsize, visualize=False):
    line_indx = np.int_(np.arange(len(wave))[(wave>=continuum_l.wave)&(wave<=continuum_r.wave)])
    windx_min = int(line_indx[0]-binsize//2)
    windx_max = int(line_indx[-1]+binsize//2)
    fitter = fitting.LinearLSQFitter()
    lin_mod = models.Linear1D()
    continuum_fit = fitter(lin_mod,[continuum_l.wave, continuum_r.wave], [continuum_l.flux, continuum_r.flux])
    continuum = continuum_fit(wave[windx_min:windx_max])
    weight = 1./error[windx_min:windx_max]
    fit = interpolate.UnivariateSpline(wave[windx_min:windx_max], flux[windx_min:windx_max]-continuum, w=weight)
    if visualize is True:
        fig = plt.figure()
        ax = fig.add_subplot(1,1,1)
        ax.plot(wave[windx_min:windx_max], flux[windx_min:windx_max]-continuum)
        s = np.sum((weight * ((flux[windx_min:windx_max]-continuum)-fit(wave[windx_min:windx_max])))**2)
        ax.errorbar(wave[windx_min:windx_max], flux[windx_min:windx_max]-continuum, error[windx_min:windx_max], fmt='.', label='spectrum', zorder=1, color='b')
        ax.plot(wave[windx_min:windx_max], flux[windx_min:windx_max]-continuum, zorder=2, color='r')
        ax.plot(wave[windx_min:windx_max], fit(wave[windx_min:windx_max]), label='fit, s={:2.2f}, len(w)={:2.2f}, med(std)={:2.2e}'.format(s, len(weight), np.median(error[windx_min:windx_max])), color='gold', zorder=3)
        min_wave = wave[line_indx][np.argmin(fit(wave[line_indx]))]
        ax.axvline(min_wave)
        knots = fit.get_knots()
        ax.vlines(knots, ymin=ax.get_ylim()[0], ymax=ax.get_ylim()[1], linestyle=':')
        ax.legend(loc='best')
    return fit

def calc_flux_variance(data, model, err_binsize):
    kernel = Box1DKernel(err_binsize)
    errors = convolve((data-model)**2, kernel, boundary=None)
    errors = errors
    errors = np.trim_zeros(errors)
    return errors

def calc_continuum_variance(wave, continuum_l, continuum_r):
    var = (1./(continuum_l.wave - continuum_r.wave))**2 * \
          ((wave - continuum_r.wave)**2 * continuum_l.error**2 +
           (wave - continuum_l.wave)**2 * continuum_r.error**2)
    return var

def calc_pew_variance(flux, continuum, delta_wave, flux_var, continuum_var, visualize=False, wave=None):
    '''
    Calculate the variance of the equivalent width
    Parameters:
    -----------
        flux: array
            flux values over which equivalent width is calculated
        continuum: array
            continuum values over which equivalent width is calculated
        delta_wave: int
            the wavelength bin size (in angstroms) used in the equivalent width calculation
        flux_var: array
            variance in the flux
        continuum_var: array
            variance in the continuum
    Output:
        variance in the equivalent width
    '''
    pew_var_indiv = ((flux/(continuum**2)*delta_wave)**2 * continuum_var) + \
                    ((delta_wave/continuum)**2*flux_var)
    pew_var = np.sum(pew_var_indiv)
    if visualize is True:
        fig = plt.figure()
        ax = fig.add_subplot(1,1,1)
        if wave is not None:
            ax.errorbar(wave, (continuum-flux)/continuum, np.sqrt(pew_err))
        else:
            print('wavelength not defined')
            wave = np.arange(len(flux))
            ax.errorbar(wave,(continuum-flux)/continuum, np.sqrt(pew_err))
    return pew_var

def calc_velocity_error(wave, flux, vel_fit, continuum = None, visualize=False):
    '''
    Calculate 1-sigma errors
    '''
    min_indx = np.argmin(vel_fit(wave))
    indx = np.argsort(wave - wave[min_indx])
    min_indx_indx = int(indx[indx==min_indx])
    cum_sum_right = 0
    if continuum is None:
        flux_sub = flux
    else:
        flux_sub = flux - continuum
    total = np.sum((flux_sub))
    for i in indx[min_indx_indx:]:
        cum_sum_right += (flux_sub)[i]
        if cum_sum_right/total > .341:
            break
    right_err = wave[i]-wave[min_indx]

    cum_sum_left = 0
    j=0
    for j in indx[:min_indx_indx][::-1]:
        cum_sum_left += (flux_sub)[j]
        if cum_sum_left/total > .341:
            break
    left_err = wave[min_indx]-wave[j]
    if visualize is True:
        from visualization import make_color_wheel
        colors = make_color_wheel(wave)
        plt.figure()
        for c, ind in zip(colors, indx):
            plt.plot(wave[ind], flux[ind], marker='o', ls='none', color=c)
        for k in indx[min_indx_indx:i]:
            plt.plot(wave[k], flux[k], marker='s', ls='none', color=colors[k])
        for k in indx[j:min_indx_indx]:
            plt.plot(wave[k], flux[k], marker='s', ls='none', color=colors[k])
        plt.axvline(wave[i], label='1 $\sigma$ right error')
        plt.axvline(wave[j], label='1 $\sigma$ left error')
        plt.xlabel('Wavelength')
        plt.ylabel('Flux')
        plt.legend(loc='best')
    return left_err, right_err

def find_edges(spectrum, feature_dict, smooth_flux, filename, vis=False, pp=None):
    '''
    Fit incremental slopes to find where spectrum turns over
    Fit quadratic to the turnover points to get spectrum edges
    '''
    #Estimate the edges and center of the feature
    blue_edge_indx = None
    red_edge_indx = None
    good_fit_blue = False
    good_fit_red = False

    adjust_binsize = feature_dict['edge_param']['binsize']
    #increase the binsize until only 3 turning points are found (edges and center)
    wmin, wmax = find_wavelength_range(feature_dict, filename)
    npts_feature = len(spectrum.wave[(spectrum.wave>=wmin) & (spectrum.wave <= wmax)])
    while ((blue_edge_indx is None) or (red_edge_indx is None) or
            (good_fit_blue is False) or (good_fit_red is False)) and \
          (adjust_binsize < 0.4*npts_feature) and \
          (adjust_binsize < feature_dict['edge_param']['binmax']):  #TODO figure out a cutoff for this
        if plt.get_fignums() is not False:
            for ifig in plt.get_fignums():
                plt.close(ifig)
        adjust_binsize += 2
        blue_edge_indx, wcenter_indx, red_edge_indx = find_boundary(spectrum.wave,
                                                                    smooth_flux,
                                                                    wmin,
                                                                    wmax,
                                                                    adjust_binsize,
                                                                    visualize=vis)
        if (blue_edge_indx is not None) and (red_edge_indx is not None):
            err_binsize = determine_error_binsize(spectrum.wave, wave_bin=100)
            #Find the feature edges and errors
            good_fit_blue, continuum_l = define_continuum(spectrum.wave, smooth_flux, blue_edge_indx, feature_dict['edge_param']['concavity_binsize'], err_binsize, absorption=True, visualize=vis)
            good_fit_red, continuum_r = define_continuum(spectrum.wave, smooth_flux, red_edge_indx, feature_dict['edge_param']['concavity_binsize'], err_binsize, absorption=True, visualize=vis)
            if continuum_l.wave > continuum_r.wave:
                print('**** WARNING: {}, left edge {} is greater than right edge {}****'.format(os.path.basename(filename), continuum_l.wave, continuum_r.wave))
                return None

    if adjust_binsize > 0.4*npts_feature:
        blue_edge_indx, wcenter_indx, red_edge_indx = find_boundary(spectrum.wave,
                                                                    smooth_flux,
                                                                    wmin,
                                                                    wmax,
                                                                    feature_dict['edge_param']['binsize'],
                                                                    visualize=True)
        print('Unable to find edges for {}, {}'.format(feature_dict['name'],os.path.basename(filename)))
        import pdb; pdb.set_trace()
        return None
    else:
        print('filename = ',os.path.basename(filename))
        print('\tinput binsize = ', feature_dict['edge_param']['binsize'])
        print('\tadjusted binsize = ', adjust_binsize)
        print('good_fit_blue={}, good_fit_red={}, combine={}'.format(good_fit_blue, good_fit_red, ((good_fit_blue is False) or (good_fit_red is False))))
        if vis is True:
            fig1 = plt.figure(1)
            pp.savefig(fig1)
            plt.close(fig1)
            fig2 = plt.figure(2)
            pp.savefig(fig2)
            plt.close(fig2)
            fig3 = plt.figure(3)
            pp.savefig(fig3)
            plt.close(fig3)
    return err_binsize, blue_edge_indx, red_edge_indx, wcenter_indx, continuum_l, continuum_r

def find_wavelength_range(feature_dict, filename):
    date = Time(fits.getval(filename, 'date-obs', 0))
    phase = date.jd - feature_dict['texpl']
    delta_wave = feature_dict['slope']*phase
    wmin = feature_dict['wmin']+delta_wave
    wmax = feature_dict['wmax']+delta_wave
    return wmin, wmax

def final_plot(wave, flux, flux_err, continuum_l, continuum_r, vel_fit, vel_min, vel_err, pew, pew_err):
    continuum = calc_continuum(wave, continuum_l, continuum_r)
    plt.ioff()
    fig = plt.figure()
    ax = fig.add_subplot(1,1,1)
    #ax.plot(wave, flux/continuum, label='Spectrum')
    ax.errorbar(wave, flux/continuum, flux_err/continuum, fmt='.', label='Spectrum')
    ax.errorbar(np.array([continuum_l.wave, continuum_r.wave]),
                np.array([continuum_l.flux, continuum_r.flux])/calc_continuum(np.array([continuum_l.wave, continuum_r.wave]), continuum_l, continuum_r),
                np.array([continuum_l.error, continuum_r.error])/calc_continuum(np.array([continuum_l.wave, continuum_r.wave]), continuum_l, continuum_r),
                fmt='o', label='Feature Edges')
    min_continuum = calc_continuum(np.array([vel_min]), continuum_l, continuum_r)[0]
    ax.plot(wave, (vel_fit(wave)+continuum)/continuum, label='Velocity fit')
    ax.errorbar(vel_min, np.min((vel_fit(vel_min)+min_continuum)/min_continuum), xerr=(vel_err,), fmt='s', label='Velocity')
    pew_collection_err = collections.BrokenBarHCollection.span_where(wave, ymin=0, ymax=1,
                                                        where= (wave>=vel_min-pew/2-np.sqrt(pew_err))&(wave<=vel_min+pew/2+np.sqrt(pew_err)),
                                                        color='k', alpha=0.1)
    pew_collection = collections.BrokenBarHCollection.span_where(wave, ymin=0, ymax=1,
                                                        where= (wave>=vel_min-pew/2)&(wave<=vel_min+pew/2),
                                                        color='k', alpha=0.1, label = 'pEW')
    ax.add_collection(pew_collection_err)
    ax.add_collection(pew_collection)
    ax.legend(loc='best')
    ax.set_xlabel('Wavelength')
    ax.set_ylabel('Continuum subtracted flux')
    #ax.set_ylim(-0.05, 1.1)
    plt.ion()
    return fig

def characterize_line(feature_dict, filename, visualization_level=0):
    final_vis = False
    intermediate_vis = False
    pp = None
    if (visualization_level == 1) or (visualization_level ==2):
        final_vis = True
    if visualization_level == 2:
        intermediate_vis = True
        pp = PdfPages(os.path.join(FIG_DIR,
                                 'line_fit_intermed_{}_{}.pdf'.format(feature_dict['name'],
                                                                      os.path.basename(filename).split('.pdf')[0])))
    #Read in spectrum
    spectrum = read_iraf_spectrum(filename)
    #Remove CR and other large deviations
    smooth_flux = smooth_signal(spectrum.flux,
                                feature_dict['smooth_param']['width'],
                                feature_dict['smooth_param']['deg'])
    edge_results = find_edges(spectrum, feature_dict, smooth_flux, filename, vis=intermediate_vis, pp=pp)
    if edge_results is not None:
        err_binsize, blue_edge_indx, red_edge_indx, wcenter_indx, continuum_l, continuum_r = edge_results
        #Calculate the pseudo equivalent widths
        pew = calc_pseudo_ew(spectrum.wave, smooth_flux, continuum_l, continuum_r, visualize=intermediate_vis)
        if intermediate_vis is True:
            pp.savefig()
            plt.close()
        #Calculate the most common velocity
        wcenter = spectrum.wave[wcenter_indx]
        vel_fit = find_velocity(spectrum.wave, smooth_flux, spectrum.error, wcenter, continuum_l, continuum_r, err_binsize, visualize=intermediate_vis)
        if intermediate_vis is True:
            pp.savefig()
            plt.close()
        #Find the error in the pseudo equivalent width
        line_indx = np.arange(len(spectrum.wave))[(spectrum.wave>=continuum_l.wave)&(spectrum.wave<=continuum_r.wave)]
        min_indx = int(np.floor(line_indx[0]-err_binsize/2))
        max_indx = int(np.ceil(line_indx[-1]+err_binsize/2+1))
        continuum_extended = calc_continuum(spectrum.wave[min_indx:max_indx], continuum_l, continuum_r)
        flux_var = calc_flux_variance(spectrum.flux[min_indx:max_indx]-continuum_extended,
                                    vel_fit(spectrum.wave[min_indx:max_indx]), err_binsize) #These don't include the errors from the continuum subtraction yet; ok for EW calc
        if len(flux_var) > len(spectrum.flux[line_indx]):
            flux_var = flux_var[1:-1]
        continuum = calc_continuum(spectrum.wave[line_indx], continuum_l, continuum_r)
        continuum_var = calc_continuum_variance(spectrum.wave[line_indx], continuum_l, continuum_r)
        delta_wave = np.median(spectrum.wave[1:]-spectrum.wave[:-1])
        pew_var = calc_pew_variance(spectrum.flux[line_indx], continuum, delta_wave, flux_var, continuum_var, wave=spectrum.wave[line_indx])
        #Find the velocity error
        vel_err = calc_velocity_error(spectrum.wave[line_indx], spectrum.flux[line_indx], vel_fit, continuum=continuum)
        vel_min = spectrum.wave[line_indx][np.argmin(vel_fit(spectrum.wave[line_indx]))]
        if final_vis is True:
            fig = final_plot(spectrum.wave[min_indx:max_indx], spectrum.flux[min_indx:max_indx],spectrum.error[min_indx:max_indx], continuum_l, continuum_r, vel_fit, vel_min, vel_err, pew, pew_var)
            fig.suptitle(os.path.basename(filename))
        if intermediate_vis:
            pp.close()
        return pew, pew_var, vel_min, vel_err, fig

reason: import numpy,from scipy,from astropy


text:
import numpy as np
import matplotlib.pyplot as plt

N = 6
Accuracy = (53.76,47.58,24.52,58.37,41.73,49.64)
Error = (100 - 53.76,100 - 47.58,100 - 24.52,100 - 58.37,100 - 41.73,100 - 49.64)
ind = np.arange(N)  # the x locations for the groups
width = 0.35       # the width of the bars

fig, ax = plt.subplots()
rects1 = ax.bar(ind, Accuracy, width, color='b')
rects2 = ax.bar(ind + width, Error, width, color='y')

# add some text for labels, title and axes ticks
ax.set_ylabel('Accuracy')
ax.set_title('Classifier Vs Accuracy')
ax.set_xticks(ind+width)
ax.set_xticklabels(('LogReg', 'DecisionTree', 'NaiveBayes', 'RandForest', 'SVM', 'SVM+Boosting'))
ax.legend((rects1[0], rects2[0]), ('Accuracy', 'Error'))

def autolabel(rects):
    # attach some text labels
    for rect in rects:
        height = rect.get_height()
        ax.text(rect.get_x() + rect.get_width()/2., 1.05*height,
                '%d' % int(height),
                ha='center', va='bottom')

autolabel(rects1)
autolabel(rects2)

plt.show()

ext: py

max_stars_repo_name: Sunhick/music-cognita

max_issues_repo_name: Sunhick/music-cognita

max_forks_repo_name: Sunhick/music-cognita

reason: import numpy


text:
"""Provide pages for Command line interface App."""

import functools

import numpy
import six
import urwid
from urwid.compat import with_metaclass

from ..core import PKG_LOGGER, get_client, validate_connection
from ..util import wrap_connection_error
from .widgets import (FormFieldHorizontal, FormFieldHorizontalEdit,
                      FormFieldHorizontalPass, TicketCell, TicketColumn)

ZENDESK_LOGO = \
    " `////////////////.  :///////////////.  \n" + \
    "  mNNNNNNNNNNNNNNN-  dNNNNNNNNNNNNNm/   \n" + \
    "  -mNNNNNNNNNNNNN+   dNNNNNNNNNNNNs`    \n" + \
    "   .smNNNNNNNNNh-    dNNNNNNNNNNh.      \n" + \
    "     `:+syyso/`  ::  dNNNNNNNNm/        \n" + \
    "               .yN+  dNNNNNNNo`         \n" + \
    "              +NNN+  dNNNNNh.           \n" + \
    "            :dNNNN+  dNNNm:             \n" + \
    "          .yNNNNNN+  dNNo`              \n" + \
    "        `oNNNNNNNN+  dy.     `          \n" + \
    "       :mNNNNNNNNN+  -  -ohmNNNmho-     \n" + \
    "     .hNNNNNNNNNNN+   -hNNNNNNNNNNNh-   \n" + \
    "   `oNNNNNNNNNNNNN+  -NNNNNNNNNNNNNNN-  \n" + \
    "  /mNNNNNNNNNNNNNN+  yNNNNNNNNNNNNNNNy\ "

class AppElementMixin(with_metaclass(urwid.MetaSuper)):
    """
    Functionality common to app elements.

    - Run `refresh_widgets` whenever `render` or `keypress` is called.
    - Handle events.
    """

    # A mapping of keys to actions (override this).
    key_actions = {}

    def refresh_widgets(self, size):
        pass

    def _mix_render(self, size, focus=False):
        """Wrap super `render` to refresh widgets."""
        PKG_LOGGER.debug('{} rendering, size={} focus={}'.format(
            self.__class__.__name__, size, focus
        ))
        self.refresh_widgets(size)

    def _mix_keypress(self, size, key):
        """Wrap super `keypress` to refresh widgets."""
        PKG_LOGGER.debug('{} keypress, size={} key={}'.format(
            self.__class__.__name__, size, repr(key)
        ))

        # TODO: replace logic with urwid.command_map ?

        if key in self.key_actions:
            getattr(self, '_action_{}'.format(self.key_actions[key]))(
                size, key
            )

        self.refresh_widgets(size)

    def _action_exit(self, *_):
        raise urwid.ExitMainLoop()

    def _get_markup(self, ticket_dict, key, formatter=None):
        formatter = formatter or id

        unformatted = ticket_dict.get(key, '')
        try:
            return formatter(unformatted)
        except UnicodeEncodeError:
            if not isinstance(unformatted, six.text_type):
                unformatted = six.text_type(unformatted)
            unformatted = (unformatted).encode('ascii', errors='ignore')
            return formatter(unformatted)

    def modal_fatal_error(self, message=None, exc=None):
        """Cause a fatal error to be displayed and the program to exit."""
        message = "Error: {}".format(message) if message else "Fatal Error"

        # This could be called from a parent frame or a page.
        parent_frame = getattr(self, 'parent_frame', self)
        if message:
            PKG_LOGGER.critical(message)
            parent_frame.pages['ERROR'].page_title = message

        details = []
        if exc:
            PKG_LOGGER.critical(exc)
            details.insert(0, str(exc))
        details.append("press ctrl-c to exit")
        parent_frame.pages['ERROR'].error_details = "\n\n".join(details)
        parent_frame.set_page('ERROR')

class AppPageMixin(AppElementMixin):
    """Provide the interface for a page within an app."""

    _usage = ""
    _title = ""

    def __init__(self):
        """Wrap super __init__ as per `urwid.MetaClass` spec."""
        assert self.parent_frame
        self.__super.__init__()

    @property
    def page_usage(self):
        """Provide the usage message for this page."""
        return self._usage

    @property
    def page_title(self):
        """Provide the title for this page."""
        return self._title

    @property
    def page_status(self):
        """Provide an optional status line for this page."""
        return ""

class BlankPage(urwid.ListBox, AppPageMixin):
    """A blank app page."""

    def __init__(self, parent_frame, *args, **kwargs):
        """Wrap super `__init__` with extra metadata."""
        self.parent_frame = parent_frame
        self.__super.__init__(urwid.SimpleListWalker([]))

def mix_render_keypress(cls):
    """
    Shorthand for adding these methods into sublasses multiple times.

    Since each subclass of `AppElementMixin` calls the mixin keypress and
    render, this saves lots of redundant code
    """
    def keypress(self, size, key):
        """Wrap super `keypress` and perform actions."""
        # Scroll regardless of if a move was made
        self._mix_keypress(size, key)
        super(cls, self).keypress(size, key)

    def render(self, size, focus=False):
        """Wrap super and mixin `render`s."""
        self._mix_render(size, focus)
        return super(cls, self).render(size, focus)

    cls.keypress = keypress
    cls.render = render
    return cls

@mix_render_keypress
class TicketListPage(urwid.Columns, AppPageMixin):
    """An app page which displays a table of ticket information."""

    _selectable = True
    header_size = 1
    footer_size = 0
    # how quickly the scrolls when paging
    page_speed = 1
    _usage = (
        u"UP / DOWN / PAGE UP / PAGE DOWN scrolls. "
        u"SPACE / ENTER selects. "
        u"F8 exits."
    )
    _title = "Ticket List"

    key_actions = {
        ' ': 'open',
        'enter': 'open',
        'up': 'scroll',
        'down': 'scroll',
        'page up': 'scroll',
        'page down': 'scroll',
    }

    def __init__(self, parent_frame, *args, **kwargs):
        """Wrap super `__init__` with extra metadata."""
        # Cache access to generator to avoid api calls
        self._ticket_cache = []
        # Offset into the generator of the first visible element
        self.offset = 0
        # Index of the highlighted element
        self.index_highlighted = 0
        # Force a space of 1 between columns
        kwargs['dividechars'] = 0
        self.parent_frame = parent_frame
        self._ticket_generator = None
        self.__super.__init__(
            self.initial_column_widgets(), *args, **kwargs
        )

    @property
    def ticket_generator(self):
        """Try and get generator of tickets from API otherwise use cache."""
        if self._ticket_generator is not None:
            return self._ticket_generator

        client = self.parent_frame.client
        cache = client.tickets.cache.mapping['ticket'].cache
        # if we are in offline mode

        if cache.__class__.__name__ != 'TTLCache':
            self._ticket_generator = cache.values().__iter__()
        else:
            self._ticket_generator = wrap_connection_error(
                functools.partial(client.tickets, timeout=5),
                attempting="Connecting to API",
                on_fail=functools.partial(self.modal_fatal_error),
                on_success=functools.partial(
                    PKG_LOGGER.info, "Connected to API"
                )
            )
        return self._ticket_generator

    @property
    def next_ticket(self):
        # get before wrap so that we don't wrap twice
        generator = self.ticket_generator
        if generator is None:
            # want to exit getting ticket_generator cleanly in case there
            # was an error that needs to be displayed
            raise StopIteration()
        response = wrap_connection_error(
            injected=functools.partial(next, generator),
            attempting="Making an API request",
            on_fail=functools.partial(self.modal_fatal_error),
        )
        if response is None:
            raise StopIteration()
        return response

    @property
    def nonbody_overhead(self):
        """Rows taken up by the header and footer."""
        return self.header_size + self.footer_size

    @AppPageMixin.page_status.getter
    def page_status(self):
        """Provide the status message for this page."""
        # TODO: paging progress ("X - Y of Z")
        return ""

    def initial_column_widgets(self):
        """
        Generate the initial list of column widgets.

        Widget size is not known until render time, so no ticket entries are
        added to the widget list initially.
        """
        # First column is a selection indicator
        widget_list = [
            ('fixed', 1, TicketColumn(
                header=urwid.Divider(),
                body=urwid.Divider(),
                key='_selected'
            ))
        ]
        # Other widget columns show ticket data
        for key, meta in self.parent_frame.column_meta.items():
            if 'list_view' in meta and not meta['list_view']:
                continue
            title = meta.get('title', key.title())
            column_widget = TicketColumn(
                header=TicketCell(title),
                body=urwid.ListBox(urwid.SimpleListWalker([])),
                key=key
            )
            if 'sizing' in meta:
                column_widget = tuple(meta['sizing'] + [column_widget])
            widget_list.append(column_widget)

        return widget_list

    def get_tickets(self, offset, limit):
        """
        Fetch `limit` tickets from the generator starting at `offset`.

        Prefetch and cache 2 * `limit` tickets from the API.

        Args
        ----
            offset (:obj:`int`): the index of the first ticket required
            limit (:obj:`int`): the number of tickets required
        """
        prefetch_index = 2 * (limit) + offset

        try:
            while self.ticket_generator is not None:
                if prefetch_index < len(self._ticket_cache):
                    break
                self._ticket_cache.append(self.next_ticket)
        except StopIteration:
            pass
        return self._ticket_cache[offset:offset + limit]

    def _get_cell_widgets(self, key, visible_tickets, index_highlighted):
        meta = self.parent_frame.column_meta.get(key, {})
        formatter = meta.get('formatter', str)
        cell_kwargs = {
            'align': meta.get('align', urwid.LEFT)
        }
        cell_widgets = []
        for index, ticket in enumerate(visible_tickets):
            cell_kwargs['markup'] = self._get_markup(
                ticket.to_dict(), key, formatter
            )
            if key == '_selected' and index == index_highlighted:
                cell_kwargs['markup'] = '>'
            cell_widget = TicketCell(**cell_kwargs)
            if index == index_highlighted:
                cell_widget = urwid.AttrWrap(cell_widget, 'important')
            cell_widgets.append(cell_widget)
        return cell_widgets

    def refresh_widgets(self, size):
        """
        Populate frame body of each column with visible tickets.

        Args
        ----
            size (:obj:`tuple` of :obj:`int`): The allowed size of the widget

        """
        PKG_LOGGER.debug('refreshing, size={}'.format(size))
        self._action_scroll(size)

        _, maxcol = size
        visible_tickets = self.get_tickets(
            self.offset, maxcol - self.nonbody_overhead
        )
        index_highlighted = numpy.clip(
            self.index_highlighted,
            0,
            min(maxcol - self.nonbody_overhead, len(visible_tickets)) - 1
        )

        for column, _ in self.contents:
            cell_widgets = self._get_cell_widgets(
                column.key, visible_tickets, index_highlighted
            )

            # TODO: test for memory leaks

            column.body = urwid.ListBox(urwid.SimpleListWalker(cell_widgets))

    def _action_scroll(self, size, key=None):
        """
        Move highlighted index by `movement`, scroll `offset` at boundaries.

        Even if movement is 0 it is useful to refresh these values since the
        widget can be resized.
        """
        PKG_LOGGER.debug('scrolling, size={} key={}'.format(
            size, key
        ))
        _, maxcol = size
        # Map key value to scroll movement amount
        page_jump = int(self.page_speed * (maxcol - self.nonbody_overhead))
        key_movements = {
            'up': -1,
            'down': 1,
            'page up': -page_jump,
            'page down': page_jump
        }
        movement = key_movements.get(key, 0)
        # move highlighted index until boundaries
        can_move_to = numpy.clip(
            self.index_highlighted + movement,
            0,
            maxcol - self.nonbody_overhead - 1
        )
        # determine remaining movement to potentially move the offset
        movement = movement - (can_move_to - self.index_highlighted)
        self.index_highlighted = can_move_to

        # offset can't exceed ticket cache
        self.offset = numpy.clip(
            self.offset + movement,
            0,
            max(len(self._ticket_cache) - 1, 0)
        )

    def _action_open(self, *_):
        """Open view of selected ticket."""
        ticket = self._ticket_cache[self.offset + self.index_highlighted]
        PKG_LOGGER.debug('Actioning ticket id={}'.format(ticket))
        self.parent_frame.pages['TICKET_VIEW'].current_ticket = ticket
        self.parent_frame.set_page('TICKET_VIEW')

@mix_render_keypress
class TicketViewPage(urwid.ListBox, AppPageMixin):
    """An app page which displays a single ticket's information."""

    _usage = (
        u"ESC is back. "
        u"F8 exits."
    )
    _title = "Ticket View"

    def __init__(self, parent_frame, *args, **kwargs):
        """Wrap super `__init__` with extra metadata."""
        self.parent_frame = parent_frame
        self.current_ticket = None
        self.__super.__init__(urwid.SimpleListWalker(
            self.initial_row_widgets()
        ))

    def initial_row_widgets(self):
        """Initialize the row widgets to be updated later."""
        widget_list = []

        for key, meta in self.parent_frame.column_meta.items():
            field_label = meta.get('title', key.title())
            field_class = meta.get('field_class', FormFieldHorizontal)
            widget_list.append(field_class(field_label, key=key))

        return widget_list

    def refresh_widgets(self, size):
        """Refresh the row widgets."""
        ticket_dict = {}
        if self.current_ticket:
            ticket_dict = self.current_ticket.to_dict()

        for wg_field in self.body.contents:
            meta = self.parent_frame.column_meta.get(wg_field.key, {})
            _, (wg_field_value, _) = wg_field.contents
            formatter = meta.get('formatter', str)
            markup = self._get_markup(ticket_dict, wg_field.key, formatter)
            if wg_field_value.text != markup:
                wg_field_value.set_text(markup)

@mix_render_keypress
class ErrorPage(urwid.Overlay, AppPageMixin):
    """An app page which displays an error."""
    _usage = (
        u"F8 / ESC exits."
    )

    key_actions = {
        'f8': 'exit',
        'esc': 'exit',
        'ctrl c': 'exit',
    }

    def __init__(self, parent_frame, *args, **kwargs):
        """Wrap super `__init__` with extra metadata."""
        self.parent_frame = parent_frame
        self._title = kwargs.get('error_message', 'Error')
        self.error_details = kwargs.get('error_details', 'An Error occured.')

        self.__super.__init__(
            urwid.AttrWrap(
                urwid.LineBox(
                    urwid.ListBox(urwid.SimpleFocusListWalker([
                        urwid.Divider(),
                        urwid.Text(self.error_details, align='center')
                    ])),
                    title=self.page_title
                ),
                'header'
            ),
            self.parent_frame.current_page,
            align='center', width=('relative', 50),
            valign='middle', height=('relative', 50),
            min_width=24, min_height=8,
        )

    @AppPageMixin.page_title.setter
    def page_title(self, value):
        self._title = value

    def refresh_widgets(self, *_):
        _, (wg_top, _) = self.contents
        wg_title = wg_top.original_widget.title_widget
        if wg_title.text != self.page_title:
            wg_title.set_text(self.page_title)
        wg_details = wg_top.original_widget.original_widget.body[1]
        if wg_details.text != self.error_details:
            wg_details.set_text(self.error_details)

@mix_render_keypress
class WelcomePage(urwid.Overlay, AppPageMixin):
    _title = "Welcome"

    key_actions = {
        'enter': 'login'
    }

    _usage = (
        u"F8 / ESC exits."
    )

    def __init__(self, parent_frame, *args, **kwargs):
        """Wrap super `__init__` with extra metadata."""
        self.parent_frame = parent_frame
        config = self.parent_frame.config
        self.form_fields = [
            cls(label_text, getattr(config, key, ''), key=key)
            for key, (cls, label_text) in {
                'subdomain': (FormFieldHorizontalEdit, 'Subdomain: '),
                'email': (FormFieldHorizontalEdit, 'Email: '),
                'password': (FormFieldHorizontalPass, 'Password: '),
            }.items()
        ]

        widget_list = [
            urwid.Divider(),
            urwid.Text(ZENDESK_LOGO, align='center'),
            urwid.Divider(),
        ] + self.form_fields + [
            urwid.Divider(),
            urwid.Button(
                "Login", on_press=self._action_login
            )
        ]

        self.__super.__init__(
            urwid.AttrWrap(
                urwid.LineBox(
                    urwid.ListBox(urwid.SimpleFocusListWalker(
                        widget_list
                    )),
                ),
                'column_header'
            ),
            self.parent_frame.current_page,
            align='center', width=('relative', 50),
            valign='middle', height=('relative', 80),
            min_width=24, min_height=8,
        )

    def _action_login(self, *args):
        for wg_field in self.form_fields:
            value = wg_field.get_value_text()
            setattr(
                self.parent_frame.config, wg_field.key, value
            )
            if wg_field.key == 'password':
                value = "*" * len(value)
            PKG_LOGGER.info("updated config[{}] = {}".format(
                wg_field.key, value
            ))

        # The Ticket Viewer should handle the API being unavailable
        wrap_connection_error(
            functools.partial(validate_connection, self.parent_frame.config),
            attempting="Validate connection",
            on_fail=functools.partial(self.modal_fatal_error),
            on_success=functools.partial(
                PKG_LOGGER.info, "Connection validated"
            )
        )

        self.parent_frame.client = wrap_connection_error(
            functools.partial(get_client, self.parent_frame.config),
            attempting="Create client",
            on_fail=functools.partial(self.modal_fatal_error),
            on_success=functools.partial(
                PKG_LOGGER.info, "Client created"
            )
        )

        # if no error screen is showing
        if self.parent_frame.current_page_id == 'WELCOME':
            self.parent_frame.set_page('TICKET_LIST')

max_stars_repo_name: derwentx/Zendesk-Ticket-Viewer

max_issues_repo_name: derwentx/Zendesk-Ticket-Viewer

max_forks_repo_name: derwentx/Zendesk-Ticket-Viewer


text:
import unittest

import numpy as np
from shapely import wkt as wktreader
from csv import DictReader

from deep_geometry.vectorizer \
    import num_points_from_wkt, vectorize_wkt, get_max_points, \
    GEO_VECTOR_LEN, IS_INNER_INDEX, IS_OUTER_INDEX, FULL_STOP_INDEX

TOPOLOGY_CSV = 'test_files/polygon_multipolygon.csv'

SOURCE_DATA = []
with open(TOPOLOGY_CSV, 'r') as csv_file:
    reader = DictReader(csv_file, )
    for row in reader:
        SOURCE_DATA.append(row)

brt_wkt = [record['brt_wkt'] for record in SOURCE_DATA]
osm_wkt = [record['osm_wkt'] for record in SOURCE_DATA]
target_wkt = [record['intersection_wkt'] for record in SOURCE_DATA]

input_geom = np.array([
    [0., 0., 1., 0., 0.],
    [0., 1., 1., 0., 0.],
    [1., 1., 1., 0., 0.],
    [1., 0., 1., 0., 0.],
    [0., 0., 0., 1., 0.],
    [0., 0., 1., 0., 0.],
    [0., -1., 1., 0., 0.],
    [-1., -1., 1., 0., 0.],
    [-1., 0., 1., 0., 0.],
    [0., 0., 0., 0., 1.],
    [0., 0., 0., 0., 0.]
])

output_geom = np.array([
    [0.0, 0.00, 1., 0., 0.],
    [0.0, 0.25, 1., 0., 0.],
    [0.0, 0.50, 1., 0., 0.],
    [0.0, 0.75, 1., 0., 0.],
    [0.0, 1.00, 1., 0., 0.],
    [0.25, 1.0, 1., 0., 0.],
    [0.50, 1.0, 1., 0., 0.],
    [1.0, 1.00, 1., 0., 0.],
    [1.0, 0.50, 1., 0., 0.],
    [1.0, 0.00, 1., 0., 0.],
    [0.5, 0.00, 1., 0., 0.],
    [0.0, 0.00, 0., 1., 0.],
    [0.0, 0.00, 1., 0., 0.],
    [0.0, -0.5, 1., 0., 0.],
    [0.0, -1.0, 1., 0., 0.],
    [-0.5, -1., 1., 0., 0.],
    [-1., -1.0, 1., 0., 0.],
    [-1., -0.5, 1., 0., 0.],
    [-1., 0.00, 1., 0., 0.],
    [-0.5, 0.0, 1., 0., 0.],
    [0.00, 0.0, 0., 0., 1.],
    [0.00, 0.0, 0., 0., 0.]
])

non_empty_geom_collection = 'GEOMETRYCOLLECTION(LINESTRING(1 1, 3 5),POLYGON((-1 -1, -1 -5, -5 -5, -5 -1, -1 -1)))'

class TestVectorizer(unittest.TestCase):
    def test_num_points_from_multipolygon_with_hole(self) -> None:
        with open('test_files/multipolygon_with_hole.txt', 'r') as file:
            wkt = file.read()
            total_points = num_points_from_wkt(wkt)
            self.assertEqual(total_points, 683)

    def test_max_points(self) -> None:
        max_pts = get_max_points(brt_wkt, osm_wkt)
        self.assertEqual(max_pts, 159)

    def test_max_points_3d(self) -> None:
        geom_3d = 'POLYGON((0 0 0, 1 1 1, 2 2 2, 0 0 0))'
        max_pts = get_max_points([geom_3d])
        self.assertEqual(max_pts, 4)

    def test_vectorize_one_wkt(self) -> None:
        max_points = 20
        input_set = target_wkt
        vectorized = []
        for index in range(len(input_set)):
            vectorized.append(vectorize_wkt(input_set[index], max_points, simplify=True))
        self.assertEqual(len(input_set), len(brt_wkt))
        self.assertEqual(vectorized[0].shape, (19, GEO_VECTOR_LEN))
        self.assertEqual(vectorized[1].shape, (1, GEO_VECTOR_LEN))

    def test_no_max_points_fixed_size(self) -> None:
        input_set = np.array(target_wkt)
        with self.assertRaises(AssertionError):
            vectorized = [vectorize_wkt(wkt, fixed_size=True) for wkt in input_set]

    def test_fixed_size(self) -> None:
        max_points = 20
        input_set = np.array(target_wkt)
        vectorized = [vectorize_wkt(wkt, max_points, simplify=True, fixed_size=True) for wkt in input_set]
        self.assertEqual(np.array(vectorized).shape, (input_set.size, 20, GEO_VECTOR_LEN))
        for vector in vectorized:
            self.assertEqual(vector[-1, FULL_STOP_INDEX], 1)

    def test_non_empty_geom_coll(self) -> None:
        with self.assertRaises(ValueError):
            vectorize_wkt(non_empty_geom_collection, 100)

    def test_point_without_max_points(self) -> None:
        vectorized = vectorize_wkt('POINT(12 14)')
        self.assertEqual(vectorized.shape, (1, GEO_VECTOR_LEN))
        self.assertEqual(vectorized[-1, FULL_STOP_INDEX], 1)  # Test full stop bit

    def test_point_with_max_points(self) -> None:
        vectorized = vectorize_wkt('POINT(12 14)', 5)
        self.assertEqual(vectorized.shape, (1, GEO_VECTOR_LEN))
        self.assertEqual(vectorized[-1, FULL_STOP_INDEX], 1)  # Test full stop bit

    def test_unsupported_geom(self) -> None:
        with self.assertRaises(Exception):
            vectorize_wkt(
                'THIS_SHOULD_THROW_AN_EXCEPTION ((10 10, 20 20, 10 40),(40 40, 30 30, 40 20, 30 10))', 16)

    def test_multipolygon(self) -> None:
        with open('test_files/multipolygon.txt', 'r') as file:
            wkt = file.read()
            vectorized = vectorize_wkt(wkt)
            self.assertEqual((333, GEO_VECTOR_LEN), vectorized.shape)
            self.assertEqual(vectorized[-1, FULL_STOP_INDEX], 1)  # Test full stop bit

    def test_polygon_with_hole(self) -> None:
        polygon_with_hole = "POLYGON((0 0, 3 0, 3 3, 0 3, 0 0), (1 1, 2 1, 2 2, 1 2, 1 1))"
        vectorized = vectorize_wkt(polygon_with_hole)
        for is_inner_bit in vectorized[:5, IS_INNER_INDEX]:
            self.assertEqual(is_inner_bit, 1)
        for is_outer_bit in vectorized[5:, IS_OUTER_INDEX]:
            self.assertEqual(is_outer_bit, 1)
        self.assertEqual(vectorized[-1, FULL_STOP_INDEX], 1)

    def test_multipolygon_with_hole(self) -> None:
        with open('test_files/multipolygon_with_hole.txt', 'r') as file:
            wkt = file.read()
            vectorized = vectorize_wkt(wkt)
            self.assertEqual((683, GEO_VECTOR_LEN), vectorized.shape)
            self.assertEqual(vectorized[-1, FULL_STOP_INDEX], 1)  # Test full stop bit

    def test_vectorize_big_multipolygon(self) -> None:
        with open('test_files/big_multipolygon_wkt.txt', 'r') as file:
            wkt = file.read()
            max_pts = get_max_points([wkt])
            vectorized = vectorize_wkt(wkt, max_pts)
            self.assertEqual((144, GEO_VECTOR_LEN), vectorized.shape)
            self.assertEqual(vectorized[-1, FULL_STOP_INDEX], 1)  # Test full stop bit

    def test_simplify_multipolygon_gt_max_points(self) -> None:
        with open('test_files/multipart_multipolygon_wkt.txt', 'r') as file:
            wkt = file.read()
            max_points = 20
            vectorized = vectorize_wkt(wkt, max_points, simplify=True)
            self.assertEqual((20, GEO_VECTOR_LEN), vectorized.shape)
            self.assertEqual(vectorized[-1, FULL_STOP_INDEX], 1)  # Test full stop bit

    def test_simplify_without_max_points(self) -> None:
        with open('test_files/multipart_multipolygon_wkt.txt', 'r') as file:
            wkt = file.read()
            with self.assertRaises(AssertionError):
                vectorize_wkt(wkt, simplify=True)

    def test_multipolygon_exceed_max_points(self) -> None:
        with open('test_files/multipart_multipolygon_wkt.txt', 'r') as file:
            wkt = file.read()
            max_points = 20
            with self.assertRaises(Exception):
                vectorize_wkt(wkt, max_points)

    def test_polygon_exceed_max_points(self) -> None:
        with open('test_files/multipart_multipolygon_wkt.txt', 'r') as file:
            wkt = file.read()
            shape = wktreader.loads(wkt)
            geom = shape.geoms[0]
            max_points = 20
            with self.assertRaises(Exception):
                vectorize_wkt(geom.wkt, max_points)

max_stars_repo_name: SPINLab/deep-geometry

max_issues_repo_name: SPINLab/deep-geometry

max_forks_repo_name: SPINLab/deep-geometry