#!/usr/bin/env python

"""Tools for calculating the complex voltage response of MWA tiles and dipoles
Randall Wayth. March 2014, based on the work of Adrian Sutinjo.
"""

import copy
import logging

import numpy

import astropy.io.fits as pyfits

from scipy import interpolate

from . import config
from . import mwa_impedance

vel_light = 2.99792e8
DQ = 435e-12 * vel_light  # delay quantum in distance light travels for 1 quantum
logging.basicConfig(format='# %(levelname)s:%(name)s: %(message)s')
logger = logging.getLogger(__name__)  # default logger level is WARNING

# locations of dipoles in tile
DIPOLE_SEP = config.DIPOLE_SEPARATION  # meters
X_OFFSETS = numpy.array([-1.5, -0.5, 0.5, 1.5,
                         -1.5, -0.5, 0.5, 1.5,
                         -1.5, -0.5, 0.5, 1.5,
                         -1.5, -0.5, 0.5, 1.5],
                        dtype=numpy.float32) * DIPOLE_SEP
Y_OFFSETS = numpy.array([1.5, 1.5, 1.5, 1.5,
                         0.5, 0.5, 0.5, 0.5,
                         -0.5, -0.5, -0.5, -0.5,
                         -1.5, -1.5, -1.5, -1.5],
                        dtype=numpy.float32) * DIPOLE_SEP

# DIPOLE_DEFAULT = None            # Will contain a default Dipole('lookup') object
# APERTURE_ARRAY_DEFAULT = None    # Will contain a default ApertureArray object

class Dipole(object):
    """Provides a generic dual pol dipole object to support MWA beam models"""

    def __init__(self,
                 atype='short',
                 height=0.29,
                 length=0.74,
                 lookup_filename=config.Jmatrix,
                 gain=None):
        """
          General dipole object. Dual pol crossed dipole.
          Assumes a groundscreen with dipole height meters above ground.
          Supported types are 'short': analytic short dipole
                              'long':  analytic real length dipole
                              'lookup': lookup table
          For a short dipole, the length of the dipole is ignored.
          Gain is a 2x2 matrix with gain and direction independent crosstalk
          for the X (upper) and Y (lower) dipole resectively
        """

        assert (atype == 'short' or atype == 'long' or atype == 'lookup'), 'Unknown atype %r' % atype
        self.atype = atype
        self.height = height
        self.length = length
        if gain is None:
            self.gain = numpy.eye(2, dtype=numpy.complex64)
        else:
            self.gain = gain
        self.interp_freq = 0.0
        self.lookup = None
        self.lookup_za = None
        self.lookup_ph = None
        self.i00_real, self.i01_real, self.i10_real, self.i11_real = (None, None, None, None)
        self.i00_imag, self.i01_imag, self.i10_imag, self.i11_imag = (None, None, None, None)
        self.j00norm, self.j01norm, self.j10norm, self.j11norm = (None, None, None, None)
        self.freqs = numpy.array([])
        if atype == 'lookup':
            self.loadLookup(lookup_filename)

    def loadLookup(self, lookup_filename):
        """Load a dipole Jones response lookup table (FITS file)"""
        # data is a direct conversion of the output of the simulation so has redundant info.
        # we do all the conversion and stuff here
        # for reference, columns are:
        # theta phi  real(Jxt(t,p)) imag(Jxt(t,p)) real(Jxp(t,p)) imag(Jxp(t,p)) real(Jyt(t,p)) imag(Jyt(t,p)) real(Jyp(t,p)) imag(Jyp(t,p)))
        try:
            logger.debug('Loading %s' % lookup_filename)
            hdulist = pyfits.open(lookup_filename)
        except IOError:
            raise Exception('Cannot load Jones matrix file %s' % lookup_filename)
        nfreqs = len(hdulist)
        self.lookup_za = numpy.unique(hdulist[0].data[:, 0])  # zenith angle
        self.lookup_ph = numpy.unique(hdulist[0].data[:, 1])  # phi angle == 90 - az
        # p = numpy.where
        nza = len(self.lookup_za)
        nph = len(self.lookup_ph)
        self.lookup = numpy.empty((nfreqs, nza, nph, 2, 2), dtype=numpy.complex64)
        freqs = []
        for i in range(nfreqs):
            hdu = hdulist[i]
            logger.debug('Loading J lookup matrix for freq ' + str(hdu.header['FREQ']))
            freqs.append(hdu.header['FREQ'])
            jxt = hdu.data[:, 2] + 1.0j * hdu.data[:, 3]
            jxp = hdu.data[:, 4] + 1.0j * hdu.data[:, 5]
            jyt = hdu.data[:, 6] + 1.0j * hdu.data[:, 7]
            jyp = hdu.data[:, 8] + 1.0j * hdu.data[:, 9]
            self.lookup[i, :, :, 0, 0] = jxt.reshape((nph, nza)).transpose()
            self.lookup[i, :, :, 0, 1] = jxp.reshape((nph, nza)).transpose()
            self.lookup[i, :, :, 1, 0] = jyt.reshape((nph, nza)).transpose()
            self.lookup[i, :, :, 1, 1] = jyp.reshape((nph, nza)).transpose()
        logger.debug('Loaded dipole Jones matrix lookup model from ' + lookup_filename + ' with ' + str(nfreqs) + ' freqs')
        self.freqs = numpy.array(freqs)
        logger.debug('Supported frequencies (MHz): ' + str(self.freqs / 1e6))
        logger.debug("There are " + str(nza) + " tabulated zenith angles: " + str(self.lookup_za))
        logger.debug("There are " + str(nph) + " tabulated phi angles: " + str(self.lookup_ph))

    def getJones(self, az, za, freq, zenith_norm=True):
        """Return the Jones matrix for arrays of az/za for a given freq
        az and za are numpy arrays of equal length.
        Results are in corrds of az/za unit vectors such that
        output[j,i,:,:] is a 2x2 Jones matrix that maps
        za onto E-W     az onto E-W
        za onto N-S     az onto N-S
        By default, normalise to the zenith
        """
        if self.atype == 'short':
            return self.getJonesShortDipole(az, za, freq, zenith_norm=zenith_norm)
        elif self.atype == 'lookup':
            return self.getJonesLookup(az, za, freq)
        else:
            raise Exception("Dipole atype %r is not implemented yet." % self.atype)

    def getJonesLookup(self, az, za, freq):
        """Return the Jones matrix for arrays of az/za for a given freq (Hz)
        this method interpolates from the tablulated numerical results loaded
        by the constructor"""
        # need to interpolate each of the 4 Jones elements separately and each
        # the real and imag separately (since interpolate.RectBivariateSpline)
        # apparently doesn't handle complex

        # find the nearest freq lookup table
        pos = numpy.argmin(numpy.abs(self.freqs - freq))
        if numpy.abs(self.freqs[pos] - freq) > 2e6:
            logger.warning("Nearest tabulated impedance matrix freq is more than 2 MHz away from desired freq.")
        logger.info("Selecting matrix for nearest freq " + str(self.freqs[pos]))

        # cache the interpolation functions
        if self.interp_freq != freq:
            self.interp_freq = freq
            logger.debug("Setting new cache lookup freq to " + str(self.freqs[pos]))
            self.i00_real = interpolate.RectBivariateSpline(self.lookup_za, self.lookup_ph,
                                                            self.lookup[pos, :, :, 0, 0].real)
            self.i00_imag = interpolate.RectBivariateSpline(self.lookup_za, self.lookup_ph,
                                                            self.lookup[pos, :, :, 0, 0].imag)
            self.i01_real = interpolate.RectBivariateSpline(self.lookup_za, self.lookup_ph,
                                                            self.lookup[pos, :, :, 0, 1].real)
            self.i01_imag = interpolate.RectBivariateSpline(self.lookup_za, self.lookup_ph,
                                                            self.lookup[pos, :, :, 0, 1].imag)
            self.i10_real = interpolate.RectBivariateSpline(self.lookup_za, self.lookup_ph,
                                                            self.lookup[pos, :, :, 1, 0].real)
            self.i10_imag = interpolate.RectBivariateSpline(self.lookup_za, self.lookup_ph,
                                                            self.lookup[pos, :, :, 1, 0].imag)
            self.i11_real = interpolate.RectBivariateSpline(self.lookup_za, self.lookup_ph,
                                                            self.lookup[pos, :, :, 1, 1].real)
            self.i11_imag = interpolate.RectBivariateSpline(self.lookup_za, self.lookup_ph,
                                                            self.lookup[pos, :, :, 1, 1].imag)
            # determine normalisation factors. The simulations include all ph angles at za=0.
            # these are not redundant, and the ph value determines the unit vector directions of
            # both axes. We should normalise by where the result will be maximal.
            # For the E-W dipoles, the projection of the ZA unit vec will be max when
            # pointing east, i.e. when ph=0. For the PH unit vec, this will be when ph=-90 or 90
            # For the N-S dipoles, projection of ZA onto N-S is max az ph=90 and
            # proj of ph onto N-S is max when ph=0

            # determine the indices of 0 and 90 degrees of phi in the tabulated values
            ph0 = numpy.where(self.lookup_ph == 0.0)
            ph90 = numpy.where(self.lookup_ph == 90.0)
            th0 = numpy.where(self.lookup_za == 0.0)

            self.j00norm = self.lookup[pos, th0, ph0, 0, 0]
            self.j01norm = -self.lookup[pos, th0, ph90, 0, 1]  # use -90, not 90.
            self.j10norm = self.lookup[pos, th0, ph90, 1, 0]
            self.j11norm = self.lookup[pos, th0, ph0, 1, 1]

        ph_deg = 90.0 - az * 180.0 / numpy.pi  # ph in degrees
        za_deg = za * 180.0 / numpy.pi
        p = ph_deg < 0
        ph_deg[p] += 360.0
        j00 = self.i00_real.ev(za_deg.flatten(), ph_deg.flatten()) + 1.0j * self.i00_imag.ev(za_deg.flatten(),
                                                                                ph_deg.flatten())
        j01 = self.i01_real.ev(za_deg.flatten(), ph_deg.flatten()) + 1.0j * self.i01_imag.ev(za_deg.flatten(),
                                                                                ph_deg.flatten())
        j10 = self.i10_real.ev(za_deg.flatten(), ph_deg.flatten()) + 1.0j * self.i10_imag.ev(za_deg.flatten(),
                                                                                ph_deg.flatten())
        j11 = self.i11_real.ev(za_deg.flatten(), ph_deg.flatten()) + 1.0j * self.i11_imag.ev(za_deg.flatten(),
                                                                                ph_deg.flatten())

        result = numpy.empty((za.shape + (2, 2)), dtype=numpy.complex64)
        result[..., 0, 0] = j00.reshape(za.shape) / self.j00norm
        result[..., 0, 1] = -j01.reshape(za.shape) / self.j01norm  # sign flip between az and phi
        result[..., 1, 0] = j10.reshape(za.shape) / self.j10norm
        result[..., 1, 1] = -j11.reshape(za.shape) / self.j11norm  # sign flip between az and phi
        return result

    def getJonesShortDipole(self, az, za, freq, zenith_norm=True):
        """Calculate the Jones matrix for a short dipole.
        This is defined by purely geometric projection of unit vectors
        on the sky onto the unit vector defined by the dipole's direction.
        """
        # check for input scalars vs arrays. If the input is scalar, then
        # convert it so that numpy array-based operations are possible
        if numpy.isscalar(az):
            logger.debug("Converting scalar az input to a matrix")
            az = numpy.asarray(az, dtype=numpy.float32)
        if numpy.isscalar(za):
            logger.debug("Converting scalar za input to a matrix")
            za = numpy.asarray(za, dtype=numpy.float32)
        assert az.shape == za.shape, "Input za and az arrays must have same dimension"

        # output array has 2x2 Jones matrix for every az/za point in input
        result = numpy.empty((za.shape + (2, 2)), dtype=numpy.complex64)
        # apply the groundscreen factor, which is independent of az
        # znorm = 1.0
        gs = self.groundScreen(za, freq)
        if zenith_norm:
            logger.debug("Normalising response to zenith")
            gs /= self.groundScreen(0.0, freq)
        result[..., 0, 0] = numpy.cos(za) * numpy.sin(az) * gs
        result[..., 0, 1] = numpy.cos(az) * gs
        result[..., 1, 0] = numpy.cos(za) * numpy.cos(az) * gs
        result[..., 1, 1] = -numpy.sin(az) * gs
        return result

    def groundScreen(self, za, freq):
        """
        Calculate the groundscreen effect for an ideal infinite groundscreen
        given the dipole's height above the screen and the frequency (Hz)
        """
        ll = vel_light / freq
        return numpy.sin(numpy.pi * (2.0 * self.height / ll) * numpy.cos(za)) * 2.0

    def __str__(self):
        return "Dipole. Type: " + self.atype + ". height: " + str(self.height) + "m. Gain: " + str(self.gain)

DIPOLE_DEFAULT = Dipole('lookup')      # A default Dipole('lookup') object. FOr thread safety, copy.deepcopy this global object

class ApertureArray(object):
    """Aperture array antenna object"""

    def __init__(self, dipoles=None, xpos=None, ypos=None):
        """Constructor for aperture array station. xpos and ypos are arrays with
        the coords of the dipoles (meters) in local coords relative to centre of
        the antenna looking down on the station. Ordering goes left to right, top to bottom,
        hence are offsets in east and north from the array phase centre.
        """
        assert dipoles is None or len(dipoles) == 16, "Expecting 16 input dipoles, got %r" % str(len(dipoles))
        if dipoles is None:
            d = Dipole(atype='short')
            dipoles = [d] * 16
        if xpos is None:
            xpos = X_OFFSETS
        if ypos is None:
            ypos = Y_OFFSETS
        self.dipoles = dipoles
        self.xpos = xpos
        self.ypos = ypos
        self.im = mwa_impedance.TileImpedanceMatrix()
        self.lna_z = mwa_impedance.LNAImpedance()

    def getPortCurrents(self, freq, delays=None):
        """
        Return the port currents on a tile given the freq (Hz) and delays (integer)
        """
        if delays is None:
            delays = numpy.zeros((2, 16), dtype=numpy.float32)
        lam = vel_light / freq
        phases = -2.0 * numpy.pi * delays * (DQ / lam)
        ph_rot = numpy.cos(phases) + 1.0j * numpy.sin(phases)
        # this code ignores any dipole gain (and crosstalk) terms.
        # should FIXME it.
        z_total = self.im.getImpedanceMatrix(freq) + numpy.eye(32) * self.lna_z.getZ(freq)
        inv_z = numpy.linalg.inv(z_total)
        port_current = numpy.dot(inv_z, ph_rot.reshape(32)).reshape(2, 16)
        return port_current

    def getArrayFactor(self, az, za, freq=155e6, delays=None):
        """
        Get the scalar array factor response of the array for a given
        freq (Hz) and delay settings.
        az and za (radian) are numpy arrays of equal length defining a set
        of points to calculate the response for.
        delays is a 2D array of integer delay steps for the Y and X pol
        respectively.
        Result are in same coords as the az/za input arrays
        """
        if delays is None:
            delays = numpy.zeros((2, 16), dtype=numpy.float32)
        lam = vel_light / freq
        port_current = self.getPortCurrents(freq, delays)

        # check for input scalars vs arrays. If the input is scalar, then
        # convert it so that numpy array-based operations are possible
        if numpy.isscalar(az):
            logger.debug("Converting scalar az input to a matrix")
            az = numpy.matrix(az, dtype=numpy.float32)
        if numpy.isscalar(za):
            logger.debug("Converting scalar za input to a matrix")
            za = numpy.matrix(za, dtype=numpy.float32)
        assert az.shape == za.shape, "Input az and za arrays must have same dimenions"

        # now calculate the array factor using these port currents
        sz = numpy.sin(za)
        kx = (2.0 * numpy.pi / lam) * numpy.sin(az) * sz
        ky = (2.0 * numpy.pi / lam) * numpy.cos(az) * sz
        ax = numpy.zeros_like(az, dtype=numpy.complex64)
        ay = numpy.zeros_like(az, dtype=numpy.complex64)
        for i in range(len(self.xpos)):
            ph = kx * self.xpos[i] + ky * self.ypos[i]
            ax += port_current[1, i] * (numpy.cos(ph) + 1.0j * numpy.sin(ph))  # X dipoles
            ay += port_current[0, i] * (numpy.cos(ph) + 1.0j * numpy.sin(ph))  # Y dipoles
        # set the points below the horizon to zero
        p = za >= numpy.pi / 2.0
        ax[p] = 0.0
        ay[p] = 0.0
        return (ax, ay)

    def getResponse(self, az, za, freq=155e6, delays=None):
        """
        Get the full Jones matrix response of the tile including the dipole
        reponse and array factor incorporating any mutual coupling effects
        from the impedance matrix. freq in Hz.
        delays in unit steps of beamformer delays as numpy array shape (2,16).
        az and za (radian) are numpy arrays of equal length defining a set
        of points to calculate the response for.
        Result is an array like az/za with [2][2] on the end for the Jones.
        """
        assert delays is None or numpy.size(delays) == 32, "Expecting 32 delays, got %r" % str(numpy.size(delays))
        if delays is None:
            delays = numpy.zeros((2, 16), dtype=numpy.float32)
        (ax, ay) = self.getArrayFactor(az, za, freq, delays)
        # get the zenith response to normalise to:
        (zax, zay) = self.getArrayFactor(numpy.array([0.0]), numpy.array([0.0]), freq)  # no delays == zenith
        ax /= numpy.abs(zax)
        ay /= numpy.abs(zay)
        d = self.dipoles[0]  # for now, assume all dipoles identical FIXME
        j = d.getJones(az, za, freq)
        j[..., 0, 0] *= ax
        j[..., 0, 1] *= ax
        j[..., 1, 0] *= ay
        j[..., 1, 1] *= ay
        return j

APERTURE_ARRAY_DEFAULT = ApertureArray(dipoles=[DIPOLE_DEFAULT] * 16)

def get_AA_Cached():
    return copy.deepcopy(APERTURE_ARRAY_DEFAULT)

def convertJonesAzEl2HaDec(az, za, lat):
    """
    Generate a converter for arrays of Jones matrices in Az/ZA to HA/DEC.
    This is a rotation of the orthogonal az/za unit vectors to
    the orthogonal ha/dec unit vectors by the parallactic angle
    plus some negative signs to account for the definition of
    coord axes. Multiply this by your az/za Jones to get a HA/DEC Jones.
    Inputs and outputs are numpy arrays of dimension [az][za][2][2]"""
    # parallactic angle needs HA, DEC and lat.
    (ha, dec) = h2e(az, za, lat)
    pa = calcParallacticAngle(ha, dec, lat)
    rot = numpy.empty((ha.shape + (2, 2)))
    crot = numpy.cos(pa * numpy.pi / 180.0)
    srot = numpy.sin(pa * numpy.pi / 180.0)
    # for clockwise rotations, use [[cos,sin],[-sin,cos]]
    rot[:, :, 0, 0] = crot
    rot[:, :, 0, 1] = srot
    rot[:, :, 1, 0] = -srot
    rot[:, :, 1, 1] = crot

def h2e(az, za, lat):
    """
    Horizon to equatorial.
    Convert az/za (radian) to HA/DEC (degrees, degrees)
    given an observatory latitude (degrees)
    """
    sa = numpy.sin(az)
    ca = numpy.cos(az)
    se = numpy.sin(numpy.pi / 2.0 - za)
    ce = numpy.cos(numpy.pi / 2.0 - za)
    sp = numpy.sin(lat * numpy.pi / 180.0)
    cp = numpy.cos(lat * numpy.pi / 180.0)

    # HA,Dec as x,y,z */
    x = - ca * ce * sp + se * cp
    y = - sa * ce
    z = ca * ce * cp + se * sp

    # To spherical */
    r = numpy.sqrt(x * x + y * y)
    ha = numpy.arctan2(y, x) * 180.0 / numpy.pi
    dec = numpy.arctan2(z, r) * 180.0 / numpy.pi
    return (ha, dec)

def e2h(ha, dec, lat):
    """
    Equatorial to horizon.
    Convert equatorial ha/dec coords (both in degs) to az,za (radian)
    given an observer latitute (degs). Returns (az,za)
    """
    ha_rad = ha * numpy.pi / 180.0
    dec_rad = dec * numpy.pi / 180.0
    lat_rad = lat * numpy.pi / 180.0
    sh = numpy.sin(ha_rad)
    ch = numpy.cos(ha_rad)
    sd = numpy.sin(dec_rad)
    cd = numpy.cos(dec_rad)
    sp = numpy.sin(lat_rad)
    cp = numpy.cos(lat_rad)
    x = - ch * cd * sp + sd * cp
    y = - sh * cd
    z = ch * cd * cp + sd * sp
    r = numpy.sqrt(x * x + y * y)
    a = numpy.arctan2(y, x)
    if a < 0.0:
        az = a + 2.0 * numpy.pi
    else:
        az = a
    el = numpy.arctan2(z, r)
    return (az, numpy.pi / 2.0 - el)

def calcParallacticAngle(ha, dec, lat):
    """
    Calculate the parallactic angle in degrees given an HA (degs)
    dec (degrees) and observatory latitude (degrees)
    """
    cl = numpy.cos(lat * numpy.pi / 180.0)
    sl = numpy.sin(lat * numpy.pi / 180.0)
    ch = numpy.cos(ha * numpy.pi / 180.0)
    sh = numpy.sin(ha * numpy.pi / 180.0)
    cd = numpy.cos(dec * numpy.pi / 180.0)
    sd = numpy.sin(dec * numpy.pi / 180.0)
    num = cl * sh
    den = sl * cd - cl * sd * ch
    return numpy.arctan2(num, den) * 180.0 / numpy.pi

def makeAZZA(npix=256):
    """
    Make azimuth and zenith angle arrays for a square image of side npix
    Projection is sine, all-sky
    Returns (az,za). Angles are in radian.
    """
    # build az and za arrays
    z = numpy.arange(npix, dtype=numpy.float32) - npix / 2
    x = numpy.empty((npix, npix), dtype=numpy.float32)
    y = numpy.empty((npix, npix), dtype=numpy.float32)
    for i in range(npix):
        y[i, 0:] = z
        x[0:, i] = z
    d = numpy.sqrt(x * x + y * y) / (npix / 2)
    # only select pixels above horizon
    t = (d <= 1.0)
    za = numpy.ones((npix, npix), dtype=numpy.float32) * numpy.pi / 2.0
    za[t] = numpy.arcsin(d[t])
    az = numpy.arctan2(y, x)
    return az, za

def makeUnpolInstrumentalResponse(j1, j2):
    """
    Form the visibility matrix in instrumental response from two Jones
    matrices assuming unpolarised sources (hence the brightness matrix is
    the identity matrix)
    Input: j1,j2: Jones matrices of dimension[za][az][2][2]
    Returns: [za][az][[xx,xy],[yx,yy]] where "X" and "Y" are defined by the receptors
    of the Dipole object used in the ApertureArray. Hence to get "XX", you want
    result[za][az][0][0] and for "YY" you want result[za][az][1][1]
    """
    result = numpy.empty_like(j1)

    result[:, :, 0, 0] = j1[:, :, 0, 0] * j2[:, :, 0, 0].conjugate() + j1[:, :, 0, 1] * j2[:, :, 0, 1].conjugate()
    result[:, :, 1, 1] = j1[:, :, 1, 0] * j2[:, :, 1, 0].conjugate() + j1[:, :, 1, 1] * j2[:, :, 1, 1].conjugate()
    result[:, :, 0, 1] = j1[:, :, 0, 0] * j2[:, :, 1, 0].conjugate() + j1[:, :, 0, 1] * j2[:, :, 1, 1].conjugate()
    result[:, :, 1, 0] = j1[:, :, 1, 0] * j2[:, :, 0, 0].conjugate() + j1[:, :, 1, 1] * j2[:, :, 0, 1].conjugate()
    return result

# def makePolInstrumentResponse(j1, j2, b):
#   """
#   Form the instrument response from two Jones matrices with an
#   arbitrary source brightness matrix, hence arbitrary polarisation
#   Returns: (xx,yy,xy,yx) where "X" and "Y" are defined by the receptors
#   of the Dipole object used in the ApertureArray
#   """
#   # FIXME: need to work out how to do this in vectorised way.

def plotDipoleJones(d, freq=155e6):
    """
    Utility to make a plot of a dipole Jones matrix for debugging
    """

    import matplotlib.pyplot as plt
    (az, za) = makeAZZA()
    logger.info("plotting dipole Jones response for atype: " + d.atype + ', freq (MHz): ' + str(freq / 1e6))
    j = d.getJones(az, za, freq)

    plt.imshow(j[:, :, 0, 0].real)
    plt.title('MWA ' + str(freq / 1e6) + 'MHz dipole J00 voltage real')
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_J00_voltage_re_' + str(freq / 1e6) + 'MHz_dipole.png')
    plt.clf()

    plt.imshow(j[:, :, 0, 1].real)
    plt.title('MWA ' + str(freq / 1e6) + 'MHz dipole J01 voltage real')
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_J01_voltage_re_' + str(freq / 1e6) + 'MHz_dipole.png')
    plt.clf()

    plt.imshow(j[:, :, 1, 0].real)
    plt.title('MWA ' + str(freq / 1e6) + 'MHz dipole J10 voltage real')
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_J10_voltage_re_' + str(freq / 1e6) + 'MHz_dipole.png')
    plt.clf()

    plt.imshow(j[:, :, 1, 1].real)
    plt.title('MWA ' + str(freq / 1e6) + 'MHz dipole J11 voltage real')
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_J11_voltage_re_' + str(freq / 1e6) + 'MHz_dipole.png')
    plt.clf()

def plotArrayJones(j, freq, za):
    """
    Utility to plot the output of tile Jones matrices
    """
    import matplotlib.pyplot as plt

    plt.imshow(numpy.abs(j[:, :, 0, 0]))
    plt.title('MWA ' + str(freq / 1e6) + 'MHz J00 voltage mag ZA=' + str(za))
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_J00_voltage_mag_' + str(freq / 1e6) + 'MHz_ZA' + str(za) + '.png')
    plt.clf()

    plt.imshow(numpy.abs(j[:, :, 0, 1]))
    plt.title('MWA ' + str(freq / 1e6) + 'MHz J01 voltage mag ZA=' + str(za))
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_J01_voltage_mag_' + str(freq / 1e6) + 'MHz_ZA' + str(za) + '.png')
    plt.clf()

    plt.imshow(numpy.abs(j[:, :, 1, 0]))
    plt.title('MWA ' + str(freq / 1e6) + 'MHz J10 voltage mag ZA=' + str(za))
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_J10_voltage_mag_' + str(freq / 1e6) + 'MHz_ZA' + str(za) + '.png')
    plt.clf()

    plt.imshow(numpy.abs(j[:, :, 1, 1]))
    plt.title('MWA ' + str(freq / 1e6) + 'MHz J11 voltage mag ZA=' + str(za))
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_J11_voltage_mag_' + str(freq / 1e6) + 'MHz_ZA' + str(za) + '.png')
    plt.clf()

def plotVisResponse(j, freq, za):
    """
    Utility to plot the visibilty XX,YY,XY and YX response of the array for
    an unpolarised 1Jy source
    Input: j a visibility matrix (complex) of dimensions [za][az][2][2]
    """
    import matplotlib.pyplot as plt

    vis = makeUnpolInstrumentalResponse(j, j)
    plt.imshow(numpy.abs(vis[:, :, 0, 0]))
    plt.title('MWA ' + str(freq / 1e6) + 'MHz XX mag ZA=' + str(za))
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_XX_mag_' + str(freq / 1e6) + 'MHz_ZA=' + str(za) + '.png')
    plt.clf()

    plt.imshow(numpy.abs(vis[:, :, 1, 1]))
    plt.title('MWA ' + str(freq / 1e6) + 'MHz YY mag ZA=' + str(za))
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_YY_mag_' + str(freq / 1e6) + 'MHz_ZA=' + str(za) + '.png')
    plt.clf()

    plt.imshow(numpy.abs(vis[:, :, 0, 1]))
    plt.title('MWA ' + str(freq / 1e6) + 'MHz XY mag ZA=' + str(za))
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_XY_mag_' + str(freq / 1e6) + 'MHz_ZA=' + str(za) + '.png')
    plt.clf()

    plt.imshow(numpy.abs(vis[:, :, 1, 0]))
    plt.title('MWA ' + str(freq / 1e6) + 'MHz YX mag ZA=' + str(za))
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_YX_mag_' + str(freq / 1e6) + 'MHz_ZA=' + str(za) + '.png')
    plt.clf()

def plotArrayFactors(ax, ay, za):
    """
    Utility to make a plot of an array factor for debugging
    """
    import matplotlib.pyplot as plt

    plt.imshow(numpy.abs(ax))
    plt.title('MWA ' + str(freq / 1e6) + 'MHz X voltage mag array factor. ZA=' + str(za))
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_array_voltage_X_mag_' + str(freq / 1e6) + 'MHz_ZA' + str(za) + '.png')
    plt.clf()

    plt.imshow(numpy.angle(ax))
    plt.title('MWA ' + str(freq / 1e6) + 'MHz X voltage phase array factor ZA=' + str(za))
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_array_voltage_X_ph_' + str(freq / 1e6) + 'MHz_ZA' + str(za) + '.png')
    plt.clf()

    plt.imshow(numpy.abs(ay))
    plt.title('MWA ' + str(freq / 1e6) + 'MHz Y voltage mag array factor ZA=' + str(za))
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_array_voltage_Y_mag_' + str(freq / 1e6) + 'MHz_ZA' + str(za) + '.png')
    plt.clf()

    plt.imshow(numpy.angle(ay))
    plt.title('MWA ' + str(freq / 1e6) + 'MHz Y voltage phase array factor ZA=' + str(za))
    plt.colorbar()
    plt.gca().invert_yaxis()
    plt.savefig('/tmp/MWA_array_voltage_Y_ph_' + str(freq / 1e6) + 'MHz_ZA' + str(za) + '.png')
    plt.clf()

# execute a series of tests if invoked from the command line
if __name__ == "__main__":

    logger.setLevel(logging.DEBUG)
    doplots = True
    lat = config.MWAPOS.latitude.deg
    freq = 216e6
    (az, za) = makeAZZA()

    # parallactic angle
    (ha, dec) = h2e(az, za, lat)
    pa = calcParallacticAngle(ha, dec, lat)

    d = Dipole('lookup')
    #    d = Dipole()
    if doplots:
        plotDipoleJones(d, freq)

    # tests with full tile
    # az=0, za=14 degs
    delays1 = numpy.array([[6, 6, 6, 6,
                            4, 4, 4, 4,
                            2, 2, 2, 2,
                            0, 0, 0, 0],
                           [6, 6, 6, 6,
                            4, 4, 4, 4,
                            2, 2, 2, 2,
                            0, 0, 0, 0]],
                          dtype=numpy.float32)
    za_delays = {'0': delays1 * 0, '14': delays1, '28': delays1 * 2}

    tile = ApertureArray(dipoles=[d] * 16)
    if doplots:
        for za_delay in za_delays:
            logger.debug("ZA is: %s. Delays are: %r" % (za_delay, za_delays[za_delay]))
            (ax0, ay0) = tile.getArrayFactor(az, za, freq, za_delays[za_delay])
            logger.info("plotting Array factor voltage for ZA " + str(za_delay))
            plotArrayFactors(ax0, ay0, za_delay)
            logger.info("plotting tile Jones response for ZA " + str(za_delay))
            j = tile.getResponse(az, za, freq, za_delays[za_delay])
            plotArrayJones(j, freq, za_delay)
            logger.info("Plotting visbility response for two identical tiles ZA " + str(za_delay))
            plotVisResponse(j, freq, za_delay)


# ║ zerowanie macierzy gdy "0" w liscie║
# ║ Transponowanie tabeli              ║
# ║                                    ║
# ╚════════════════════════════════════╝
from typing import List
class Solution:
    def setZeroes(self, matrix: List[List[int]]) -> None:
        """
        Do not return anything, modify matrix in-place instead.
        """
        # print(matrix, id(matrix))
        zerolist = []
        for i in range(len(matrix)):
            for j in range(len(matrix[i])):
                if matrix[i][j] == 0:
                    zerolist.append(j)
            if 0 in matrix[i]:
                matrix[i] = [0]*len(matrix[i])
                # print(matrix[i])
        # print(zerolist)

        for x in zerolist:
            for i in range(len(matrix)):
                matrix[i][x] = 0
                # print(x,i)

        # print(matrix, id(matrix))

Wynik = Solution()

matrix = [[1,1,1],[1,0,1],[1,1,1]]
Wynik1 = Wynik.setZeroes(matrix= matrix)
print("******\nWynik1", Wynik1, "\n******")

LeetCode/Top Interview 150/Matrix/73. Set Matrix Zeroes.py


class Pila:
    def __init__(self):
        self.items = ["Z0"]

    def esta_vacia(self):
        return self.items == []

    def push(self, item):
        self.items.append(item)

    def pop(self):
        if not self.esta_vacia():
            return self.items.pop()
        return None

    def peek(self):
        if not self.esta_vacia():
            return self.items[-1]
        return None

    def __str__(self):
        return str(self.items)

class AutomataDePila:
    def __init__(self):
        self.estados = {'q0', 'q1', 'q2'}
        self.estado_actual = 'q0'
        self.entrada = ""
        self.pila = Pila()
        self.estado_aceptacion = {'q2'}

    def reset(self):
        self.estado_actual = 'q0'
        self.pila = Pila()

    def procesar_caracter(self, caracter):
        cima_pila = self.pila.peek()
        if self.estado_actual == 'q0':
            if caracter == '0' and (cima_pila == 'Z0' or cima_pila == 'X'):
                self.pila.pop()
                self.pila.push('X')
                self.pila.push('X')
            elif caracter == '1' and cima_pila == 'X':
                self.estado_actual = 'q1'
                self.pila.pop()
            else:
                return False
        elif self.estado_actual == 'q1':
            if caracter == '1' and cima_pila == 'X':
                self.pila.pop()
            elif caracter == '':
                if cima_pila == 'Z0':
                    self.estado_actual = 'q2'
                    return True
            else:
                return False
        return True

    def procesar_entrada(self, entrada):
        self.reset()
        self.entrada = entrada
        for caracter in self.entrada:
            if not self.procesar_caracter(caracter):
                return False
        if self.estado_actual == 'q1':
            return self.procesar_caracter('')  # procesa el carácter vacío final
        return self.estado_actual in self.estado_aceptacion

    def acepta(self, entrada):
        return self.procesar_entrada(entrada)

# Instancia del autómata de pila
aut = AutomataDePila()

# Ejemplos de cadenas
cadenas = ["00", "0001", "01", "001", "00011", "11"]

# Probar el autómata con las cadenas
for cadena in cadenas:
    resultado = aut.acepta(cadena)
    print(f"La cadena '{cadena}' {'es aceptada' if resultado else 'no es aceptada'} por el autómata.")


from nltk.corpus import wordnet
from nltk.tokenize import word_tokenize
from random import randint
import nltk.data

text = "He is a good boy, that speaks English fast."

output = ""

tokenizer = nltk.data.load('tokenizers/punkt/english.pickle')
tokenized = tokenizer.tokenize(text)
words = word_tokenize(text)
tagged = nltk.pos_tag(words)
print(tagged)
for i in range(0, len(words)):
    replacements = []

    for syn in wordnet.synsets(words[i]):

        if tagged[i][1] == 'NN' or tagged[i][1] == 'DT' or tagged[i][1] == 'PRP' or tagged[i][1] == 'WP' or tagged[i][1] == 'MD' or tagged[i][1] == 'VBZ':
            break

        word_type = tagged[i][1][0].lower()

        if syn.name().find("." + word_type + "."):
            r = syn.name()[0:syn.name().find(".")]

            replacements.append(r)

    if len(replacements) > 0:
        replacement = replacements[randint(0, len(replacements) - 1)]
        output = output + " " + replacement
    else:

        output = output + " " + words[i]

print(output)


import numpy as np

from gpaw.utilities.blas import gemm
from gpaw.utilities import pack, unpack2
from gpaw.utilities.timing import nulltimer

class EmptyWaveFunctions:
    def __nonzero__(self):
        return False

    def set_eigensolver(self, eigensolver):
        pass

    def set_orthonormalized(self, flag):
        pass

    def estimate_memory(self, mem):
        mem.set('Unknown WFs', 0)

class WaveFunctions(EmptyWaveFunctions):
    """...

    setups:
        List of setup objects.
    symmetry:
        Symmetry object.
    kpt_u:
        List of **k**-point objects.
    nbands: int
        Number of bands.
    nspins: int
        Number of spins.
    dtype: dtype
        Data type of wave functions (float or complex).
    bzk_kc: ndarray
        Scaled **k**-points used for sampling the whole
        Brillouin zone - values scaled to [-0.5, 0.5).
    ibzk_kc: ndarray
        Scaled **k**-points in the irreducible part of the
        Brillouin zone.
    weight_k: ndarray
        Weights of the **k**-points in the irreducible part
        of the Brillouin zone (summing up to 1).
    kpt_comm:
        MPI-communicator for parallelization over **k**-points.
    """
    def __init__(self, gd, nvalence, setups, bd, dtype,
                 world, kd, timer=None):
        if timer is None:
            timer = nulltimer

        self.gd = gd
        self.nspins = kd.nspins
        self.nvalence = nvalence
        self.bd = bd
        self.nbands = self.bd.nbands #XXX
        self.mynbands = self.bd.mynbands #XXX
        self.dtype = dtype
        self.world = world
        self.kd = kd
        self.band_comm = self.bd.comm #XXX
        self.timer = timer
        self.rank_a = None

        # XXX Remember to modify aseinterface when removing the following
        # attributes from the wfs object
        self.gamma = kd.gamma
        self.kpt_comm = kd.comm
        self.bzk_kc = kd.bzk_kc
        self.ibzk_kc = kd.ibzk_kc
        self.ibzk_qc = kd.ibzk_qc
        self.weight_k = kd.weight_k
        self.symmetry = kd.symmetry
        self.nibzkpts = kd.nibzkpts

        self.kpt_u = kd.create_k_points(self.gd)

        self.eigensolver = None
        self.positions_set = False

        self.set_setups(setups)

    def set_setups(self, setups):
        self.setups = setups

    def set_eigensolver(self, eigensolver):
        self.eigensolver = eigensolver

    def __nonzero__(self):
        return True

    def calculate_density_contribution(self, nt_sG):
        """Calculate contribution to pseudo density from wave functions."""
        nt_sG.fill(0.0)
        for kpt in self.kpt_u:
            self.add_to_density_from_k_point(nt_sG, kpt)
        self.band_comm.sum(nt_sG)
        self.kpt_comm.sum(nt_sG)

        self.timer.start('Symmetrize density')
        for nt_G in nt_sG:
            self.symmetry.symmetrize(nt_G, self.gd)
        self.timer.stop('Symmetrize density')

    def add_to_density_from_k_point(self, nt_sG, kpt):
        self.add_to_density_from_k_point_with_occupation(nt_sG, kpt, kpt.f_n)

    def get_orbital_density_matrix(self, a, kpt, n):
        """Add the nth band density from kpt to density matrix D_sp"""
        ni = self.setups[a].ni
        D_sii = np.zeros((self.nspins, ni, ni))
        P_i = kpt.P_ani[a][n]
        D_sii[kpt.s] += np.outer(P_i.conj(), P_i).real
        D_sp = [pack(D_ii) for D_ii in D_sii]
        return D_sp

    def calculate_atomic_density_matrices_k_point(self, D_sii, kpt, a, f_n):
        if kpt.rho_MM is not None:
            P_Mi = kpt.P_aMi[a]
            #P_Mi = kpt.P_aMi_sparse[a]
            #ind = get_matrix_index(kpt.P_aMi_index[a])
            #D_sii[kpt.s] += np.dot(np.dot(P_Mi.T.conj(), kpt.rho_MM),
            #                       P_Mi).real
            rhoP_Mi = np.zeros_like(P_Mi)
            D_ii = np.zeros(D_sii[kpt.s].shape, kpt.rho_MM.dtype)
            #gemm(1.0, P_Mi, kpt.rho_MM[ind.T, ind], 0.0, tmp)
            gemm(1.0, P_Mi, kpt.rho_MM, 0.0, rhoP_Mi)
            gemm(1.0, rhoP_Mi, P_Mi.T.conj().copy(), 0.0, D_ii)
            D_sii[kpt.s] += D_ii.real
            #D_sii[kpt.s] += dot(dot(P_Mi.T.conj().copy(),
            #                        kpt.rho_MM[ind.T, ind]), P_Mi).real
        else:
            P_ni = kpt.P_ani[a]
            D_sii[kpt.s] += np.dot(P_ni.T.conj() * f_n, P_ni).real

        if hasattr(kpt, 'c_on'):
            for ne, c_n in zip(kpt.ne_o, kpt.c_on):
                d_nn = ne * np.outer(c_n.conj(), c_n)
                D_sii[kpt.s] += np.dot(P_ni.T.conj(), np.dot(d_nn, P_ni)).real

    def calculate_atomic_density_matrices(self, D_asp):
        """Calculate atomic density matrices from projections."""
        f_un = [kpt.f_n for kpt in self.kpt_u]
        self.calculate_atomic_density_matrices_with_occupation(D_asp, f_un)

    def calculate_atomic_density_matrices_with_occupation(self, D_asp, f_un):
        """Calculate atomic density matrices from projections with
        custom occupation f_un."""
        # Varying f_n used in calculation of response part of GLLB-potential
        for a, D_sp in D_asp.items():
            ni = self.setups[a].ni
            D_sii = np.zeros((self.nspins, ni, ni))
            for f_n, kpt in zip(f_un, self.kpt_u):
                self.calculate_atomic_density_matrices_k_point(D_sii, kpt, a,
                                                               f_n)
            D_sp[:] = [pack(D_ii) for D_ii in D_sii]
            self.band_comm.sum(D_sp)
            self.kpt_comm.sum(D_sp)

        self.symmetrize_atomic_density_matrices(D_asp)

    def symmetrize_atomic_density_matrices(self, D_asp):
        if len(self.symmetry.op_scc) > 1:
            all_D_asp = []
            for a, setup in enumerate(self.setups):
                D_sp = D_asp.get(a)
                if D_sp is None:
                    ni = setup.ni
                    D_sp = np.empty((self.nspins, ni * (ni + 1) // 2))
                self.gd.comm.broadcast(D_sp, self.rank_a[a])
                all_D_asp.append(D_sp)

            for s in range(self.nspins):
                D_aii = [unpack2(D_sp[s]) for D_sp in all_D_asp]
                for a, D_sp in D_asp.items():
                    setup = self.setups[a]
                    D_sp[s] = pack(setup.symmetrize(a, D_aii,
                                                    self.symmetry.a_sa))

    def set_positions(self, spos_ac):
        self.positions_set = False
        rank_a = self.gd.get_ranks_from_positions(spos_ac)

        """
        # If both old and new atomic ranks are present, start a blank dict if
        # it previously didn't exist but it will needed for the new atoms.
        if (self.rank_a is not None and rank_a is not None and
            self.kpt_u[0].P_ani is None and (rank_a == self.gd.comm.rank).any()):
            for kpt in self.kpt_u:
                kpt.P_ani = {}
        """

        if self.rank_a is not None and self.kpt_u[0].P_ani is not None:
            self.timer.start('Redistribute')
            requests = []
            mynks = len(self.kpt_u)
            flags = (self.rank_a != rank_a)
            my_incoming_atom_indices = np.argwhere(np.bitwise_and(flags, \
                rank_a == self.gd.comm.rank)).ravel()
            my_outgoing_atom_indices = np.argwhere(np.bitwise_and(flags, \
                self.rank_a == self.gd.comm.rank)).ravel()

            for a in my_incoming_atom_indices:
                # Get matrix from old domain:
                ni = self.setups[a].ni
                P_uni = np.empty((mynks, self.mynbands, ni), self.dtype)
                requests.append(self.gd.comm.receive(P_uni, self.rank_a[a],
                                                     tag=a, block=False))
                for myu, kpt in enumerate(self.kpt_u):
                    assert a not in kpt.P_ani
                    kpt.P_ani[a] = P_uni[myu]

            for a in my_outgoing_atom_indices:
                # Send matrix to new domain:
                P_uni = np.array([kpt.P_ani.pop(a) for kpt in self.kpt_u])
                requests.append(self.gd.comm.send(P_uni, rank_a[a],
                                                  tag=a, block=False))
            self.gd.comm.waitall(requests)
            self.timer.stop('Redistribute')

        self.rank_a = rank_a

        if self.symmetry is not None:
            self.symmetry.check(spos_ac)

    def allocate_arrays_for_projections(self, my_atom_indices):
        if not self.positions_set and self.kpt_u[0].P_ani is not None:
            # Projections have been read from file - don't delete them!
            pass
        else:
            for kpt in self.kpt_u:
                kpt.P_ani = {}
            for a in my_atom_indices:
                ni = self.setups[a].ni
                for kpt in self.kpt_u:
                    kpt.P_ani[a] = np.empty((self.mynbands, ni), self.dtype)

    def collect_eigenvalues(self, k, s):
        return self.collect_array('eps_n', k, s)

    def collect_occupations(self, k, s):
        return self.collect_array('f_n', k, s)

    def collect_array(self, name, k, s, subset=None):
        """Helper method for collect_eigenvalues and collect_occupations.

        For the parallel case find the rank in kpt_comm that contains
        the (k,s) pair, for this rank, collect on the corresponding
        domain a full array on the domain master and send this to the
        global master."""

        kpt_u = self.kpt_u
        kpt_rank, u = self.kd.get_rank_and_index(s, k)

        if self.kpt_comm.rank == kpt_rank:
            a_nx = getattr(kpt_u[u], name)

            if subset is not None:
                a_nx = a_nx[subset]

            # Domain master send this to the global master
            if self.gd.comm.rank == 0:
                if self.band_comm.size == 1:
                    if kpt_rank == 0:
                        return a_nx
                    else:
                        self.kpt_comm.ssend(a_nx, 0, 1301)
                else:
                    b_nx = self.bd.collect(a_nx)
                    if self.band_comm.rank == 0:
                        if kpt_rank == 0:
                            return b_nx
                        else:
                            self.kpt_comm.ssend(b_nx, 0, 1301)

        elif self.world.rank == 0 and kpt_rank != 0:
            # Find shape and dtype:
            a_nx = getattr(kpt_u[0], name)
            shape = (self.nbands,) + a_nx.shape[1:]
            dtype = a_nx.dtype
            b_nx = np.zeros(shape, dtype=dtype)
            self.kpt_comm.receive(b_nx, kpt_rank, 1301)
            return b_nx

    def collect_auxiliary(self, value, k, s, shape=1, dtype=float):
        """Helper method for collecting band-independent scalars/arrays.

        For the parallel case find the rank in kpt_comm that contains
        the (k,s) pair, for this rank, collect on the corresponding
        domain a full array on the domain master and send this to the
        global master."""

        kpt_u = self.kpt_u
        kpt_rank, u = self.kd.get_rank_and_index(s, k)

        if self.kpt_comm.rank == kpt_rank:
            if isinstance(value, str):
                a_o = getattr(kpt_u[u], value)
            else:
                a_o = value[u] # assumed list

            # Make sure data is a mutable object
            a_o = np.asarray(a_o)

            if a_o.dtype is not dtype:
                a_o = a_o.astype(dtype)

            # Domain master send this to the global master
            if self.gd.comm.rank == 0:
                if kpt_rank == 0:
                    return a_o
                else:
                    self.kpt_comm.send(a_o, 0, 1302)

        elif self.world.rank == 0 and kpt_rank != 0:
            b_o = np.zeros(shape, dtype=dtype)
            self.kpt_comm.receive(b_o, kpt_rank, 1302)
            return b_o

    def collect_projections(self, k, s):
        """Helper method for collecting projector overlaps across domains.

        For the parallel case find the rank in kpt_comm that contains
        the (k,s) pair, for this rank, send to the global master."""

        kpt_rank, u = self.kd.get_rank_and_index(s, k)

        natoms = len(self.rank_a) # it's a hack...
        nproj = sum([setup.ni for setup in self.setups])

        if self.world.rank == 0:
            if kpt_rank == 0:
                P_ani = self.kpt_u[u].P_ani
            mynu = len(self.kpt_u)
            all_P_ni = np.empty((self.nbands, nproj), self.dtype)
            for band_rank in range(self.band_comm.size):
                nslice = self.bd.get_slice(band_rank)
                i = 0
                for a in range(natoms):
                    ni = self.setups[a].ni
                    if kpt_rank == 0 and band_rank == 0 and a in P_ani:
                        P_ni = P_ani[a]
                    else:
                        P_ni = np.empty((self.mynbands, ni), self.dtype)
                        world_rank = (self.rank_a[a] +
                                      kpt_rank * self.gd.comm.size *
                                      self.band_comm.size +
                                      band_rank * self.gd.comm.size)
                        self.world.receive(P_ni, world_rank, 1303 + a)
                    all_P_ni[nslice, i:i + ni] = P_ni
                    i += ni
                assert i == nproj
            return all_P_ni

        elif self.kpt_comm.rank == kpt_rank: # plain else works too...
            P_ani = self.kpt_u[u].P_ani
            for a in range(natoms):
                if a in P_ani:
                    self.world.ssend(P_ani[a], 0, 1303 + a)

    def get_wave_function_array(self, n, k, s):
        """Return pseudo-wave-function array.

        For the parallel case find the rank in kpt_comm that contains
        the (k,s) pair, for this rank, collect on the corresponding
        domain a full array on the domain master and send this to the
        global master."""

        kpt_rank, u = self.kd.get_rank_and_index(s, k)
        band_rank, myn = self.bd.who_has(n)

        size = self.world.size
        rank = self.world.rank

        if self.kpt_comm.rank == kpt_rank:
            psit1_G = self._get_wave_function_array(u, myn)
            if size == 1:
                return psit1_G
            if self.band_comm.rank == band_rank:
                psit_G = self.gd.collect(psit1_G)

                if kpt_rank == 0 and band_rank == 0:
                    if rank == 0:
                        return psit_G

                # Domain master send this to the global master
                if self.gd.comm.rank == 0:
                    self.world.send(psit_G, 0, 1398)

        if rank == 0:
            # allocate full wavefunction and receive
            psit_G = self.gd.empty(dtype=self.dtype, global_array=True)
            world_rank = (kpt_rank * self.gd.comm.size *
                          self.band_comm.size +
                          band_rank * self.gd.comm.size)
            self.world.receive(psit_G, world_rank, 1398)
            return psit_G

    def _get_wave_function_array(self, u, n):
        raise NotImplementedError


"""707. Design Linked List"""

from typing import Optional

class Node:
    def __init__(self, val: int, next_node: Optional["Node"] = None) -> None:
        self.val = val
        self.next = next_node

class MyLinkedList:
    def __init__(self):
        self.head: Optional[Node] = None

    def get(self, index: int) -> int:
        if not self.head:
            return -1

        node = self.head

        for _ in range(index):
            if not node.next:
                return -1

            node = node.next

        return node.val

    def addAtHead(self, val: int) -> None:
        self.head = Node(val, self.head)

    def addAtTail(self, val: int) -> None:
        if not self.head:
            self.head = Node(val)
            return

        node = self.head

        while node.next:
            node = node.next

        node.next = Node(val)

    def addAtIndex(self, index: int, val: int) -> None:
        if index == 0:
            self.addAtHead(val)
            return

        node = self.head

        for _ in range(index):
            if not node:
                return

            prev = node
            node = node.next

        prev.next = Node(val, node)

    def deleteAtIndex(self, index: int) -> None:
        if not self.head:
            return

        if index == 0:
            self.head = self.head.next
            return

        node = self.head

        for _ in range(index):
            if not node.next:
                return

            prev = node
            node = node.next

        prev.next = node.next

# Your MyLinkedList object will be instantiated and called as such:
# obj = MyLinkedList()
# param_1 = obj.get(index)
# obj.addAtHead(val)
# obj.addAtTail(val)
# obj.addAtIndex(index,val)
# obj.deleteAtIndex(index)


import math

import torch

def fn(u):
    return math.cos((u + 0.008) / 1.008 * math.pi / 2)**2

def beta_schedule(schedule,
                  num_timesteps=1000,
                  init_beta=None,
                  last_beta=None):
    '''
    This code defines a function beta_schedule that generates a sequence of beta values based on the given input
    parameters. These beta values can be used in video diffusion processes. The function has the following parameters:
        schedule(str): Determines the type of beta schedule to be generated. It can be 'linear', 'linear_sd',
                      'quadratic', or 'cosine'.
        num_timesteps(int, optional): The number of timesteps for the generated beta schedule. Default is 1000.
        init_beta(float, optional): The initial beta value. If not provided, a default value is used based on the
                                    chosen schedule.
        last_beta(float, optional): The final beta value. If not provided, a default value is used based on the
                                    chosen schedule.
    The function returns a PyTorch tensor containing the generated beta values.
    The beta schedule is determined by the schedule parameter:
        1.Linear: Generates a linear sequence of beta values betweeninit_betaandlast_beta.
        2.Linear_sd: Generates a linear sequence of beta values between the square root of init_beta and the square root
                     oflast_beta, and then squares the result.
        3.Quadratic: Similar to the 'linear_sd' schedule, but with different default values forinit_betaandlast_beta.
        4.Cosine: Generates a sequence of beta values based on a cosine function, ensuring the values are between 0
                  and 0.999.
    If an unsupported schedule is provided, a ValueError is raised with a message indicating the issue.
    '''
    if schedule == 'linear':
        scale = 1000.0 / num_timesteps
        init_beta = init_beta or scale * 0.0001
        last_beta = last_beta or scale * 0.02
        return torch.linspace(
            init_beta, last_beta, num_timesteps, dtype=torch.float64)
    elif schedule == 'linear_sd':
        return torch.linspace(
            init_beta**0.5, last_beta**0.5, num_timesteps,
            dtype=torch.float64)**2
    elif schedule == 'quadratic':
        init_beta = init_beta or 0.0015
        last_beta = last_beta or 0.0195
        return torch.linspace(
            init_beta**0.5, last_beta**0.5, num_timesteps,
            dtype=torch.float64)**2
    elif schedule == 'cosine':
        betas = []
        for step in range(num_timesteps):
            t1 = step / num_timesteps
            t2 = (step + 1) / num_timesteps
            betas.append(min(1.0 - fn(t2) / fn(t1), 0.999))
        return torch.tensor(betas, dtype=torch.float64)
    else:
        raise ValueError(f'Unsupported schedule: {schedule}')


##############################
def calc_length(in_Npts,in_d,in_ratio):
    if(in_ratio==1):
        length=in_d*(in_Npts-1)
    else:
        length=in_d*(in_ratio**(in_Npts-1)-1)/(in_ratio-1)
    return length
##############################
def calc_npts_and_ratio(in_L,in_d,in_D):
    assert(in_L>0)
    from numpy import log
    if(in_D==in_d):
        ratio=1
        npts=in_L/in_d
    else:
        ratio=(in_L/in_d-1)/(in_L/in_d-in_D/in_d)
        npts=log(in_D/in_d)/log(ratio)+2
    return npts,ratio
##############################
def calc_ratio_func(in_r,in_L,in_Npts,in_d):
    res=0
    if(in_r==1):
        res=(in_Npts-1)*in_d-in_L
    else:
        res=in_d*((in_r**(in_Npts-1)-1)/(in_r-1))-in_L
    return res
def calc_ratio(in_L,in_Npts,in_d,ratio_0=1.1):
    assert(in_L>0)
    from scipy.optimize import root
    from numpy import abs
    sol=root(calc_ratio_func,ratio_0,args=(in_L,in_Npts,in_d),\
        method='hybr')
    ratio=sol.x[0]
    length=calc_length(in_Npts,in_d,ratio)
    if(abs(length-in_L) > 1E-9):
        sol=root(calc_ratio_func,ratio_0,args=(in_L,in_Npts,in_d),\
            method='lm')
        ratio=sol.x[0]
        length=calc_length(in_Npts,in_d,ratio)
        if(abs(length-in_L) > 1E-9):
            print("L=%e,L_calc=%e,N=%d,d=%e,r=%e"%(in_L,length,in_Npts,in_d,ratio))
            print(sol)
            from traceback import print_stack
            print_stack()
            exit()
    return ratio
##############################
def calc_npts(in_L,in_d,in_Ratio,in_npts_bin=[2,100]):
    if in_Ratio==1:
        npts=in_L/in_d+1
    else:
        from numpy import arange,ones,abs,argmin
        npts_arr=arange(in_npts_bin[1]-in_npts_bin[0]+1)+in_npts_bin[0]
        res_arr=ones(npts_arr.shape)*in_L/in_d
        for i in range(npts_arr.size):
            res_arr[i]=abs(res_arr[i]-(in_Ratio**(npts_arr[i]-1)-1)/(in_Ratio-1))
        npts=npts_arr[argmin(res_arr)]
    return npts


"""
@param: item to have its type found
@pre: none
@post: returns 0 if its an integer, 1 if it's an operator(+,-,*,/,=), else 2
@complexity: best = worse = O(1)
"""
def find_type(item):
    try:
        item = int(item)
        item_type = 0
    except:
        if item in ("+", "-", "*", "/", "="):
            item_type = 1
        else:
            item_type = 2
    return item_type

"""
@param none
@pre find_type is defined
@post returns True if find_type is working properly,
            False if not
@complexity best = worse = O(1)
"""
def test_find_type():
    correct = True
    if 0 != find_type('1'):
        correct = False
    if 1 != find_type('+'):
        correct = False
    if 2 != find_type('+sa'):
        correct = False
    return correct

intro to comp science/python functions/find_type.py


import pkg_resources

from xblock.core import XBlock
from xblock.fields import Scope, Integer, Boolean
from xblock.fragment import Fragment

class XDoneXBlock(XBlock):

    """

    Button to set mark a current section as passed.

    """

    done = Boolean(
        default=False, scope=Scope.user_state,
        help="The mark a student finished task"
    )

    def resource_string(self, path):
        """Handy helper for getting resources from our kit."""
        data = pkg_resources.resource_string(__name__, path)
        return data.decode("utf8")

    def student_view(self, context=None):
        """
        The primary view of the XDoneXBlock, shown to students
        when viewing courses.
        """
        html = self.resource_string("static/html/xdone.html")
        frag = Fragment(html.format(self=self))
        frag.add_css(self.resource_string("static/css/xdone.css"))
        frag.add_javascript(self.resource_string("static/js/src/xdone.js"))
        frag.initialize_js('XDoneXBlock')
        return frag

    @XBlock.json_handler
    def done_submit(self, data, suffix=''):
        """
        handler, mark section as done
        """

        self.done = True
        return {"done": self.done}

    @staticmethod
    def workbench_scenarios():
        """A canned scenario for display in the workbench."""
        return [
            ("XDoneXBlock",
             """<xdone/>
             """),
            ("Multiple XDoneXBlock",
             """<vertical_demo>
                <xdone/>
                <xdone/>
                <xdone/>
                </vertical_demo>
             """),
        ]


import sys
import string

def Caesar(vl_plaint,vl_clave):
    vl_ciphert=""
    for letra in vl_plaint:
        if letra.isalpha():
            if letra.islower():
                vl_ciphert+= string.ascii_lowercase[(string.ascii_lowercase.index(letra)+vl_clave)%26]
            else:
                vl_ciphert+= string.ascii_uppercase[(string.ascii_uppercase.index(letra)+vl_clave)%26]
        else:
            vl_ciphert+=letra
    return vl_ciphert

arg=sys.argv
vl_lebarg=len(arg)

if vl_lebarg==1:
     print('No ha ingresado la clave')
elif vl_lebarg==2:
    vl_clave=int(arg[1])
    vl_plaint = input("Plain text:")
    vl_ciphert = Caesar(vl_plaint,vl_clave)
    print(vl_ciphert)
else:
    print('Ha ingresado demasiados parámetros')


""" jobs.py provides data structures and management for HPC jobs """

def new_job():
    """ Create and return an empty job data structure

    Returns:
    A dictionary with jobs structure

    """

    job = dict()
    job['name'] = ''
    job['creation-time'] = 0
    job['queuing-time'] = 0
    job['execution-time'] = 0
    job['completion-time'] = 0
    job['accounting-time'] = 0
    job['time-spent-queue'] = 0
    job['requested-mem'] = 0
    job['nb-reserved-nodes'] = 0
    job['requested-features'] = []
    job['process-id'] = 0
    job['nb-cores-allocated'] = 0
    job['nb-unique-nodes'] = 0
    job['exit-code'] = 0
    job['used-cput'] = 0
    job['power-consumption'] = 0.0
    job['used-mem'] = 0.0
    job['used-vmem'] = 0.0
    job['used-walltime'] = 0
    job['requested-resources'] = ''
    job['status'] = 'unknown' # Job status
    job['id'] = 'noId' # Job ID
    job['invalid-requested-resources'] = False # Flag to 1 if got invalid requested resources
    job['is-invalid'] = False # Flag to 1 is job is invalid
    job['nb-requested-nodes'] = 0
    job['nb-requested-cores'] = 0
    job['running-hosts'] = dict()

    return job


# pylint: disable=protected-access
from unittest.mock import patch

import pytest
from pendulum import now

from gsy_e.models.strategy.state import SmartMeterState
from gsy_e.models.strategy.state.base_states import UnexpectedStateException
from tests.test_state.test_consumption_state import TestConsumptionState
from tests.test_state.test_production_state import TestProductionState

class TestSmartMeterState(TestConsumptionState, TestProductionState):
    """Test PVState class."""

    @staticmethod
    def _setup_base_configuration():
        return SmartMeterState(), now()

    def _setup_default_state(self):
        smart_meter_state, past_time_slot = self._setup_base_configuration()
        current_time_slot = past_time_slot.add(minutes=15)
        smart_meter_state._available_energy_kWh[past_time_slot] = 1
        smart_meter_state._available_energy_kWh[current_time_slot] = 1
        smart_meter_state._desired_energy_Wh[past_time_slot] = 1
        smart_meter_state._desired_energy_Wh[current_time_slot] = 1
        smart_meter_state._energy_production_forecast_kWh[past_time_slot] = 1
        smart_meter_state._energy_production_forecast_kWh[current_time_slot] = 1
        smart_meter_state._energy_requirement_Wh[past_time_slot] = 1
        smart_meter_state._energy_requirement_Wh[current_time_slot] = 1
        return past_time_slot, current_time_slot, smart_meter_state

    def test_delete_past_state_values_market_slot_not_in_past(self):
        past_time_slot, current_time_slot, smart_meter_state = self._setup_default_state()
        with patch("gsy_e.gsy_e_core.util.ConstSettings.SettlementMarketSettings."
                   "ENABLE_SETTLEMENT_MARKETS", True):
            smart_meter_state.delete_past_state_values(current_time_slot)
            assert smart_meter_state.get_available_energy_kWh(
                past_time_slot, default_value=None) is not None
            assert smart_meter_state.get_energy_production_forecast_kWh(
                past_time_slot, default_value=None) is not None
            assert smart_meter_state.get_energy_requirement_Wh(
                past_time_slot, default_value=None) is not None
            assert smart_meter_state.get_desired_energy_Wh(
                past_time_slot, default_value=None) is not None

    def test_delete_past_state_values_market_slot_in_past(self):
        past_time_slot, current_time_slot, smart_meter_state = self._setup_default_state()
        with patch("gsy_e.gsy_e_core.util.ConstSettings.SettlementMarketSettings."
                   "ENABLE_SETTLEMENT_MARKETS", False):
            smart_meter_state.delete_past_state_values(current_time_slot)
            assert smart_meter_state.get_available_energy_kWh(
                past_time_slot) == 0.0
            assert smart_meter_state.get_energy_production_forecast_kWh(
                past_time_slot) == 0.0
            assert smart_meter_state.get_energy_requirement_Wh(
                past_time_slot) == 0.0
            assert smart_meter_state.get_desired_energy_Wh(
                past_time_slot) == 0.0

    def test_get_energy_at_market_slot_production_and_consumption_raise_error(self):
        _, current_time_slot, smart_meter_state = self._setup_default_state()
        smart_meter_state.set_desired_energy(energy=0.0,
                                             time_slot=current_time_slot)
        with pytest.raises(UnexpectedStateException):
            smart_meter_state.get_energy_at_market_slot(time_slot=current_time_slot)


# 서로 다른 무게의 볼링공을 고르는 경우의 수
# 입력받기
n, m = map(int, input().split())
data = list(map(int, input().split()))
data.sort()

# 같은 무게의 공이 몇개인지
count = 1
result = 0

for i in range(0,len(data)-1):
    # 다음 공의 무게가 다른 경우
    if(data[i] != data[i+1]):
        # 같은 무게의 공이 있는 경우
        if(count!=1):
            result += (len(data) - i - 1) * count
        # 같은 무게의 공이 없는 경우
        else:
            result += len(data)-i-1
        count = 1
    # 다음 공의 무게가 같은 경우
    else:
        count += 1

# # 책에서 재시한 해결책
# # 1부터 10까지의 무게를 담을 수 있는 리스트 생성
# array = [0] * 11

# for x in data:
#     # 각 무게에 해당하는 볼링공의 개수 카운트
#     array[x] += 1

# # 1부터 m까지의 각 무게에 대해 처리
# for i in range(1, m+1):
#     n -= array[i] # 무게가 i인 볼링공의 개수(A가 선택할 수 있는 개수) 제외
#     result += array[i] * n # B가 선택하는 경우의 수와 곱하기

print(result)

그리디/볼링공_고르기.py


import requests
from models import models
import json
import base64
import os
import logging
import glob

default_file_name = "test"
settings="settings.txt"
status_empty=-1
logging.basicConfig(filename="log.txt", level=logging.DEBUG)

def listdir_nohidden(path):
    return glob.glob(os.path.join(path, '*'))

def main():   #get auth token, connect, and pull initial repo
    try:
        files_in_local_repo=""

        with open(settings, "r") as settings_file:
            #get settings
            file_contents=settings_file.readlines()
        model =   models(models.format(json.loads(file_contents[0])))
        # file_structure = {}
        # for root,dir_names,files in os.walk(model.file_path):
        #         file_structure.update(roots=dir_names)
        #         print(f"dirs: {dir_names}")
        #         print(f"files in folder: {files}")
        #         print(file_structure)
        ##TODO: Add recusively files/folders to file_structure
        #get all files in current directory
        files_in_local_repo = []
        for path in listdir_nohidden(model.file_path):
            #ignore file with personal data
            if(path.startswith(".") or os.path.basename(path) =="settings.txt"
                or os.path.basename(path) == "__pycache__"
                or os.path.basename(path) == "log.txt"):
                pass
            else:
                print(os.path.basename(path))
                files_in_local_repo.append(os.path.basename(path))
    except Exception as e:
        logging.debug(e.args)
    finally:
        update_file(files_in_local_repo,model)
    print(files_in_local_repo)

def put_(url,model,data):
    print("data for post: " , data)
    return requests.put(url = url, auth = ("{model.username}",model.token),json = data)

def get_(url,headers):
    return  requests.get(url, headers)

def post_(headers,model):
    payload = {
                "name":model.repo_name
                }
    return requests.post(url =" https://api.github.com/user/repos",
                                      headers=json.loads(headers), json= payload)

def get_repo(update_url,data=None,model=any):
    try:
        header_types = ["Accept","Authorization","X-GitHub-Api-Version"]
        header_data = ["application/vnd.github+json","Bearer {0}".format(model.token),"2022-11-28"]
        headers= json.dumps(dict(zip(header_types,header_data)))
        print(f"repo url: {update_url}")
        response = None
        if(data):
            response=put_(update_url,model,data)
        else:
            response=get_(update_url,headers)
        print("status: " + str(response.status_code))

        #create repo if none exists
        if(response.status_code == 404  or
           json.loads(response.text)["message"] == "This repository is empty."
           or json.loads(response.text)["message"] == "Not Found"):
           post_(headers,model)
           return status_empty
    except Exception as e:
        logging.debug(e.args)
    finally:
        model_list=[] #if response is a list there we have more than 1 file
                    #and we need to create different models separately
        if(isinstance(response,list)):
            for file in response:
                model_list.append(models(models.format(file))) #converts to object
                #write_to_file(f"{file_name}.py",response_json)
            return model_list
        else:
            d=json.loads(response.text)
           # test=models(models.format(d))
            return models(models.format(d))

def write_to_file(filename=f"{default_file_name}.py",data=None):
    if(data):
        with open(filename, 'w') as file:
            file.write(json.dumps(data,sort_keys=True,indent=4))
            #file.write("def "+ "get_"+ a +":" + "\n" + "\t" +"return " +b  )
            file.close()

def convert_file_base64(filename=None,path=None):
    try:
        with open(filename,'rb') as file:
            converted_file_contents = base64.b64encode(file.read())
            file.close()
        return converted_file_contents.decode("utf-8") #have to decode to utf-8 because
                                        #the stupid b symbol in front of byte string
    except Exception as e:
        logging.debug(e.args)

# def isBase64(s):
#     try:
#         return base64.b64encode(base64.b64decode(s)) == s
#     except Exception:
#         return False

###branch_url = "https://api.github.com/repos/%s/%s/branches/main"
###branch_models = get_repo(branch_url % (f"{user_name}",f"{repo_name}"))
###sha = branch_models.commit["sha"]

def update_file(files_in_local_repo=[],model = any):
    try:
        #get all files sha's in repo
        repo_url= f"https://api.github.com/repos/{model.username}/{model.repo_name}/contents/"
        payload = {
                                "message" : "commit",
                                "branch" : "master",
                                }

        if(len(files_in_local_repo) > 0):
            files_in_remote= get_repo(repo_url,model=model)
            if  files_in_remote != status_empty:
                # we need to get all the remote file names beforehand
                remote_files=[]
                try:
                    for k,v in files_in_remote.__dict__.items():
                        remote_files.append(v.name)
                except Exception as e:
                    logging.debug(e.args)

            for local_file in files_in_local_repo:
                #remove this from production, only for testing
                #so we dont upload git token
                if(local_file == "settings.txt"):
                    pass
                else:
                    if(files_in_remote !=status_empty and local_file in remote_files):
                        for key,remote_file in files_in_remote.__dict__.items():
                                #if file matches remote we update it
                                if local_file ==  remote_file.name :
                                    try:
                                        print(f"sha for file: {remote_file.name} {remote_file.sha}")
                                        path = f'''{model.file_path}\\{local_file}'''
                                        content =  convert_file_base64(path)
                                        #add sha for existing files
                                        payload.update({"sha" :  remote_file.sha,
                                        "content": content})
                                        repo_contents = get_repo(repo_url + local_file,payload,model=model)
                                    except Exception as e:
                                        print(e)
                    else:
                            #no sha for new files
                            print("new file to be added to remote: " + local_file)
                            path = f"{local_file}"
                            content =  convert_file_base64(path)
                            payload.update({"content": content})
                            repo_contents = get_repo(repo_url + local_file,payload,model=model)
    except Exception as e:
        print(repr(e.args))

#must be called at the beginning of the program
main()


# !/usr/bin/env python
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSDurabilityPolicy, QoSHistoryPolicy, QoSReliabilityPolicy
from rclpy.qos import QoSProfile
from geometry_msgs.msg import Twist
from sensor_msgs.msg import LaserScan

distance = float()

class LinearVelPublisher(Node):

    def __init__(self):

        super().__init__('linear_vel_publisher')
        self.publisher_ = self.create_publisher(Twist, '/cmd_vel', 10)

        timer_period = 0.5
        self.timer = self.create_timer(timer_period, self.timer_callback)

    def timer_callback(self):

        vel = Twist()

        vel.linear.x = 0.07
        vel.angular.z = 0.0

        self.publisher_.publish(vel)
        self.get_logger().info('Publishing velocity: linear x: "%f"' % vel.linear.x)

class AngularVelPublisher(Node):

    def __init__(self):

        super().__init__('angular_vel_publisher')
        self.publisher_ = self.create_publisher(Twist, '/cmd_vel', 10)

        timer_period = 0.5
        self.timer = self.create_timer(timer_period, self.timer_callback)

    def timer_callback(self):

        vel = Twist()

        vel.linear.x = 0.0
        vel.angular.z = 0.7

        self.publisher_.publish(vel)
        self.get_logger().info('Publishing velocity: angular z: "%f"' % vel.angular.z)

class LaserSubscriber(Node):

    def __init__(self):
        super().__init__('laser_subscriber')

        qos_profile = QoSProfile(depth=10)
        qos_profile.reliability = QoSReliabilityPolicy.BEST_EFFORT
        qos_profile.durability = QoSDurabilityPolicy.VOLATILE

        self.subscription = self.create_subscription(LaserScan, '/scan', self.listener_callback, qos_profile)
        self.subscription
        self.scan = LaserScan()

    def listener_callback(self, msg):

        global distance
        distance = msg.ranges[90]

        self.get_logger().info('I heard: min "%f" max "%f"'% (msg.ranges[90], msg.range_max))

def main(args=None):
    global distance

    rclpy.init(args=args)

    linear_vel_publisher = LinearVelPublisher()
    angular_vel_publisher = AngularVelPublisher()

    laser_subscriber = LaserSubscriber()

    #rclpy.spin(laser_subscriber)

    while rclpy.ok():
        rclpy.spin_once(laser_subscriber)
        rclpy.spin_once(linear_vel_publisher)
        print(distance)
        if distance <= 1.0:
            rclpy.spin_once(angular_vel_publisher)

    linear_vel_publisher.destroy_node()
    angular_vel_publisher.destroy_node()
    laser_subscriber.destroy_node()
    rclpy.shutdown()

if __name__ == "__main__":
    main()


import pygame
import board
import constants
import UserInterface

class Game:

    def __init__(self, screen):

        self.board = board.Board(screen)
        self.player = constants.RED
        self.winner = None
        self.running = True
        self.screen  = screen
        self.drawBoard()

    def reset(self):

        self.board.clear_board()
        self.player = constants.RED
        self.winner = None
        self.running = True

    def manage_drop(self, col):

        self.winner = self.board.check_win(self.board.board)

        if self.winner == None:
            sucess = self.board.manage_drop(col, self.player)

        if self.winner != None:

            pos = lambda row, col: (col * constants.SQUARE_SIZE + constants.SQUARE_SIZE // 2,
                                   (row+1) * constants.SQUARE_SIZE + constants.SQUARE_SIZE // 2)
            self.board.winLine(
                pos( self.winner[1][0][0], self.winner[1][0][1] ),
                pos( self.winner[1][1][0], self.winner[1][1][1] ) )

            sucess = False

            if self.winner == 10:
                UserInterface.Text("YELLOW WINS", constants.SQUARE_SIZE, 0, 0, constants.YELLOW).blit_text(self.screen)

            if self.winner == -10:
                UserInterface.Text("RED WINS", constants.SQUARE_SIZE, 0, 0, constants.RED).blit_text(self.screen)

            pygame.display.flip()

        if sucess:
            self.changePlayer()

        self.drawBoard()

    def changePlayer(self):

        if self.player == constants.RED:
            self.player = constants.YELLOW

        else:
            self.player = constants.RED

    def drawBoard(self):
        self.board.draw_board()

    def ai_move(self):
        return self.board.ai_move()


#!/usr/bin/python3
# -*- coding:utf-8 -*-
# @Author lvmengting
# @Time   2021/2/26 10:28
# @Email  13554221497@163.com
# @File   hand_file.py

"""
如何流式读取大文件
"""

import time
from functools import wraps

def cal_time(func):
    """
    计算程序运行时间的装饰器
    """
    @wraps(func)
    def do_something(*args, **kwargs):
        start_time = time.time()
        func(*args, **kwargs)
        print('time use', func.__name__, time.time() - start_time, sep=':  ')
    return do_something

# 标准做法
def count_nine(fname):
    """
    计算文件中含有多少个数字'9'
    """
    count = 0
    with open(fname) as file:
        for line in file:
            count += line.count('9')
    return count
    # 如何一个文件没有换行符, 大小为5个G, 占用内存会很大.

# 使用read方法分块读取, file.read(chunk_size)
def count_nine_v2(fname):
    """
    经过测试, 该方法虽然降低了内存, 但是运行时间提高.
    """
    count = 0
    block_size = 1024 * 8
    with open(fname) as fp:
        while True:
            chunk = fp.read(block_size)
            if not chunk:
                break
            count += chunk.count('9')
    return count

# 对方法2的代码进行解耦合
def chunked_file_reader(fp, block_size=1024 * 8):
    """
    利用生成器函数
    """
    while True:
        chunk = fp.read(block_size)
        if not chunk:
            break
        yield chunk

def count_nine_v3(fname):
    count = 0
    with open(fname) as fp:
        for chunk in chunked_file_reader(fp):
            count += chunk.count('9')
    return count

@cal_time
def Test():
    for i in range(1, 90000):
        count_nine('demo.txt')

@cal_time
def Test_v2():
    for i in range(1, 90000):
        count_nine_v2('demo.txt')

@cal_time
def Test_v3():
    for i in range(1, 90000):
        count_nine_v3('demo.txt')

Test()
Test_v2()
Test_v3()


# You are given the head of a linked list A. Check if the data inside it exists in non-decreasing order.

# Problem Constraints
# 1 <= size of linked list <= 105
# 1 <= value of nodes <= 109

# Input Format
# The first argument A is the head of a linked list.

# Output Format
# Return 1 if the data of nodes is non-decreasing and 0 otherwise

# Example Input
# Input 1:
# A = 1 -> 2 -> 3 -> 3
# Input 2:
# A = 4 -> 3 -> 2 -> 1

# Example Output
# Output 1:
# 1
# Output 2:
# 0

# Example Explanation
# For Input 1:
# The data of nodes follow a non-decreasing order.
# For Input 2:
# The data of nodes are not in non-decreasing order.

# Definition for singly-linked list.
# class ListNode:
#    def __init__(self, x):
#        self.val = x
#        self.next = None

class Solution:
    # @param A : head node of linked list
    # @return an integer
    def solve(self, a):
        temp=a
        i=0
        while(temp.next!=None):
            if temp.val>temp.next.val:
                return 0
            temp=temp.next
        return 1


from turtle import Screen
from paddle import Paddle
from mid_line import Mid_line
from Computer_Scoreboard import Scoreboard
from Ball import Ball
from layout import Layout
import time
from music import Music

screen = Screen()
screen.bgcolor("red")
screen.title("Pong Game")
screen.setup(width=800, height=600)
screen.tracer(0)

layout = Layout()
mid_line = Mid_line()
l_paddle = Paddle((-360, 0))
r_paddle = Paddle((360, 0))
scoreboard = Scoreboard()
ball = Ball()
music = Music()

winning_score = screen.numinput(title="Winning Score", prompt="Set the Winning Score(Default: 5): ", default=5)

win_score = int(winning_score)

screen.listen()

screen.onkey(r_paddle.go_up, "Up")
screen.onkey(r_paddle.go_down, "Down")

game_is_On = True
is_paused = False

def pause():
    global is_paused
    if is_paused:
        is_paused = False
    else:
        is_paused = True

screen.onkey(pause, "space")

while game_is_On:
    if not is_paused:
        time.sleep(0.06)
        screen.update()
        ball.move()

        if l_paddle.ycor() < ball.ycor() and abs(l_paddle.ycor() - ball.ycor()) > 10:
            l_paddle.go_up()
        elif l_paddle.ycor() > ball.ycor() and abs(l_paddle.ycor() - ball.ycor()) > 10:
            l_paddle.go_down()

        # collision with upper and lower Wall
        if ball.ycor() > 250 or ball.ycor() < -250:
            ball.y_bounce()

        # Collision with Paddle
        if ball.distance(l_paddle) < 50 or ball.xcor() < -350 or ball.distance(r_paddle) < 50 or ball.xcor() > 350:
            ball.increase_speed()
            ball.x_bounce()
            music.play_Paddle_Collision()

        # Increasing Score
        if ball.xcor() > 350:
            music.play_Wall_Collision()
            ball.reset_position()
            scoreboard.l_points()

        if ball.xcor() < -350:
            music.play_Wall_Collision()
            ball.reset_position()
            scoreboard.r_points()

        # Game Over and Final Result
        if scoreboard.l_score >= win_score:
            is_paused = False
            music.play_Winning_Sound()
            Scoreboard.game_over(ball,scoreboard.player1, scoreboard.l_score)
        if scoreboard.r_score >= win_score:
            is_paused = False
            music.play_Winning_Sound()
            Scoreboard.game_over(ball,scoreboard.player2, scoreboard.r_score)
    else:
        screen.update()


class Test:
	class inner:
		def f1(self):
			print('f1')
	def f2(self):
			print('f2')

t=Test()
o=t.inner()
o.f1()
t.f2()
#or--------second method
#a=Test.inner()
#a.f1()


import pickle
from copy import copy

import numpy as np
import pytest

import nengo
from nengo import spa
from nengo.exceptions import NetworkContextError, NotAddedToNetworkWarning
from nengo.params import IntParam, NdarrayParam, iter_params
from nengo.utils.numpy import is_array_like
from nengo.utils.progress import TerminalProgressBar

def assert_is_copy(cp, original):
    assert cp is not original  # ensures separate parameters
    for param in iter_params(cp):
        param_inst = getattr(cp, param)
        if isinstance(param_inst, (nengo.solvers.Solver, nengo.base.NengoObject)):
            assert param_inst is getattr(original, param)
        elif is_array_like(param_inst):
            assert np.all(param_inst == getattr(original, param))
        else:
            assert param_inst == getattr(original, param)

def assert_is_deepcopy(cp, original):
    assert cp is not original  # ensures separate parameters
    for param in iter_params(cp):
        param_inst = getattr(cp, param)
        if isinstance(param_inst, (nengo.solvers.Solver, nengo.base.NengoObject)):
            assert_is_copy(param_inst, getattr(original, param))
        elif is_array_like(param_inst):
            assert np.all(param_inst == getattr(original, param))
        else:
            assert param_inst == getattr(original, param)

def make_ensemble():
    with nengo.Network():
        e = nengo.Ensemble(10, 1, radius=2.0)
    return e

def make_probe():
    with nengo.Network():
        e = nengo.Ensemble(10, 1)
        p = nengo.Probe(e, synapse=0.01)
    return p

def make_node():
    with nengo.Network():
        n = nengo.Node(np.min, size_in=2, size_out=2)
    return n

def make_connection():
    with nengo.Network():
        e1 = nengo.Ensemble(10, 1)
        e2 = nengo.Ensemble(10, 1)
        c = nengo.Connection(e1, e2, transform=2.0)
    return c

def make_function_connection():
    with nengo.Network():
        e1 = nengo.Ensemble(10, 1)
        e2 = nengo.Ensemble(10, 1)
        c = nengo.Connection(e1, e2, function=lambda x: x**2)
    return c

def make_learning_connection():
    with nengo.Network():
        e1 = nengo.Ensemble(10, 1)
        e2 = nengo.Ensemble(11, 1)
        c = nengo.Connection(e1.neurons, e2.neurons, transform=np.ones((11, 10)))
        c.learning_rule_type = nengo.PES()
        nengo.Connection(e2, c.learning_rule)
    return c

def make_network():
    with nengo.Network() as model:
        e1 = nengo.Ensemble(10, 1)
        e2 = nengo.Ensemble(10, 1)
        nengo.Connection(e1, e2, transform=2.0)
        nengo.Probe(e2)
    return model

def test_neurons_reference_copy():
    original = make_ensemble()
    cp = original.copy(add_to_container=False)
    assert original.neurons.ensemble is original
    assert cp.neurons.ensemble is cp

def test_learningrule_reference_copy():
    original = make_learning_connection()
    cp = original.copy(add_to_container=False)
    assert original.learning_rule.connection is original
    assert cp.learning_rule.connection is cp

def test_copy_in_network_default_add():
    original = make_network()

    with nengo.Network() as model:
        cp = original.copy()
    assert cp in model.all_objects

    assert_is_deepcopy(cp, original)

def test_copy_outside_network_default_add():
    original = make_network()
    cp = original.copy()
    assert_is_deepcopy(cp, original)

def test_network_copies_defaults():
    original = nengo.Network()
    original.config[nengo.Ensemble].radius = 1.5
    original.config[nengo.Connection].synapse = nengo.Lowpass(0.1)

    cp = original.copy()
    assert cp.config[nengo.Ensemble].radius == 1.5
    assert cp.config[nengo.Connection].synapse == nengo.Lowpass(0.1)

def test_network_copy_allows_independent_manipulation():
    with nengo.Network() as original:
        nengo.Ensemble(10, 1)
    original.config[nengo.Ensemble].radius = 2.0

    cp = original.copy()
    with cp:
        e2 = nengo.Ensemble(10, 1)
    cp.config[nengo.Ensemble].radius = 1.0

    assert e2 not in original.ensembles
    assert original.config[nengo.Ensemble].radius == 2.0

def test_copies_structure():
    with nengo.Network() as original:
        e1 = nengo.Ensemble(10, 1)
        e2 = nengo.Ensemble(10, 1)
        nengo.Connection(e1, e2)
        nengo.Probe(e2)

    cp = original.copy()

    assert cp.connections[0].pre is not e1
    assert cp.connections[0].post is not e2
    assert cp.connections[0].pre in cp.ensembles
    assert cp.connections[0].post in cp.ensembles

    assert cp.probes[0].target is not e2
    assert cp.probes[0].target in cp.ensembles

def test_network_copy_builds(Simulator):
    with Simulator(make_network().copy()):
        pass

def test_copy_obj_view():
    with nengo.Network():
        ens = nengo.Ensemble(10, 4)
        original = ens[:2]

    cp = original.copy()

    assert cp is not original
    assert cp.obj is ens
    assert cp.slice == original.slice

def test_copy_obj_view_in_connection():
    with nengo.Network() as original:
        node = nengo.Node([0.1, 0.2])
        ens = nengo.Ensemble(10, 2)
        nengo.Connection(node[0], ens[1])
        nengo.Connection(node[1], ens[0])

    cp = original.copy()
    assert cp.nodes[0] is not node
    assert cp.ensembles[0] is not ens
    assert cp.connections[0].pre.obj is cp.nodes[0]
    assert cp.connections[1].pre.obj is cp.nodes[0]
    assert cp.connections[0].post.obj is cp.ensembles[0]
    assert cp.connections[1].post.obj is cp.ensembles[0]

def test_pickle_obj_view():
    with nengo.Network():
        ens = nengo.Ensemble(10, 4)
        original = ens[:2]

    cp = pickle.loads(pickle.dumps(original))

    assert cp is not original
    assert cp.obj is not ens
    assert cp.obj.n_neurons == 10
    assert cp.obj.dimensions == 4
    assert cp.slice == original.slice

def test_pickle_obj_view_in_connection():
    with nengo.Network() as original:
        node = nengo.Node([0.1, 0.2])
        ens = nengo.Ensemble(10, 2)
        nengo.Connection(node[0], ens[1])
        nengo.Connection(node[1], ens[0])

    cp = pickle.loads(pickle.dumps(original))
    assert cp.nodes[0] is not node
    assert cp.ensembles[0] is not ens
    assert cp.connections[0].pre.obj is cp.nodes[0]
    assert cp.connections[1].pre.obj is cp.nodes[0]
    assert cp.connections[0].post.obj is cp.ensembles[0]
    assert cp.connections[1].post.obj is cp.ensembles[0]

@pytest.mark.parametrize(
    ("make_f", "assert_f"),
    [
        (make_ensemble, assert_is_copy),
        (make_probe, assert_is_copy),
        (make_node, assert_is_copy),
        (make_connection, assert_is_copy),
        (make_function_connection, assert_is_copy),
        (make_learning_connection, assert_is_copy),
        (make_network, assert_is_copy),
    ],
)
class TestCopy:
    """A basic set of tests that should pass for all objects."""

    def test_copy_in_network(self, make_f, assert_f):
        original = make_f()

        with nengo.Network() as model:
            cp = original.copy(add_to_container=True)
        assert cp in model.all_objects

        assert_f(cp, original)

    def test_copy_in_network_without_adding(self, make_f, assert_f):
        original = make_f()

        with nengo.Network() as model:
            cp = original.copy(add_to_container=False)
        assert cp not in model.all_objects

        assert_f(cp, original)

    def test_copy_outside_network(self, make_f, assert_f):
        original = make_f()
        with pytest.raises(NetworkContextError, match="either be created inside a"):
            original.copy(add_to_container=True)

    def test_copy_outside_network_without_adding(self, make_f, assert_f):
        original = make_f()
        cp = original.copy(add_to_container=False)
        assert_f(cp, original)

    def test_python_copy_warns_abt_adding_to_network(self, make_f, assert_f):
        original = make_f()
        copy(original)  # Fine because not in a network
        with nengo.Network():
            with pytest.warns(NotAddedToNetworkWarning):
                copy(original)

@pytest.mark.parametrize(
    "make_f", (make_ensemble, make_probe, make_node, make_connection, make_network)
)
class TestPickle:
    """A basic set of tests that should pass for all objects."""

    def test_pickle_roundtrip(self, make_f):
        """Tests making a deepcopy using pickling and unpickling."""
        original = make_f()
        cp = pickle.loads(pickle.dumps(original))
        assert_is_deepcopy(cp, original)

    def test_unpickling_warning_in_network(self, make_f):
        """Tests a not added to network warning when unpickling inside a network when
        the original was not in the network."""
        original = make_f()
        pkl = pickle.dumps(original)
        with nengo.Network():
            with pytest.warns(NotAddedToNetworkWarning):
                pickle.loads(pkl)

@pytest.mark.parametrize(
    "original",
    [
        nengo.learning_rules.PES(),
        nengo.learning_rules.BCM(),
        nengo.learning_rules.Oja(),
        nengo.learning_rules.Voja(),
        nengo.processes.WhiteNoise(),
        nengo.processes.FilteredNoise(),
        nengo.processes.BrownNoise(),
        nengo.processes.PresentInput([0.1, 0.2], 1.0),
        nengo.synapses.LinearFilter([0.1, 0.2], [0.3, 0.4], True),
        nengo.synapses.Lowpass(0.005),
        nengo.synapses.Alpha(0.005),
        nengo.synapses.Triangle(0.005),
    ],
)
class TestFrozenObjectCopies:
    def test_copy(self, original):
        assert_is_copy(copy(original), original)

    def test_pickle_roundtrip(self, original):
        assert_is_deepcopy(pickle.loads(pickle.dumps(original)), original)

def test_copy_spa(Simulator):
    with spa.SPA() as original:
        original.state = spa.State(16)
        original.cortex = spa.Cortical(spa.Actions("state = A"))

    cp = original.copy()

    # Check that it still builds.
    with Simulator(cp):
        pass

    # check vocab instance param is set
    for node in cp.all_nodes:
        if node.label in ["input", "output"]:
            assert cp.config[node].vocab is not None

def test_copy_instance_params():
    with nengo.Network() as orig_net:
        orig_net.config[nengo.Ensemble].set_param(
            "test", IntParam("test", optional=True)
        )
        orig_net.config[nengo.Ensemble].set_param(
            "test2", NdarrayParam("test2", optional=True)
        )

        orig_ens = nengo.Ensemble(10, 1)
        orig_net.config[orig_ens].test = 42
        orig_net.config[orig_ens].test2 = np.array([49])

    # test copy function
    copy_net = orig_net.copy()
    copy_ens = copy_net.ensembles[0]
    assert copy_net.config[copy_ens].test == 42
    assert np.array_equal(copy_net.config[copy_ens].test2, np.array([49]))

    # test copy via pickle
    copy_net = pickle.loads(pickle.dumps(orig_net))
    copy_ens = copy_net.ensembles[0]
    assert copy_net.config[copy_ens].test == 42
    assert np.array_equal(copy_net.config[copy_ens].test2, np.array([49]))

def test_pickle_model(Simulator, seed, allclose):
    t_run = 0.5
    simseed = seed + 1

    with nengo.Network(seed=seed) as network:
        u = nengo.Node(nengo.processes.WhiteSignal(t_run, 5))
        a = nengo.Ensemble(100, 1)
        b = nengo.Ensemble(100, 1)
        nengo.Connection(u, a, synapse=None)
        nengo.Connection(a, b, function=np.square)
        up = nengo.Probe(u, synapse=0.01)
        ap = nengo.Probe(a, synapse=0.01)
        bp = nengo.Probe(b, synapse=nengo.Alpha(0.01))

    with Simulator(network, seed=simseed) as sim:
        sim.run(t_run)
        t0, u0, a0, b0 = sim.trange(), sim.data[up], sim.data[ap], sim.data[bp]
        pkls = pickle.dumps({"model": sim.model, "up": up, "ap": ap, "bp": bp})

    # reload model
    del network, sim, up, ap, bp
    pkl = pickle.loads(pkls)
    up, ap, bp = pkl["up"], pkl["ap"], pkl["bp"]

    with Simulator(None, model=pkl["model"], seed=simseed) as sim:
        sim.run(t_run)
        t1, u1, a1, b1 = sim.trange(), sim.data[up], sim.data[ap], sim.data[bp]

    tols = {"atol": 1e-5}
    assert allclose(t1, t0, **tols)
    assert allclose(u1, u0, **tols)
    assert allclose(a1, a0, **tols)
    assert allclose(b1, b0, **tols)

def test_copy_convolution():
    x = nengo.Convolution(2, (2, 3, 4), channels_last=False, groups=2)
    y = copy(x)

    assert x.n_filters == y.n_filters
    assert x.input_shape == y.input_shape
    assert x.channels_last == y.channels_last
    assert x.groups == y.groups

@pytest.mark.parametrize(
    "optimize, dt, progress_bar", [(True, 0.001, False), (False, 0.002, True)]
)
def test_pickle_sim(Simulator, seed, allclose, optimize, dt, progress_bar):
    trun0 = 0.5
    trun1 = 0.5
    simseed = seed + 1
    progress = TerminalProgressBar() if progress_bar else progress_bar

    with nengo.Network(seed=seed) as network:
        u = nengo.Node(nengo.processes.WhiteSignal(trun0 + trun1, high=5))
        a = nengo.Ensemble(100, 1)
        b = nengo.Ensemble(100, 1)
        nengo.Connection(u, a, synapse=None)
        nengo.Connection(a, b, function=np.square)
        up = nengo.Probe(u, synapse=0.01)
        ap = nengo.Probe(a, synapse=0.01)
        bp = nengo.Probe(b, synapse=nengo.Alpha(0.01))

    with Simulator(
        network, seed=simseed, dt=dt, optimize=optimize, progress_bar=progress
    ) as sim:
        sim.run(trun0)
        pkls = pickle.dumps({"sim": sim, "up": up, "ap": ap, "bp": bp})

        sim.run(trun1)
        t0, u0, a0, b0 = sim.trange(), sim.data[up], sim.data[ap], sim.data[bp]

    # reload model
    del network, sim, up, ap, bp
    pkl = pickle.loads(pkls)
    up, ap, bp = pkl["up"], pkl["ap"], pkl["bp"]

    with pkl["sim"] as sim:
        assert sim.seed == simseed
        assert sim.dt == dt
        assert sim.optimize == optimize
        if progress_bar:
            assert isinstance(sim.progress_bar, TerminalProgressBar)
        else:
            assert sim.progress_bar is progress_bar

        sim.run(trun1)
        t1, u1, a1, b1 = sim.trange(), sim.data[up], sim.data[ap], sim.data[bp]

    tols = {"atol": 1e-5}
    assert allclose(t1, t0, **tols)
    assert allclose(u1, u0, **tols)
    assert allclose(a1, a0, **tols)
    assert allclose(b1, b0, **tols)

    # check that closed status is preserved
    pkl = pickle.loads(pickle.dumps({"sim": sim}))
    assert pkl["sim"].closed


import logging as log
import os
import errno

from threading import Thread

try:
    from urllib import urlopen
    from urlparse import urlparse
    import unicodecsv as csv
except ImportError:
    from urllib.request import urlopen
    from urllib.parse import urlparse
    import csv

class FolderCreationError(Exception):
    pass

class BrokenImageError(Exception):
    pass

def make_path(path):
    """Ensures all the folders in path exists.
    Raises FolderCreationError if failed to create the required folders.
    """
    try:
        os.makedirs(path)
    except OSError as e:
        if e.errno != errno.EEXIST or not os.path.isdir(path):
            raise FolderCreationError('Failed to create folder <{}>'.format(path))

class Record(object):

    def __init__(self, filename, schema):
        self.filename = filename + '.csv'
        # if there's a folder in the filename make sure it exists
        if (os.path.dirname(self.filename)):
            make_path(os.path.dirname(self.filename))
        self.file = open(self.filename, 'w')
        self.writer = csv.DictWriter(self.file, fieldnames=schema, quoting=csv.QUOTE_ALL, strict=True)
        self.writer.writeheader()
        log.info('Created a new record file at: %s', self.filename)

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

    def add_record(self, data):
        self.writer.writerow(data)

class Album(object):

    def __init__(self, name, descriptions=False):
        make_path(name)
        self.name = name
        self.record = None
        if descriptions:
            self.record = Record(name, ['filename', 'description', 'permalink'])

    def add_image(self, url):
        Thread(target=self._image_dl(url), args=(url)).start()

    def _image_dl(self, url):
        try:
            img_bin = urlopen(url).read()
            if img_bin:
                filename = urlparse(url).path.split('/')[-1]
                fullpath = os.path.join(self.name, filename)
                with open(fullpath, 'wb') as f:
                    f.write(img_bin)
            return
        except IOError:
            # there was some issue with the connection
            # raise the Broken Image Error
            pass
        raise BrokenImageError

    def add_description(self, imgurl, desc, perma):
        if self.record:
            filename = urlparse(imgurl).path.split('/')[-1]
            self.record.add_record({
                'filename': filename,
                'description': desc,
                'permalink': perma,
            })


#!/usr/bin/env python3

'''
    【简介】
	PyQT5中 QTreeWidget 例子

'''

from PyQt5.QtWidgets import *
import sys

class TreeWidgetDemo(QMainWindow):
	def __init__(self,parent=None):
		super(TreeWidgetDemo,self).__init__(parent)
		self.setWindowTitle('TreeWidget 例子')
		self.tree = QTreeWidget()
        # 设置列数
		self.tree.setColumnCount(2)
        # 设置头的标题
		self.tree.setHeaderLabels(['Key','Value'])
		root= QTreeWidgetItem(self.tree)
		root.setText(0,'root')
		root.setText(1,'0')

		child1 = QTreeWidgetItem(root)
		child1.setText(0,'child1')
		child1.setText(1,'1')

		child2 = QTreeWidgetItem(root)
		child2.setText(0,'child2')
		child2.setText(1,'2')

		child3 = QTreeWidgetItem(root)
		child3.setText(0,'child3')
		child3.setText(1,'3')

		child4 = QTreeWidgetItem(child3)
		child4.setText(0,'child4')
		child4.setText(1,'4')

		child5 = QTreeWidgetItem(child3)
		child5.setText(0,'child5')
		child5.setText(1,'5')

		self.tree.addTopLevelItem(root)
		self.tree.clicked.connect( self.onTreeClicked )

		self.setCentralWidget(self.tree)

	def onTreeClicked(self, qmodelindex):
		item = self.tree.currentItem()
		print("key=%s ,value=%s" % (item.text(0), item.text(1)))

if __name__ == '__main__':
	app = QApplication(sys.argv)
	tree = TreeWidgetDemo()
	tree.show()
	sys.exit(app.exec_())