#! /usr/bin/env python
# 2018
# L. Strolger
# Adapted from code written for FFSN
# by S. Rodney

# 2010
# S.Rodney
#
# Defines a python class for Sextractor catalogs
# and functions to run sextractor on an image
# with an input bad pixel mask, and to clean the
# resulting catalog by rejecting obvious bad sources
# 2012.07.17 master zeropoint dicts kept in tools.hstsnphot
"""
sextract.py [--options] image.fits [sex-options]

Command-line access to run sextractor with HSTSNpipeline defaults
on a single image.  If the image has suffix _sci.fits or _masked.fits
and the corresponding _weight.fits and _badpix.fits files are present,
then they are automatically used.

--verbose X   :  run with verbose level X, 0-10
--checkaper   :  make a checkimage showing apertures
--clobber     :  overwrite existing images

Any options given after the image name are passed directly to sextractor.
e.g.:   sextract.py --checkaper some_image_sci.fits -DETECT_MINAREA 3 -DETECT_THRESH 2
"""

class Catalog(object):
    """ A catalog of astronomical sources """
    def __init__( self, catfile=None, imfile=None ):
        """ Initiate a catalog of sources with
        (x,y) pixel coordinates and/or (ra,dec)
        sky coordinates, as well as measured
        magnitudes and/or fluxes.
        If provided, the 'imfile' is an image file
        from which to extract WCS information for
        connecting (x,y) to (ra,dec) and the
        'catfile' is the name of an ASCII catalog
        file for reading/writing the list of sources.
        """
        self.coldata = {} # dictionary of data, keyed by column name
        self.colindex = {}
        if catfile :
            import os
            catfile = os.path.abspath( catfile )
            self.readcatfile( catfile )
            self.filename = catfile

    @property
    def colnames( self ):
        """ return a list of column names """
        return( [ colname for colname in self.coldata.keys() ] )

    @property
    def nrows( self ):
        if len( self.coldata.values() )>0  :
            return( len( self.coldata.values()[0] ) )
        else :
            return( 0 )

    @property
    def ncols( self ):
        return( len( self.coldata.keys() ) )

    def getrow( self, rownum ) :
        """ extract the given row from the catalog data"""
        row = []
        for colindex in range(self.ncols) :
            row.append( self.data[colindex][rownum] )
        return( row )

    def getrowdict( self, rownum ) :
        """ extract the given row from the catalog data
        as a dictionary keyed by column names"""
        rowdict = {}
        for colname in self.colnames :
            rowdict[colname] = self.data[colname][rownum]
        return( rowdict )

    @property
    def ra( self ):
        for k in ['X_WORLD','RA','ra','R.A.','ALPHA_J2000']:
            if k in self.colnames :
                return( self.coldata[k] )
        return( [] )

    @property
    def dec( self ):
        for k in ['Y_WORLD','DEC','dec','Dec','Decl','DELTA_J2000']:
            if k in self.colnames :
                return( self.coldata[k] )
        return( [] )

    @property
    def x( self ):
        if 'X_IMAGE' in self.colnames :
            return( self.coldata['X_IMAGE'] )
        else :
            return( [] )

    @property
    def y( self ):
        if 'Y_IMAGE' in self.colnames :
            return( self.coldata['Y_IMAGE'] )
        else :
            return( [] )

    # # TODO : UPDATE THIS
    # def appendrow( self, row ):
    #     """ append a single row of data,
    #     The row of data may be given as a
    #     python list or a dict, keyed by
    #     column names.
    #     """
    #     self.data.insert( 1, -1 )
    #     self.updaterow( self.nrows-1, row )

    # # TODO : UPDATE THIS
    # def updaterow( self, index, row ):
    #     """ update a single row of data at
    #     the given row index.
    #     The row of data may be given as a
    #     python list or a dict, keyed by
    #     column names.
    #     """
    #     if type(row) == list :
    #         for colindex in range(len(row)) :
    #             self.data[colindex][index] = row[colindex]
    #     elif type(row) == dict :
    #         for column in row.keys()  :
    #             self.data[column][index] = row[column]

    def delrow( self, index ):
        self.data.delete( index )

    def readcatfile( self, catfile ) :
        """ read a sextractor catalog """
        # TODO : handle different table types
        from numpy import loadtxt, array, genfromtxt
        # read the sex-cat header
        fin = open( catfile, 'r' )
        header = fin.readlines()
        fin.close()
        # TODO : handle catalogs with a single row of data !!
        #catdat = loadtxt( catfile )
        catdat = genfromtxt( catfile )
        self.coldata = {}
        self.colindex = {}
        # guard against single-line catalogs:
        if len(catdat.shape)==1 : catdat = array([ catdat ])
        idfirst=True
        for line in header :
            if not line.startswith('#'): break
            try :
                icol = int( line.split()[1] ) - 1
            except ValueError :
                try:
                    # is this a reverse-formatted header?
                    idfirst = False
                    icol = int( line.split()[-1].strip('()') ) - 1
                except e:
                    continue
            if idfirst : colname =  line.split()[2]
            elif '(' in line :
                colname = line.lstrip('#').split('(')[0].strip().replace(' ','_')
            self.colindex[icol+1] = colname

            if colname in ['NUMBER', 'FLAGS'] :
                coldata = catdat[:,icol].astype( int )
            else :
                coldata = catdat[:,icol].astype( float )
            self.coldata[colname] = coldata
        #self.data = asciidata.open( catfile )

    def writecatfile( self, catfile, colInfo=None, headComment=None, selectrows=[] ) :
        """ write out catalog as a sextractor .cat file.
        selectrows : a sequence of row numbers to select for output.
            If empty, all rows are printed.
        """
        fout = open( catfile, 'w' )
        iclist = self.colindex.keys()
        iclist.sort()
        for icol in iclist :
            print >> fout,"# %3i %15s "%(icol, self.colindex[icol] )
        for irow in range(self.nrows) :
            if len(selectrows) and irow not in selectrows : continue
            outstring = '  '
            for icol in iclist :
                colname = self.colindex[icol]
                datum = self.coldata[colname][irow]
                if datum % 1 :
                    if abs(datum)>1e-4 :
                        outstring += ' %14.7f'%datum
                    else :
                        outstring += ' %14.7e'%datum
                else :
                    outstring += ' %-14i'%datum
            print >> fout, outstring
        fout.close()
        return( catfile )

    def wds9reg( self, regfile=None, color='red', circlesize=10, width=1,
                 xcol='X_IMAGE', ycol='Y_IMAGE', textcol=None, criteria=None):
        """ write sources to a ds9 region file.
        circlesize may be a scalar or an array with one size
        for each object in the catalog. Criteria can now be used to only write certian regions
        to the file from the catelog. criteria is a tuple with values (colname,comparison_operator,value).
        Eg.     criteria = ('EXT_IMAGE','==',1)
        """
        import numpy as np
        import os
        #import exceptions

        if( xcol not in self.colnames or ycol not in self.colnames ):
            print("ERROR: cannot make a ds9 region file without %s,%s"%(xcol,ycol))
            return( -1 )

        xcolname = xcol.lower()
        convertRADEC=False
        if ('alpha' in xcolname) or ('ra' in xcolname) or ('r.a.' in xcolname) or ('world' in xcolname) :
            if isinstance( self.coldata[xcol][0], float) :
                try :
                    from astLib import astCoords
                    print( "Converting RA,DEC from decimal degrees to hh:mm:ss,dd:mm:ss")
                    convertRADEC=True
                except :
                    print( "Can't load astLib for RA,DEC conversion!" )
                    return(-1)

        if not regfile :
            regfile = os.path.splitext( self.filename )[0] +'.reg'

        fout = open(regfile,'w')
        print>>fout,"# Region file format: DS9 version 5.4"
        print>>fout,"""global color=%s font="times 16 bold" select=1 edit=1 move=0 delete=1 include=1 fixed=0 width=%i"""%(color,width)
        for irow in range(self.nrows):
            if np.iterable(circlesize): cs = circlesize[irow]
            else : cs = circlesize

            textstr = ''
            if textcol :
                # textcol may be a comma-sep'd list of col names.
                textvals = []
                for tc in textcol.split(',') :
                    tval = self.coldata[tc][irow]
                    if type(tval)in [int,np.int,np.int16,np.int32,np.int64,np.int8,np.int_] :
                        textvals.append('%i'%tval)
                    elif type(tval) in [float,np.float,np.float32,np.float64,np.float128,np.float_]  :
                        textvals.append('%.2f'%tval)
                    else : textvals.append(str(tval))
                textstr = '# text={%s}'%':'.join( textvals )

            #Only write the line to the file if it passes the criteria passed to this function.
            #This allows us to separate extensions for makeing flt region files in the fakes
            if criteria is None: write_line = True
            else:write_line = eval("self.coldata['%s'][irow] %s %s"%criteria)

            if write_line:
                cstr = '%.1f'%cs
                x, y = self.coldata[xcol][irow], self.coldata[ycol][irow]
                if isinstance( x, float) :
                    xstr, ystr = '%.8f'%x, '%.8f'%y
                    if convertRADEC :
                        xstr = astCoords.decimal2hms( x, delimiter=':' )
                        ystr = astCoords.decimal2dms( y, delimiter=':' )
                        if cs > 1 : cs = min( 0.75, cs*0.06 )
                        cstr = '%.3f"'%( cs )
                print>>fout,'circle(%s,%s,%s)%s'%( xstr, ystr, cstr, textstr )

        fout.close()
        return( regfile )

    def wxycat( self, catfile ):
        """ write sources to a simple x,y catalog file.
        (useful for use with the Vista marktv command)

        NEEDS UPDATING !!!
        """
        from numpy import iterable

        XCOL,YCOL = 'X_IMAGE','Y_IMAGE'
        if XCOL in self.colnames :
            xcol = XCOL
        else :
            for colname in self.colnames:
                if colname.find(XCOL)>=0 :
                    xcol = colname
                    break
        if YCOL in self.colnames :
            ycol = YCOL
        else :
            for colname in self.colnames:
                if colname.find(YCOL)>=0 :
                    ycol = colname
                    break
        fout = open(catfile,'w')
        for i in range(self.nrows):
            print>>fout,"%12.1f %12.1f"%(
                self.coldata[xcol][i],
                self.coldata[ycol][i] )

def wht2var( infile, outfile=None, platescale='default',
             verbose=False ):
    """
    convert a _wht.fits weight map from multi-drizzle
    into a variance map suitable for use as a
    sextractor weight map.
    following Casertano et al 2006.
    """
    #import astropy.io.fits as pyfits
    from astropy.io import fits as pyfits
    from numpy import nan_to_num

    # size of the input pixels in arcsec
    # default is to read it from the config file
    if platescale=='default': pscalein = PLATESCALE

    # get size of the output pixel in arcsec
    # and pixfrac in units of input pixels from header
    try :
        pscaleout = 3600*(
            abs(pyfits.getval(infile,'CD1_1')) +
            abs(pyfits.getval(infile,'CD2_2')))/2.
    except KeyError :
        pscaleout = 3600*(
            abs(pyfits.getval(infile,'CD1_1',ext=1)) +
            abs(pyfits.getval(infile,'CD2_2',ext=1)))/2.
    try :
        pixfrac = pyfits.getval(infile,'PIXFRAC')
    except KeyError:
        try: pixfrac = pyfits.getval(infile,'D001PIXF')
        except KeyError:
            # TODO : FIX THIS WHEN MAKING WHT IMAGES!!!!
            pixfrac=1

    l = pscaleout / pscalein
    p = pixfrac
    if l<=p : Fa = (l/p)**2 * (1-l/(3*p))**2
    else : Fa = (1 - p/(3*l))**2

    # open file, check for multiple extensions
    inim = pyfits.open( infile )
    extnames = [ext.name for ext in inim]
    whtext=0
    if len(inim)>1:
        if 'WHT' in extnames : whtext = extnames.index('WHT')
        else : whtext=2

    # make the variance map
    vardata = nan_to_num(  Fa / inim[whtext].data )
    varhdu = pyfits.ImageHDU(
        data=vardata, header=inim[whtext].header, name='VAR' )
    varhdu.header['EXTNAME']='VAR'

    # write it to a new extension or a new file
    if outfile==None or outfile==infile :
        if 'VAR'  in extnames : inim['VAR'] = varhdu
        else : inim.insert( whtext+1, varhdu )
        if verbose : print("adding VAR ext to %s"%infile)
        inim.writeto( infile, clobber=True )
    else :
        if verbose : print("writing var map to %s"%outfile)
        pyfits.writeto( outfile, vardata, clobber=True )

def runsex(imfile, catfile=None, weightfile=None, flagfile=None,
           paramstring='', detimfile=None, returncat=False, verbose=False, clobber=False):
    '''
    Run sextractor on a single image.
    if returncat==True : Returns a Catalog object.
    else : returns a string with the catalog filename
    specify detimfile to run in double-image mode.
    '''
    import shutil
    import os
    import subprocess
    #import exceptions
    #import astropy.io.fits as pyfits
    from astropy.io import fits as pyfits
    ## from tools import hstsnphot

    # locate the sextractor binary and the config file dir
    try : sexbindir = os.environ['SNPATH']
    except KeyError : sexbindir = ''
    try :
        # if we got here via startpipe.py :
        sexconfigdir = os.path.join(os.environ['SNSOURCE'],'Sex')
    except KeyError :
        # if we're working from a command-line call using the pipeline installation:
        sexconfigdir = os.path.join( os.path.dirname(os.path.abspath(os.path.realpath(__file__))),'sexconfig')
    if not os.path.isdir( sexconfigdir ) :

        # otherwise, fingers crossed...
        sexconfigdir = ''

    if verbose>2 :
        print( 'using %s with config files from %s'%(
                os.path.join(sexbindir,'sex'), sexconfigdir ) )

    imfile = os.path.abspath( imfile )
    if catfile==None :
        imroot = os.path.splitext( imfile )[0]
        catfile = imroot + ".sexcat"

    if os.path.isfile( catfile ) and not clobber :
        print("%s exists. Not clobbering."%catfile )
        if returncat :
            sexcat = Catalog( imfile=imfile, catfile=catfile )
            return( sexcat )
        else :
            return( catfile )

    # Using some generic cofiguration files, good for ground work.
    configfile = os.path.join( sexconfigdir,"default.sex")
    paramfile = os.path.join( sexconfigdir,"daofind.param")
    convfile = os.path.join( sexconfigdir,"default.conv")
    nnwfile = os.path.join( sexconfigdir,"default.nnw")

    # run sextractor
    sexcommand = [ os.path.join( sexbindir, "sex" ) ]
    if detimfile :
        sexcommand.append( "%s,%s"%(detimfile,imfile) )
    else :
        sexcommand.append("%s"%imfile)
    sexcommand.append("-c %s"%configfile)
    sexcommand.append("-FILTER_NAME %s"%convfile)
    sexcommand.append("-STARNNW_NAME %s"%nnwfile)
    sexcommand.append("-CATALOG_NAME %s"%catfile)
    sexcommand.append("-CATALOG_TYPE ASCII_HEAD")

    if weightfile :
        sexcommand.append("-WEIGHT_IMAGE %s"%weightfile)
        sexcommand.append("-WEIGHT_TYPE MAP_WEIGHT")
    else :
        sexcommand.append("-WEIGHT_TYPE NONE")

    if flagfile :
        paramfile = os.path.join( sexconfigdir,"daofind_flags.param")
        sexcommand.append("-FLAG_IMAGE %s"%flagfile)
    else :
        sexcommand.append("-FLAG_IMAGE NONE")
    sexcommand.append("-PARAMETERS_NAME %s"%paramfile)

    if verbose>3:
        sexcommand.append("-VERBOSE_TYPE FULL")
    elif verbose>1:
        sexcommand.append("-VERBOSE_TYPE NORMAL")
    else :
        sexcommand.append("-VERBOSE_TYPE QUIET")
    sexcommand.append(paramstring)

    if verbose : print( ' '.join(sexcommand) )
    subprocess.call( ' '.join(sexcommand), shell=True )

    if not os.path.isfile( catfile ) :
        raise RuntimeError(
            "sextractor failed to produce expected catalog file %s!!"%catfile)

    # # read in the catalog
    if returncat :
        sexcat = Catalog( imfile=imfile, catfile=catfile )
        return( sexcat )
    else :
        return( catfile )

def cleancat( catfile, newcatfile=None,
              fwhm_min=1.8, flag_max=1, Bmin=1,
              isogood_min=0.3,
              verbose=False, clobber=False ):
    """
    read in the given catfile,
    clean it of obvious cosmic rays and bad sources
    write it out to newcatfile and return the
    name of the newcatfile.
    """
    # TODO : this is painfully slow
    # should be far faster with numpy
    import os
    cat = Catalog( catfile=catfile )
    newcat = Catalog( )

    if not newcatfile :
        newcatfile = os.path.splitext( catfile )[0] +"_clean.cat"

    if os.path.isfile( newcatfile ) and not clobber:
        print("%s already exists. Not clobbering."%newcatfile)
        return( newcatfile )

    newcat.filename = newcatfile
    for colname in cat.colnames :
        newcat.coldata[colname] = []
    for icol in cat.colindex :
        newcat.colindex[icol] = cat.colindex[icol]
    for irow in range(cat.nrows) :
        if fwhm_min :
            if cat.coldata['FWHM_IMAGE'][irow] < fwhm_min : continue
        if flag_max :
            if cat.coldata['FLAGS'][irow] > flag_max : continue
        if Bmin :
            if cat.coldata['B_IMAGE'][irow] < Bmin : continue

        if 'NIMAFLAGS_ISO' in cat.colnames and 'ISOAREA_IMAGE' in cat.colnames :
            isogood=float(cat.coldata['NIMAFLAGS_ISO'][irow]) / cat.coldata['ISOAREA_IMAGE'][irow]
            if isogood < isogood_min :  continue

        for colname in newcat.colnames :
            newcat.coldata[colname].append( cat.coldata[colname][irow] )
    newcat.writecatfile( newcatfile )
    return( newcatfile )

def astromatch(cat1, cat2, radius, coords='sky'):
    '''
    Simple astrometric catalog matching between two catalog objects.
    Returns an Nx4 array with the N sources in cat1 that have matches
    within 'radius' arcesconds in cat2.
    The coords keyword specifies the coordinate system used for
    the return array :
       coords='sky'  :   [ra1,dec1,ra2,dec2]
       coords='xy'   :   [x1, y1, x2, y2]
    '''
    import numpy as np
    if type( cat1 ) == str :
        cat1 = Catalog( cat1 )
    if type( cat2 ) == str :
        cat2 = Catalog( cat2 )
    ra1, dec1, x1, y1  = cat1.ra,cat1.dec,cat1.x,cat1.y
    ra2, dec2, x2, y2  = cat2.ra,cat2.dec,cat2.x,cat2.y

    outlist = []
    for i in range(0,cat1.nrows):
        ra_scale = np.cos( dec1[i]*np.pi/180.)
        delta_ra = 3600.*(ra1[i] - ra2)*ra_scale
        delta_dec = 3600.*(dec1[i] - dec2)
        dist = np.sqrt(delta_ra**2+delta_dec**2)
        dmin = np.amin(dist)
        jmin = np.where(dist==dmin)
        if dmin <= radius:
            if coords=='sky' :
                outlist.append( [ra1[i],dec1[i],ra2[jmin],dec2[jmin]] )
            elif coords=='xy' :
                outlist.append( [x1[i],y1[i],x2[jmin],y2[jmin]] )
    return np.array( outlist )

if __name__ == '__main__':
    import sys
    import getopt
    import os

    verbose=1
    checkaper=False
    clobber=False

    # read in arguments and options
    try:
        opt,arg = getopt.getopt(
            sys.argv[1:],"v,h",
            longopts=["verbose=","quiet","help","checkaper","clobber" ] )
    except getopt.GetoptError:
        print ("Error : incorrect option or missing argument.")
        print (__doc__)
        sys.exit(1)
    for o, a in opt:
        if o in ["-h", "--help"]:
            print (__doc__)
            sys.exit(0)
        elif o == "-v" :
            verbose = 2
        elif o == "--quiet" :
            verbose = False
        elif o == "--verbose" :
            try: verbose = int(a)
            except ValueError :
                print("Error : must specify verbosity level as an integer 1-10.")
                sys.exit(0)
        elif o == '--checkaper' :
            checkaper=True
        elif o == '--clobber' :
            clobber=True

    # first arg is the image filename
    imfile = arg[0]

    ## # Check for _weight.fits and _badpix.fits images
    ## sfx = '.fits'
    ## if imfile.endswith('_sci.fits') : sfx = '_sci.fits'
    ## elif imfile.endswith('_masked.fits') : sfx = '_masked.fits'
    ## whtfile = imfile.replace( sfx, '_weight.fits')
    ## if not os.path.isfile( whtfile ) :
    ##     whtfile = imfile.replace( sfx, '_wht.fits')
    ##     if not os.path.isfile( whtfile ) :
    ##         whtfile = None
    ## bpxfile = imfile.replace( sfx, '_badpix.fits')
    ## if not os.path.isfile( bpxfile ) :
    ##     bpxfile = imfile.replace( sfx, '_bpx.fits')
    ##     if not os.path.isfile( bpxfile ) :
    ##         bpxfile = None

    whtfile = None
    bpxfile = None

    # any additional args get passed straight to sextractor
    paramstring = ''
    if len(arg) > 1 :
        paramstring = ' '.join(arg[1:])

    if checkaper :
        paramstring+=' -CHECKIMAGE_TYPE APERTURES -CHECKIMAGE_NAME %s '%imfile.replace('.fits','_aper.fits')

    runsex(imfile, catfile=None, weightfile=whtfile, flagfile=bpxfile,
           paramstring=paramstring, detimfile=None, verbose=verbose, clobber=clobber)


from flask import Flask, render_template, request
# from sklearn.feature_extraction.text import TfidfVectorizer
# from sklearn.linear_model import PassiveAggressiveClassifier
import pickle
import pandas as pd
import re

app = Flask(__name__)

# making cleaning function
def find_the (text):
    while len(re.findall('##' , text)) != 0:
        text = re.sub('##' , '#' ,text)
    while len(re.findall('@@' , text)) != 0:
        text = re.sub('@@' , '@' ,text)

    no_h = len(re.findall('#',text))
    #make sure the @ is changed for at
    no_a = len(re.findall('@',text))

    text = re.sub('@(\s|\.|\?|,|;|:|!|\(|\))', 'at' , text)
    text = re.sub('\s' , ' ' , text)
    text = re.sub('\S@' , ' @' , text)
    text = re.sub('\S#' , ' #' , text)

    #while len(re.findall('\s\s' , text)) != 0:
    #    text = re.sub('\s\s' , '\s' ,text)

    #none_repeat = text
    text_a = text.split()
    return text

def removing_url(text):
#replace URL of a text
    cltext = re.sub('http[^\s]+','',text)
    return cltext

def text_processing(text):
    review = re.sub('[^a-zA-Z]', ' ', text)
    review = review.lower()
    review = review.split()
    review = [ps.stem(word) for word in review if not word in stop_words]
    review = ' '.join(review)
    return review

def text_cleaning(text):
  text = find_the(text)
  text = removing_url(text)
  text = text_processing(text)

  return text

# load the model from disk
model = pickle.load(open('svc_model.pkl', 'rb'))
vector = pickle.load(open('vectorizer.pkl', 'rb'))
stop_words = pickle.load(open('stopwords.pkl', 'rb'))
ps = pickle.load(open('stem.pkl', 'rb'))

class_labels = {0:'Negative',
                1: 'Positive'}

def sentiment_analysis(message):
    input_data = [message]
    vectorized_input_data = vector.transform(input_data)
    prediction = model.predict(vectorized_input_data)
    prediction = class_labels[prediction[0]]
    return prediction

@app.route('/')
def home():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    if request.method == 'POST':
        message = request.form['message']
        pred = sentiment_analysis(message)
        print(pred)
        return render_template('index.html', prediction=pred)
    else:
        return render_template('index.html', prediction="Something went wrong")

if __name__ == '__main__':
    app.run(debug=True)


import matplotlib.pyplot as plt
import csv

with open('data_out.csv', 'r', newline='') as csvfile:
    reader = csv.reader(csvfile, delimiter=' ', quotechar='|', quoting=csv.QUOTE_MINIMAL)
    right_motor = next(reader)
    left_sensor = next(reader)
    right_sensor = next(reader)
    left_motor = next(reader)

# print(left_motor)
float_right_motor = []
for i in range(len(right_motor)):
    float_right_motor.append(float(right_motor[i]))
float_left_motor = []
for i in range(len(left_motor)):
    float_left_motor.append(float(left_motor[i]))

float_right_sensor = []
for i in range(len(right_sensor)):
    float_right_sensor.append(float(right_sensor[i]))
float_left_sensor = []
for i in range(len(left_sensor)):
    float_left_sensor.append(float(left_sensor[i]))

t = [i for i in range(len(left_sensor))]
plt.plot(t, float_left_sensor)
plt.plot(t, float_right_sensor)
plt.plot(t, float_left_motor)
plt.plot(t, float_right_motor)

plt.legend(["left sensor", "right sensor", "left motor", "right motor"])
plt.title("Motor and Sensor Values")
plt.xlabel("Data Points")
plt.ylabel("Motor Speeds and Sensor Values")

plt.show()


from 声音模组 import *
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['Microsoft JhengHei']
plt.rcParams['axes.unicode_minus'] = False

声音1 = 开启wav档('钢琴声音.wav')

声音2 = 声音1.reverse()

阵列1 = 声音转阵列(声音1)
阵列2 = 声音转阵列(声音2)

plt.subplot(1,2,1)
plt.plot(阵列1)
plt.title('原本音讯', fontsize=18)

plt.subplot(1,2,2)
plt.plot(阵列2)
plt.title('反转音讯', fontsize=18)

plt.show()

課程資料/陣列與影音數位化/反轉說明.py


from django.conf.urls import url, include
#from rest_framework.routers import DefaultRouter
from . import views

# in case you want a browsable API (sort of)
#router = DefaultRouter()
#router.register('setup', views.SetUpGameView, base_name='setup')
#router.register('next', views.NextTurnView, base_name='next')
#router.register('draw', views.DrawDieView, base_name='draw')
#router.register('pick', views.PickDestView, base_name='pick')

urlpatterns = [
#    url(r'^',
#        include(router.urls)),
        url(r'^setup/$',
            views.SetUpGameView.as_view(),
            name='setupgame'),
        url(r'^next/$',
            views.NextTurnView.as_view(),
            name='next'),
        url(r'^draw/$',
            views.DrawDieView.as_view(),
            name='draw'),
        url(r'^pick/$',
            views.PickDestView.as_view(),
            name='pick'),
]


from rtda.local_vars import LocalVars
from rtda.operand_stack import OperandStack

class Frame(object):

    def __init__(self,maxLocals,maxStack):
        self.localVars = LocalVars(maxLocals) # LocalVars
        self.operandStack = OperandStack(maxStack) # * OperandStack
        self.thread = []
        self.nextPC = None


import tkinter
from tkinter import ttk
import requests
import re
import urllib.request,urllib.error
import bs4
import os
import urllib.parse
from fake_useragent import UserAgent
import time
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.figure import Figure
import tkinter as tk
#设置cookie
cookie = {
    '_uid': '1097055',
    '__client_id': '6d44000048a6adae6f6f7b4405948d69419b9c06',
}
data=[]
ProblemUrl = "https://www.luogu.com.cn/problem/"
ListUrl = "https://www.luogu.com.cn/problem/list?"
UrlSolution = "https://www.luogu.com.cn/problem/solution/"
savePath = r"C:\Users\12979\Desktop\luogu_crawler\project\\"

# 反爬设置，使用fake_useragent库随机生成一个User-Agent
ua = UserAgent()
random_user_agent = ua.chrome

#获取题目内容：背景，描述。。。
def get_question(html):
    bs = bs4.BeautifulSoup(html,"html.parser")
    core = bs.select("article")[0]
    md = str(core)
    ques=""
    content_pid = re.compile(r'
.*?
')
    matches = re.findall(content_pid, md)
    ques=matches[0]
    ques = re.sub("
", "",ques)
    ques = re.sub("
", "", ques)
    md = re.sub("
","# ",md)
    md = re.sub("
","## ",md)
    md = re.sub("
","#### ",md)
    md = re.sub("</?[a-zA-Z]+[^<>]*>","",md)
    return md,ques
#获取题解
def get_answer(html):
    soup = bs4.BeautifulSoup(html, "html.parser")
    lock_F = soup.find('script')
    lock_S = lock_F.text
    begin = lock_S.find('"')
    endword = lock_S.find('"', begin + 1)
    lock_S = lock_S[begin + 1:endword]
    #解码
    #1.对字符串进行URL解码，将特殊字符转义为原始字符。然后，
    #2.将字符串中的\ / 替换为 /，以还原正常的URL格式。接着，将字符串编码为UTF - 8格式，
    #3.使用unicode_escape进行Unicode转义字符的解码。
    decode = urllib.parse.unquote(lock_S)
    decode = decode.replace('\/', '/')
    decode = decode.encode('utf-8').decode('unicode_escape')
    #截取题解
    begin = decode.find('"content":"')
    endword = decode.find('","type":"题解"')
    decode = decode[begin + 11:endword]
    return decode
#文件保存
def save_file(data,filename):
  file = open(filename,"w",encoding="utf-8")
  for d in data:
    file.writelines(d)
  file.close()
class GUI(ttk.Frame):
    #初始化页面
    def __init__(self, parent):
        super().__init__(parent, padding=15)
        self.columnconfigure(0, weight=1)
        self.GUI_elements()
    #GUI页面绘制
    def GUI_elements(self):
        #关键词输入框
        self.keyword_label = ttk.Label(self, text="请输入关键词：",font=('Arial', 20,'bold'))
        self.keyword_label.grid(row=0, column=0, padx=5,pady=10, sticky="w")
        self.keyword_entry = tk.Entry(self,width=50,font=('Arial', 20))
        self.keyword_entry.grid(row=1, column=0, padx=5, pady=10, sticky="ew")
        #难度选择框
        self.diff_label = ttk.Label(self, text="请选择难度：", font=('Arial', 20, 'bold'))
        self.diff_label.grid(row=2, column=0, padx=5, pady=10, sticky="w")
        diff_list = ["题目难度","全部","暂未评定", "入门", "普及-", "普及&提高-", "普及+&提高", "提高+&省选-", "省选&NOI-", "NOI&NOI+&CTSC"]
        self.combobox1 = ttk.Combobox(self,values=diff_list,font=("Arial", 20))
        self.combobox1.current(0)
        self.combobox1.grid(row=3, column=0, padx=5, pady=10, sticky="ew")
        self.combobox1.configure(font=("Arial", 20))
        #开始按钮
        style = ttk.Style()
        style.configure("Custom.TButton", font=("Arial", 20))
        self.button = ttk.Button(self, text="开始爬取", width=15,style="Custom.TButton")
        self.button.grid(row=6, column=0, padx=5, pady=10, sticky="ew")
        # 爬取情况文本框
        self.process_text = tkinter.StringVar()
        self.process_label = ttk.Label(self, textvariable=self.process_text)
        self.process_label.grid(row=7, column=0, padx=5, pady=(10, 0), sticky="w")
        self.time_text = tkinter.StringVar()
        self.time_label = ttk.Label(self, textvariable=self.time_text)
        self.time_label.grid(row=8, column=0, padx=5, pady=(10, 0), sticky="w")
        self.total_text = tkinter.StringVar()
        self.total_label = ttk.Label(self, textvariable=self.total_text)
        self.total_label.configure(font=('Arial', 20))
        self.total_label.grid(row=10, column=0, padx=5, pady=(10, 0), sticky="w")
        #实时更新绘制折线
        self.fig = Figure(figsize=(5, 4), dpi=100)
        self.ax = self.fig.add_subplot(111)
        self.canvas = FigureCanvasTkAgg(self.fig, master=self)
        self.canvas.draw()
        self.canvas.get_tk_widget().grid(row=11, column=0, pady=10)
        #开始爬取函数
        def click():
            self.total_text.set("总用时：秒")
            def update_graph(count_list, time_list):
                self.ax.clear()
                self.ax.plot(count_list, time_list)
                self.ax.set_xlabel("answer count")
                self.ax.set_ylabel("time (s)")
                self.ax.set_title("answer count vs time")
                self.canvas.draw()
            global keyword, difficulty
            keyword = self.keyword_entry.get()
            diff_select = self.combobox1.get()
            savePath = r"C:\Users\12979\Desktop\luogu_crawler\project\\"
            savePath = savePath + str(diff_select) + '-' + keyword
            print(savePath)
            #创建文件夹
            if not os.path.exists(savePath):
                os.mkdir(savePath)
            #获取难度
            for i in range(len(diff_list)):
                if diff_list[i] == diff_select:
                    difficulty = i - 2
                    #print(type(difficulty))
                    if difficulty <0:
                        difficulty = ''
            Url_final = ListUrl + "&difficulty=" + str(difficulty) + "&keyword=" + str(keyword) + "&type=B%7CP"
            headers = {
                "User-Agent": random_user_agent
            }
            resp = requests.get(Url_final, headers=headers)
            pattern = re.compile(r'()*?') #匹配 HTML 锚标签  的起始部分
            pid = re.findall(pattern, resp.text)
            #获取题目名称与编号
            pid_list = []
            for x in pid:
                if (x != ""):
                    pid_list.append(x[9:14])#9到14位刚好就是题目id
            count = 0
            #存放时间
            time_list = []
            count_list = []
            time_per = []
            start_time = time.time()
            for i in range(len(pid_list) - 2):
                start_time_per = time.time()
                print(len(pid_list) - 2)
                #爬取15道题目
                if (count == len(pid_list) - 2 or count >= 15):
                    break
                count += 1
                #计算总时间
                end_time = time.time()
                duration = end_time - start_time
                print(duration)
                time_list.append(duration)
                count_list.append(i + 1)
                #实时更新图表
                update_graph(count_list, time_list)
                #实时更新文本
                title = "正在爬取第{}题,题目编号为".format(i + 1)+pid_list[i]
                self.process_text.set(title)
                self.process_label.configure(font=('Arial', 20))
                self.process_label.update()
                #根据pid分别获取题目内容和题解
                Problemurl = ProblemUrl + pid_list[i]
                resp = requests.get(Problemurl, headers=headers, cookies=cookie)
                ques_html = resp.text
                Solutionurl = UrlSolution + pid_list[i]
                resp = requests.get(Solutionurl, headers=headers, cookies=cookie)
                ans_html = resp.text
                #判断是否爬取成功
                if ques_html == "error":
                    print("爬取失败")
                else:
                    print("正在爬取第{}题,题目编号为".format(i + 1) + pid_list[i])
                    questionMD, ques = get_question(ques_html)
                    solutionMD = get_answer(ans_html)
                    print("爬取成功！正在保存\n", end="")
                    print(f"（用时：{duration:.2f}秒）")
                    #根据系统修改文件名中的非法字符
                    cleaned_pid = re.sub(r'[\\/:*?"<>|\[\]【】]', '-', pid_list[i])
                    cleaned_ques = re.sub(r'[\\/:*?"<>|\[\]【】]', '-', ques)
                    #题目与题解分开保存
                    folder_path = os.path.join(savePath, f"{cleaned_pid}-{cleaned_ques}")
                    if not os.path.exists(folder_path):
                        os.mkdir(folder_path)
                    problem_file_path = os.path.join(folder_path, f"{cleaned_pid}-{cleaned_ques}.md")
                    solution_file_path = os.path.join(folder_path, f"{cleaned_pid}-{cleaned_ques}-题解.md")
                    save_file(questionMD, problem_file_path)
                    save_file(solutionMD, solution_file_path)
                    print("保存成功!\n")
                    #计算每道题目的时间
                    end_time_per = time.time()
                    duration_per = end_time_per - start_time_per
                    time_per.append(duration_per)
                    title2 = "第{}题爬取成功题目编号为".format(i + 1) + pid_list[i]+ "（用时：{:.2f}秒）".format(duration_per)
                    self.time_text.set(title2)
                    self.time_label.configure(font=('Arial', 20))
                    self.time_label.update()
            self.process_text.set("爬取完毕!结果请前往" + savePath + "查看")
            self.total_text.set("总用时：{:.2f}秒".format(duration))
        self.button.config(command=click)
class Crawler(ttk.Frame):
    def __init__(self, parent):
        super().__init__(parent, padding=15)
        GUI(self).grid(row=0, column=1, rowspan=2, padx=10, pady=(10, 0), sticky="nsew")
def main():
    root = tkinter.Tk()
    root.title("洛谷爬虫")
    Crawler(root).pack(expand=True)
    root.mainloop()
if __name__ == "__main__":
    main()


'''
PROBLEM DESCRIPTION:
Given a string containing just the characters '(', ')', '{', '}', '[' and ']', determine if the input string is valid.

An input string is valid if:

Open brackets must be closed by the same type of brackets.
Open brackets must be closed in the correct order.
Note that an empty string is also considered valid.
'''

class Solution:
    def isValid(self, s: str) -> bool:

        l = len(s)
        if l % 2 != 0:
            return False
        if l == 0:
            return True

        while "{}" in s or "()" in s or "[]" in s:
            s = s.replace('{}', '').replace('()', '').replace('[]', '')
        return s == ''


# -*- coding: utf-8 -*-

import pandas as pd # This is the standard

# Reading a csv into Pandas.
# df1 = pd.read_csv('Scraped Data/RugbyData_Matchup_england v france.csv', header=0)

# Getting first x rows.
# df1.head(5)
# df1.tail(5)

# df1.columns
# df1['HomeTeam']
# len(df1)

# Finding out basic statistical information on your dataset.
# pd.options.display.float_format = '{:,.3f}'.format # Limit output to 3 decimal places.
# df1.describe()

# df1[df1.HomeTeam == 'France']

# df2 = pd.read_csv('Scraped Data/RugbyData_Matchup_england v ireland.csv', header=0)
# df1.to_csv('total.csv')
# df2.to_csv('total.csv')

# with open('total.csv', 'a') as f:
#     df2.to_csv(f, header=0)

# Combine all the matchups data together into total.csv
countries = ['england', 'france', 'ireland', 'italy', 'scotland', 'wales']
for c1 in countries:
    for c2 in countries:
        try:
            df = pd.read_csv('Scraped Data/RugbyData_Matchup_'+c1+' v '+c2+'.csv', header=1)
            df.to_csv('allmatchup.csv', mode='a', header = False)
        except:
            print('wrong opening file with:'+c1+' v '+c2)
            continue

matchup_all = pd.read_csv('allmatchup.csv', names = ["num", "Unnamed", "HomeTeam", "Score", "AwayTeam", "Date", "Ground"])
matchup_all.head(5)
matchup_all.columns
# matchup_all = matchup_all.drop("num",1)
# matchup_all.reset_index()
# matchup_all["num"][5]

matchup_all.to_csv("allmatchup.csv")

# For question 1.1 What are the possible matchups (Team A vs Team B)
# And how many times they matchuped up as this fixtures

matchup_count = matchup_all.groupby(["HomeTeam","AwayTeam"]).size().reset_index().rename(columns = {0:'total'})

# Split the Score column into HomeTeam Score and AwayTeam Score

# matchup_all["Score"].str.split(' - ',1).tolist()
matchup_all["HomeTeamScore"] = matchup_all["Score"].str.split(' - ',1).str[0]
matchup_all["AwayTeamScore"] = matchup_all["Score"].str.split(' - ',1).str[1]

# check the splitted scores
matchup_all.columns
matchup_all[["Score", "HomeTeamScore", "AwayTeamScore"]]

# mark HomeTeam or AwayTeam is the winner
# WhoWin = 1 : HomeTeam Win
# WhoWin = 2 : AwayTeam Win
# WhoWin = 0 : Draw
matchup_all.loc[matchup_all["HomeTeamScore"] > matchup_all["AwayTeamScore"],"WhoWin"] = int(1)
matchup_all.loc[matchup_all["HomeTeamScore"] < matchup_all["AwayTeamScore"],"WhoWin"] = int(2)
matchup_all.loc[matchup_all["HomeTeamScore"] == matchup_all["AwayTeamScore"],"WhoWin"] = int(0)

# For question 1.2 What has historically been the winning percentage of each matchup (Team A vs Team B).
win_count = matchup_all.groupby(["HomeTeam","AwayTeam", "WhoWin"]).size().reset_index().rename(columns = {0:'count'})
win_all = pd.merge(win_count, matchup_count, on=['HomeTeam', 'AwayTeam']) #inner join two dataframes
win_all["pctg"] = win_all["count"]/win_all["total"]

win_all.to_csv("winningPercentage.csv")

# calculate in each setting, the winning possibility of each team
win_stats = win_all.groupby(["HomeTeam","AwayTeam"]).apply(lambda t:t[t.pctg == t.pctg.max()])

win_stats.to_csv("winningPercentage.csv", mode = 'a')

#################################3
matchup_all["Date"]

matchup_all.sort_values(['Date'],ascending=1).Date.to_datetime
matchup_all['day'] = matchup_all["Date"].str.split(' ',2).str[0]
matchup_all['month'] = matchup_all["Date"].str.split(' ',2).str[1]
matchup_all['year'] = matchup_all["Date"].str.split(' ',2).str[2]

matchup_all.month = matchup_all.month.replace(['Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec'],
                          [1,2,3,4,5,6,7,8,9,10,11,12])

toDateConv = matchup_all[['year', 'month', 'day']]
matchup_all["DateNew"] = pd.to_datetime(toDateConv)

# for the recent N competition of each matchup, what are the results
recent3 = matchup_all[['HomeTeam','AwayTeam','WhoWin','Score','DateNew']].sort_values(by=['HomeTeam', 'AwayTeam', 'DateNew'],
                        ascending = False).groupby(['HomeTeam','AwayTeam']).head(3)

recent3_cnt = recent3.groupby(["HomeTeam","AwayTeam", "WhoWin"]).size().reset_index().rename(columns = {0:'count'})
recent3 = pd.merge(recent3, recent3_cnt, on=['HomeTeam','AwayTeam','WhoWin'])
recent3["pctg"] = recent3["count"]/3
recent3["period"] = "recent3"
recent3.to_csv("recentWinning.csv")

recent5 = matchup_all[['HomeTeam','AwayTeam','WhoWin','Score','DateNew']].sort_values(by=['HomeTeam', 'AwayTeam', 'DateNew'],
                        ascending = False).groupby(['HomeTeam','AwayTeam']).head(5)
recent5_cnt = recent5.groupby(["HomeTeam","AwayTeam", "WhoWin"]).size().reset_index().rename(columns = {0:'count'})
recent5 = pd.merge(recent5, recent5_cnt, on=['HomeTeam','AwayTeam','WhoWin'])
recent5["pctg"] = recent5["count"]/5
recent5["period"] = "recent5"
recent5.to_csv("recentWinning.csv", mode='a')

recent10 = matchup_all[['HomeTeam','AwayTeam','WhoWin','Score','DateNew']].sort_values(by=['HomeTeam', 'AwayTeam', 'DateNew'],
                        ascending = False).groupby(['HomeTeam','AwayTeam']).head(10)
recent10_cnt = recent10.groupby(["HomeTeam","AwayTeam", "WhoWin"]).size().reset_index().rename(columns = {0:'count'})
recent10 = pd.merge(recent10, recent10_cnt, on=['HomeTeam','AwayTeam','WhoWin'])
recent10["pctg"] = recent10["count"]/10
recent10["period"] = "recent10"
recent10.to_csv("recentWinning.csv", mode='a')

recentWinning = pd.concat([recent3, recent5,recent10], sort = False)
recentWinning = recentWinning.drop_duplicates(subset=['HomeTeam','AwayTeam','WhoWin','pctg','count','period'])

recentWinning.to_csv("recentWinning.csv")


# 처음에는 list에 각 개수를 입력했는데,
# list에 저장할 필요 없이 바로 출력하면 빠름

money = [50000, 10000, 5000, 1000, 500, 100, 50, 10]

T = int(input())
for t in range(1, T+1):
    n = int(input())
    print(f'#{t}')
    for i in range(len(money)):
        print(n // money[i], end=" ")
        n %= money[i]
    print()


# WilliamsRsi
"""
@author: caitlin
This strategy uses the Williams%R indicator. This momentum indicator oscillates between 0 and -100, and shows
how the current price compares to a 14 day look back period, where a reading near -100 indicates the market is
near the lowest low and a reading near 0 indicates the market is near the highest high. This strategy combines
an 8 period rsi.
"""

def backtest_strategy(stock, start_date ):
    """
    Function to backtest a strategy
    """

    csv_file = "./data/{}_1d.csv".format( stock )

    # Get today's date
    today = datetime.datetime.now().date()

    # if the file was downloaded today, read from it
    #if  ( ( os.path.exists ( csv_file ) ) and ( datetime.datetime.fromtimestamp ( os.path.getmtime ( csv_file ) ).date() == today ) ):
    if os.path.exists(csv_file) and (lambda file_path: datetime.datetime.now() - datetime.datetime.fromtimestamp(os.path.getmtime(file_path)) < datetime.timedelta(minutes=60))(csv_file):
        data = pd.read_csv ( csv_file, index_col='Date' )
    else:
        # Download data
        data = yf.download(stock, start=start_date, progress=False)
        data.to_csv ( csv_file )

    # Calculate indicators
    data = __RSI ( data, 8 )
    data = __WR  ( data, 14 )

    # Set initial conditions
    position = 0
    buy_price = 0
    sell_price = 0
    returns = []

    # Loop through data
    for i in range(len(data)):

        # Buy signal
        if ( position == 0 ) and ( ( ( data['WR_14'][i - 2] < -70) | (data['WR_14'][i - 1] < -70) | (data['WR_14'][i] < -70) | (data['WR_14'][i - 3] < -70)) & (( data['RSI_8'][i] < 30) | (data['RSI_8'][i - 1] < 30) | (data['RSI_8'][i - 2] < 30)) ):
            position = 1
            buy_price = data["Adj Close"][i]
            #print(f"Buying {stock} at {buy_price}")

        # Sell signal
        elif ( position == 1 ) and ( ( ( data['WR_14'][i - 2] > -30) | (data['WR_14'][i - 1] > -30) | (data['WR_14'][i] > -30) | (data['WR_14'][i - 3] > -30)) & ((data['RSI_8'][i] > 70) | (data['RSI_8'][i - 1] > 70) | (data['RSI_8'][i - 2] > 70)) ):
            position = 0
            sell_price = data["Adj Close"][i]
            #print(f"Selling {stock} at {sell_price}")

            # Calculate returns
            returns.append((sell_price - buy_price) / buy_price)

    # Calculate total returns
    total_returns = (1 + sum(returns)) * 100000
    percentage = ( ( (total_returns - 100000) / 100000) * 100)
    percentage = "{:.0f}".format ( percentage )

    return percentage + '%'

high = data['High']
low = data['Low']
close = data['Adj Close']

data = __RSI ( data, 8 )
data = __WR ( data, 14 )

# BUY signal: when williams indicator is less than -70 and rsi is less than 30 within the last 3 candles
if ( ( ( data['WR_14'][-3] < -70) | (data['WR_14'][-2] < -70) | (data['WR_14'][-1] < -70) | (data['WR_14'][-4] < -70))
    & (( data['RSI_8'][-1] < 30) | (data['RSI_8'][-2] < 30) | (data['RSI_8'][-3] < 30)) ):
    print_log ( 'williams_rsi.py', 'LONG', [ 'RSI_8', 'WR_14' ] , backtest_strategy ( ticker , '2020-01-01' ))

# SELL signal: when williams indicator is greater than -30 and rsi is greater than 70 within last 3 candles
if ( ( ( data['WR_14'][-3] > -30) | (data['WR_14'][-2] > -30) | (data['WR_14'][-1] > -30) | (data['WR_14'][-4] > -30))
    & ((data['RSI_8'][-1] > 70) | (data['RSI_8'][-2] > 70) | (data['RSI_8'][-3] > 70)) ):
    print_log ( 'williams_rsi.py', 'SHORT', [ 'RSI_8', 'WR_14' ] , backtest_strategy ( ticker , '2020-01-01' ) )


import json

from nonebot import on_command, on_notice
from nonebot.adapters.onebot.v11 import (
    MessageSegment,
    GROUP_ADMIN,
    GROUP_OWNER,
    NoticeEvent,
    Message,
    Event,
    Bot, )
from nonebot.permission import SUPERUSER

path = 'F:\document\OneDrive - 南京农业大学\My_codes\python\my_bot1\src\plugins\plugin1\db'#数据文件的存储位置
file = path + '\GroupInOrDecreaseNotice.json'

GroupMemNotice = on_notice(priority=5)
GroupNoticeswitch = on_command('进退群提醒',priority=5,block=True,permission=SUPERUSER | GROUP_ADMIN | GROUP_OWNER)
GroupNoticeQuery = on_command('查询进退群提醒',priority=5,block=True,permission=SUPERUSER | GROUP_ADMIN | GROUP_OWNER)

@GroupMemNotice.handle() # 进退群提醒处理
async def GroupMemNotice_(bot:Bot,event:NoticeEvent):
    group_id = str(event.get_session_id()).split('_')[1]
    user_id = event.get_user_id()
    user_info = await bot.get_stranger_info(user_id=int(user_id))
    head_url = f"https://q.qlogo.cn/g?b=qq&nk={user_id}&s=640"#QQ头像地址
    print(user_info)
    try:
        await create_document()
        f = open(file,'r')
        print('<<<'+group_id+'>>>')
        switch = str(json.load(f)[group_id])
        print('<<<'+switch+'>>>')

        if switch == '1':
            description = json.loads(event.get_event_description().replace("'",'"'))
            if description['notice_type'] == 'group_increase':
                await GroupMemNotice.finish(
                    MessageSegment.at(user_id=user_id)+Message(f"欢迎新成员{user_info['nickname']}的加入!") + \
                    MessageSegment.image(head_url) + Message('既然来了就不要走了哦੭ ᐕ)੭*⁾⁾')
                )
            elif description['notice_type'] == 'group_decrease':
                if description['sub_type'] == 'leave':
                    await GroupMemNotice.finish(
                        MessageSegment.image(head_url) + Message(f"{user_info['nickname']}咪咪的离开了我们...")
                    )
                elif description['sub_type'] == 'kick':
                    op_user_id = description['operator_id']
                    await GroupMemNotice.finish(MessageSegment.image(head_url) +
                                            "发现超级可爱的 " + MessageSegment.at(op_user_id) + f"({op_user_id})"
                                            f" 面无表情地把 {user_info['nickname']}({user_id}) 踹了出去，当时害怕极了喵。。")
    except:
        pass

@GroupNoticeswitch.handle()#进退群提醒开关
async def GroupNoticeswitch_(event:Event):
    group_id = str(event.get_session_id()).split('_')[1]
    print(group_id)
    try:
        with open(file,'r') as f:
            content = json.load(f)
    except FileNotFoundError:
        open(file,'w')
        content = {}
    f.close

    try:
        switch = content[group_id]
    except KeyError:
        switch = '0'

    if switch == '1':
        switch = '0'
    else:
        switch = '1'

    data = {group_id:switch}
    content.update(data)

    with open(file,'w') as f_new:
        json.dump(content,f_new)

    if switch == '1':
        await GroupNoticeswitch.finish('进退群提醒已开启')
    elif switch == '0':
        await GroupNoticeswitch.finish('进退群提醒已关闭')
    else:
        await GroupNoticeswitch.finish('进退群提醒错误,请检查日志!')

@GroupNoticeQuery.handle()#查询是否开启进退群提醒
async def GroupNoticeQuery_(event:Event):
    group_id = str(event.get_session_id()).split('_')[1]
    try:
        await create_document()
        f = open(file,'r')
        print('<<<'+group_id+'>>>')
        switch = str(json.load(f)[group_id])
        print('<<<'+switch+'>>>')
        if switch == '1':
            await GroupNoticeQuery.finish('已开启进退群提醒!')
        else:
            await GroupNoticeQuery.finish('未开启进退群提醒!')
    except KeyError:
        await GroupNoticeQuery.finish('未开启进退群提醒!')

async def create_document():#检查有没有数据库,没有就创建,有就pass
    try:
        import os
        os.mkdir(path)
        print('<reate floder cuccesful>')
    except FileExistsError:
        pass
    try:
        f = open(path+'\\GroupInOrDecreaseNotice.json','r')
    except FileNotFoundError:
        f = open(path+'\\GroupInOrDecreaseNotice.json','w')
        f.close()
        print('<create documen cuccesful>')
    print('<create_document() handed>')


def task(x1, y1, x2, y2, x3, y3, x4, y4):
    if x3 > x1 and x3 < x2 and y3 > y2 and y3 < y4:
        lenght = x2 - x3
        width = y1 - y3
        s = lenght * width
        return s
    else:
        return "Не входит"
print(task(1,1,2,2,3,3,4,4))


from processing.url_processing import extract_uri
import json

class TrafficData:
    def __init__(self, row):
        if row is not None:
            (
                self.session,
                self.url,
                self.created_at,
                self.action_target,
                self.action_data,
                self.action,
                self.timestamp
            ) = row
        self.url = extract_uri(self.url)

    def __repr__(self, seperator=","):
        return json.dumps({
            "url":self.url,
            "action": self.action,
            "action_data": self.action_data
        })


from framework.base_page import BasePage
import time

class SportNewsHomePage(BasePage):
    # NBA入口
    nba_link = "xpath=>//div[@class='schedule clearfix']//li[1]"
    time.sleep(1)

    def click_nba_link(self):
        self.click(self.nba_link)
        self.sleep(2)


from collections import defaultdict
from operator import itemgetter
import re

def isRealRoom(name, checksum):
    if len(checksum) != 5:
        raise Exception
    totals = defaultdict(int)
    for c in name:
        if c != '-':
            totals[c] += 1
    pairs = zip(totals.keys(), totals.values())
    alphaPairs = sorted(pairs, key=itemgetter(0))
    freqPairs = sorted(alphaPairs, key=itemgetter(1), reverse=True)
    genCheckSum = ''
    for a, b in freqPairs:
        genCheckSum += a
    return genCheckSum[:5] == checksum

def main():
    f = open('input.txt', 'r')
    sectorSum = 0
    for line in f:
        room, metadata = line.rsplit('-', 1)
        match = re.search(r'(\d+)\[(.{5})\]', metadata)
        sector = int(match.group(1))
        checksum = match.group(2)
        if(isRealRoom(room, checksum)):
            sectorSum += sector
    print(sectorSum)

if __name__ == "__main__":
    main()


from model_vgg_ctc import SpeechModel
import os
import platform as plat

#在此可补一个显存分配的设置
os.environ["CUDA_VISIBLE_DEVICES"]='0'
temp=os.environ["CUDA_VISIBLE_DEVICES"]
gpu_nums=len(temp.split(','))

S_M=SpeechModel('dataset')
filename='QAQ'

#保存模型图（说实话丑死了）
if not (os.path.exists(os.path.join('model_image','_ctc.png')) or \
    os.path.exists(os.path.join('model_image',filename+'_data.png'))):
    S_M.VisualModel(filename)

#S_M.PredModel(filename)
S_M.TestModel(filename,1000,type='train')


#!/home/benemenadmin/.virtualenvs/qlik_reloader/bin/python

# IMPORTS
import json
import sys
from websocket import create_connection
import ssl
sys.path.append("/home/benemenadmin/PyResources") # Directory for Python classes
from gpconnect import GP_Connect, Timestamp_Update, Log_Table, Qlik_Connect
import emailer
import multiprocessing
from multiprocessing import Pool, Value, Queue
import datetime
from random import randint
import signal
from contextlib import contextmanager
import requests
import time
import uuid
counter = None

#WEBSOCKET INFO
folder = "/home/benemenadmin/PyResources/Qlik_Websocket/"
certs = ({"ca_certs": folder + "root.pem",
            "certfile": folder + "client.pem",
            "keyfile": folder + "client_key.pem",
            "cert_reqs":ssl.CERT_NONE,
            "server_side": False
            })

def call_esb():
	call = requests.get('http://80.88.187.171:9090/qlik/updateLoaded')
	if call.status_code == 200:
		print("ESB reload call successfull.")
		return True
	else:
		print("ESB reload failed due to:\nStatus Code:" + str(call.status_code) + "\n" + "Reason:" + str(call.reason) + "\n" + "Response:" + str(call.text))

#CLASS FOR POOL PROCESSING
class Reloader(object):

	def __init__(self, path):

		self.n = randint(5,6) #RANDOME INT (5 or 6) TO SWITCH BETWEEN THE TWO AVAILABLE IP ADDRESSES (80.88.187.45 or 80.88.187.46)
		self.url = "wss://80.88.187.4"+str(self.n)+":4747/app/"+path #CREATES UNIQUE WEBSOCKET CONNECTION FOR EACH APP TO BE RELOADED
		self.error = True
		self.is_error = ""

		#FOR LOGGING DURATION
		self.start_time = datetime.datetime.utcnow()
		self.end_time = datetime.datetime.utcnow()

		#JSON REQUEST CREATION AND FULL/PARTIAL RELOAD VARIABLE
		self.request ={}
		self.full = False

		#WEBSOCKET CREATION
		userDirectory, userId = "INTERNAL","apiuser"
		ssl.match_hostname = lambda cert, hostname: True
		self.ws = create_connection(self.url, sslopt=certs, header={'X-Qlik-User: UserDirectory=%s; UserId=%s'% (userDirectory, userId)})

		#ADDITIONAL BASE VARIABLES FOR LATER USE IN THIS CLASS
		self.path = ""
		self.string = ""

	def main(self,path, full):

		#REDUNDANTLY DEFINES BOTH PATH(UUID) AND FULL/PARTIAL RELOAD
		self.full = full
		self.path = path

		#BASE JSON REQUEST
		self.request ={
				"jsonrpc": "2.0",
				"id": 0,
				"method": "OpenDoc",
				"handle": -1,
				"params": [
					path,
					""
				]
			}

		#DEFINES START TIME FOR DURATION
		self.start_time = datetime.datetime.utcnow()

		#BEGINS PROCESSING
		self.open_app()

	def open_app(self):

		#OPENS APP IN QLIKSENSE API ENGINE
		try:
			#SENDS JSON REQUEST
			self.ws.send(json.dumps(self.request))

			#REQUEST RESPONSE GATHERING
			response = self.ws.recv()
			response = json.loads(self.ws.recv())

			self.check_app()

		except Exception as e:
			print("Error while opening the app.")
			print(e)

#			if len(self.is_error) == 0:
#				self.is_error = self.path
			self.close(True) #True TELLS close() THERE HAS BEEN AN ERROR

	def check_app(self):

		#CHECKS THE APP PREVIOUSLY REQUESTED
		try:
			#MODIFIES BASE JSON REQUEST
			request = self.request
			request["id"]=1
			request["method"]="GetActiveDoc"
			request["params"]=[]

			#SENDS REQUEST
			self.ws.send(json.dumps(request))

			#REQUEST RESPONSE GATHERING
			response = json.loads(self.ws.recv())
			self.reload_app()

		except Exception as e:
			print("Error while checking app status.")
			print(e)
#			if len(self.is_error) == 0:
#				self.is_error = self.path
			self.close(True) #True TELLS close() THERE HAS BEEN AN ERROR

	def reload_app(self):

		#RELOADS APP
		try:
			#FULL RELOAD
			request = self.request
			request["id"]=2
			request["method"]="DoReloadEx"
			request["params"]={}
			request["handle"]=1

			if self.full == False: #BASED ON SQL TABLE
				#PARTIAL RELOAD
				request["params"] = [{"qMode": 1, "qPartial": True, "qDebug": False}] #CHANGED qMode 2 -> 1

			#SENDS MODIFIED REQUEST AND SAVES RESPONSE
			self.ws.send(json.dumps(request))
			response =  json.loads(self.ws.recv())
			self.count_up() # COUNTER FOR JOB NUMBER VS TOTAL JOBS

			#CREATES OUTPUT BASED ON RESPONSE FROM API ENGINE
			# IF FAILURE RESPONSE
			if response['result']['qResult']['qSuccess']== False:
				logFile = response["result"]["qResult"]["qScriptLogFile"] #GATHERS LOGFILE INFO
				self.string += str(" FAILED " +logFile) #PRINTS LOGFILE INFO
				self.is_error = str(logFile)

			#IF SUCCESS RESPONSE
			elif response['result']['qResult']['qSuccess']== True:
				logFile = response["result"]["qResult"]["qScriptLogFile"] #GATHERS LOGFILE INFO
				self.string += str(" Success "+"80.88.187.4"+str(self.n)+" - " + logFile) #PRINTS LOGFILE INFO

			self.save_app()

		except Exception as e:
			print("Error while starting reload on:")
			print(self.path)
			print(e)

			self.close(True) #True TELLS close() THERE HAS BEEN AN ERROR

		#CHECKS RELOAD STATUS
		try:
			status = response["result"]["qResult"]["qSuccess"]

		except:
			print("Error. Reload failed on: ")
			print(self.path)
			try:
				logFile = response["result"]["qResult"]["qScriptLogFile"] #GATHERS LOGFILE INFO
				print("See log at ", logFile) #PRINTS LOGFILE INFO
				self.is_error += str(logFile)
			except:
				print("No log was generated")
				self.is_error += str(self.path)
			self.close(True) #True TELLS close() THERE HAS BEEN AN ERROR

	def save_app(self):

		#SAVES APP
		try: #JSON REQUEST MODIFICATION
			request = self.request
			request["id"] = 6
			request["method"] = "DoSave"
			request["handle"] = 1  #HANDLE CHANGED FROM -1 -> 1
			request["params"] = []

			#SENDS MODIFIED REQUEST AND SAVES RESPONSE
			self.ws.send(json.dumps(request))
			response =  json.loads(self.ws.recv())

			self.close()

		except Exception as e:
			print("Error while saving the app.")
			print(e)
			self.close(True) #True TELLS close() THERE HAS BEEN AN ERROR

	def close(self, err=False):

		#LOGS END TIME FOR DURAION
		self.end_time = datetime.datetime.utcnow()

		#LOGS ALL INFO TO SQL TABLE
		l = Log_Table().main(self.path, self.start_time, self.end_time, self.full, err, task_id)

		#PRINTS ALL INFO GATHERED IN PROCESS
		print(self.string)
		if err and len(self.is_error)>0:
			failed.put(str("- ID: " + self.path + ", LOGFILE: "+ self.is_error +"\n"))
		#CLOSES WEBSOCKET CONNECTION
		self.ws.close()

		#RETURNS LOGTABLE DATA IN CASE OF LATER NEED
		return l

	def count_up(self): #COUNTER FOR JOB/TOTAL_JOBS

		global total_jobs
		global counter
		with counter.get_lock():
			counter.value += 1
			self.string += str(counter.value)
			self.string += str("/")
			self.string += str(total_jobs)

def process(paths): #PROCESS CREATION METHOD

	app1 = Reloader(paths[0])
	app1.main(paths[0],paths[1])

	if app1.main: #IF RELOADER SUCCESSFUL, UPDATES TIMESTAMPS
		t.main(paths[0])

class TimeoutException(Exception): pass

@contextmanager
def time_limit(seconds):
    def signal_handler(signum, frame):
        raise TimeoutException("Timed out!")
    signal.signal(signal.SIGALRM, signal_handler)
    signal.alarm(seconds)
    try:
        yield
    finally:
        signal.alarm(0)

def log_table(total_jobs, end_time, task_id):

	for_analysis = Log_Table().qs_python_tasks(total_jobs,start_time, end_time, task_id)

def main(task_id):

	#MULTIPROCESSING
	if call_esb():
		p = Pool(multiprocessing.cpu_count()*4, initargs = ((counter, ), [failed]))
		p.map(process, paths)
		p.close()
		p.join()
	else:
		print("ESB Process Failed.")

	#PROCESSING TIME
	end_time = datetime.datetime.utcnow()
	duration = end_time - start_time
	log_table(total_jobs, end_time, task_id)
	print("Runtime: "+str(duration)+"\n")
	failed_jobs = "{} Failed Job(s):\n".format(failed.qsize())
	print(failed_jobs)
	while failed.qsize()>0:
		print(failed.get())

if __name__=="__main__":

	failed = multiprocessing.Queue()
	counter = Value('i',0) #JOB COUNTER
	start_time = datetime.datetime.utcnow() #START TIME
	task_id = uuid.uuid4()

	#GREENPLUM CONNECTION AND TIMESTAMP
	gp = GP_Connect()
	t = Timestamp_Update()

	#CREATES LIST FROM SQL TABLE AND DETERMINES # OF JOBS
	paths = gp.main(command = "SELECT id, full_reload FROM benereports.v_app_reloads ORDER BY template_name,short_name;")
	total_jobs = len(paths)

	try:
		with time_limit(898): #Set at 2 seconds less than your desired time as it will take approx that long to start the time limit
			main(task_id)
	except TimeoutException as e:
		print("Timed out!")
		subject = "Qlik App Reloader long execution time"
		msg = "The Qlik App Reloader has exceeded the 15 minute execution time safety limit."
		emailer.raise_error(subject, msg)
		end_time = datetime.datetime.utcnow()
		log_table(total_jobs, end_time, task_id)


'''
Created on Jan 8, 2014

@author: mkiyer
'''
import os
import sys
import argparse
import logging
import subprocess
import itertools
import shutil
from multiprocessing import Process, Queue

import pysam

from assemblyline.lib.base import which
from assemblyline.lib.seq import DNA_reverse_complement
from assemblyline.lib.transcript import parse_gtf, strand_int_to_str, strand_str_to_int, NEG_STRAND, NO_STRAND, POS_STRAND, Exon
from assemblyline.lib.batch_sort import batch_sort

SIGNALP_HEADER = ['name', 'Cmax', 'Cpos', 'Ymax', 'Ypos', 'Smax', 'Spos',
                  'Smean', 'D', '?', 'Dmaxcut', 'Networksused']
PFAM_HEADER = ['seq_id', 'alignment_start', 'alignment_end',
               'envelope_start', 'envelope_end', 'hmm_acc', 'hmm_name',
               'type', 'hmm_start', 'hmm_end', 'hmm_length', 'bit_score',
               'e_value', 'significance', 'clan']

# batch sort configuration
SORT_BUFFER_SIZE = 32000

PFAM_FILES = ['Pfam-A.hmm',
              'Pfam-A.hmm.dat',
              'Pfam-A.hmm.h3f',
              'Pfam-A.hmm.h3i',
              'Pfam-A.hmm.h3m',
              'Pfam-A.hmm.h3p',
              'Pfam-B.hmm',
              'Pfam-B.hmm.dat',
              'Pfam-B.hmm.h3f',
              'Pfam-B.hmm.h3i',
              'Pfam-B.hmm.h3m',
              'Pfam-B.hmm.h3p']

def check_pfam_dir(path):
    for pfam_file in PFAM_FILES:
        if not os.path.exists(os.path.join(path, pfam_file)):
            return False
    return True

def to_fasta(name, seq):
    return '>%s\n%s' % (name, seq)

def parse_fasta_file(fasta_file):
    with open(fasta_file) as f:
        while True:
            lines = list(itertools.islice(f,2))
            if lines:
                yield lines
            else:
                break

def split_fasta_file(fasta_file, n, prefix):
    # open files
    filehandles = []
    for x in xrange(n):
        filename = '%s%d.fasta' % (prefix, x)
        filehandles.append(open(filename))
    x = 0
    for lines in parse_fasta_file(fasta_file):
        filehandles[x].writelines(lines)
        x  = (x + 1) % n
    # close files
    for f in filehandles:
        f.close()

GENETIC_CODE = {
    'ATA':'I', 'ATC':'I', 'ATT':'I', 'ATG':'M',
    'ACA':'T', 'ACC':'T', 'ACG':'T', 'ACT':'T',
    'AAC':'N', 'AAT':'N', 'AAA':'K', 'AAG':'K',
    'AGC':'S', 'AGT':'S', 'AGA':'R', 'AGG':'R',
    'CTA':'L', 'CTC':'L', 'CTG':'L', 'CTT':'L',
    'CCA':'P', 'CCC':'P', 'CCG':'P', 'CCT':'P',
    'CAC':'H', 'CAT':'H', 'CAA':'Q', 'CAG':'Q',
    'CGA':'R', 'CGC':'R', 'CGG':'R', 'CGT':'R',
    'GTA':'V', 'GTC':'V', 'GTG':'V', 'GTT':'V',
    'GCA':'A', 'GCC':'A', 'GCG':'A', 'GCT':'A',
    'GAC':'D', 'GAT':'D', 'GAA':'E', 'GAG':'E',
    'GGA':'G', 'GGC':'G', 'GGG':'G', 'GGT':'G',
    'TCA':'S', 'TCC':'S', 'TCG':'S', 'TCT':'S',
    'TTC':'F', 'TTT':'F', 'TTA':'L', 'TTG':'L',
    'TAC':'Y', 'TAT':'Y', 'TAA':'*', 'TAG':'*',
    'TGC':'C', 'TGT':'C', 'TGA':'*', 'TGG':'W',
}

class ORFInfo(object):
    FIELDS = ['transcript_id', 'gene_id', 'orf_id', 'frame',
              'chrom', 'start', 'end', 'strand', 'exons', 'seq']
    SEQ_COL_NUM = 11

    @staticmethod
    def make_orf_id(x):
        return 'ORF%09d' % (x)
    @staticmethod
    def cmp_orf_id(x, y):
        x = int(x.split('ORF')[1])
        y = int(y.split('ORF')[1])
        return cmp(x, y)

    def __init__(self):
        for field in ORFInfo.FIELDS:
            setattr(self, field, None)

    def to_table(self):
        return [self.transcript_id,
                self.gene_id,
                str(self.orf_id),
                str(self.frame),
                str(len(self.seq)),
                self.chrom,
                str(self.start),
                str(self.end),
                self.strand,
                ','.join(map(str, (e.start for e in self.exons))),
                ','.join(map(str, (e.end for e in self.exons))),
                self.seq]

    @staticmethod
    def from_table(line):
        fields = line.strip().split('\t')
        self = ORFInfo()
        self.transcript_id = fields[0]
        self.gene_id = fields[1]
        self.orf_id = fields[2]
        self.frame = int(fields[3])
        self.chrom = fields[5]
        self.start = int(fields[6])
        self.end = int(fields[7])
        self.strand = fields[8]
        exon_starts = map(int, fields[9].split(','))
        exon_ends = map(int, fields[10].split(','))
        self.exons = [Exon(x,y) for x,y in zip(exon_starts, exon_ends)]
        self.seq = fields[ORFInfo.SEQ_COL_NUM]
        return self

    def to_bed(self, name=None):
        assert all(self.exons[0].start < x.start for x in self.exons[1:])
        assert all(self.exons[-1].end > x.end for x in self.exons[:-1])
        assert self.start == self.exons[0].start
        assert self.end == self.exons[-1].end
        block_sizes = []
        block_starts = []
        for e in self.exons:
            block_starts.append(e.start - self.start)
            block_sizes.append(e.end - e.start)
        if name is None:
            name = '%s|%s|frame=%d|length=%d' % (self.transcript_id, self.gene_id, self.frame, len(self.seq))
        # make bed fields
        fields = [self.chrom,
                  str(self.start),
                  str(self.end),
                  name,
                  '0',
                  self.strand,
                  str(self.start),
                  str(self.end),
                  '0',
                  str(len(self.exons)),
                  ','.join(map(str,block_sizes)) + ',',
                  ','.join(map(str,block_starts)) + ',']
        return fields

def get_transcript_dna_sequence(t, ref_fa):
    exon_seqs = []
    error_occurred = False
    for e in t.exons:
        seq = ref_fa.fetch(t.chrom, e.start, e.end)
        if (not seq) or (len(seq) < (e.end - e.start)):
            logging.warning("transcript %s exon %s:%d-%d not found in reference" %
                            (t.attrs['transcript_id'], t.chrom, e.start, e.end))
            error_occurred = True
            break
        exon_seqs.append(seq)
    if error_occurred:
        return None
    # make fasta record
    seq = ''.join(exon_seqs).upper()
    # reverse complement negative stranded sequences
    if t.strand == NEG_STRAND:
        seq = DNA_reverse_complement(seq)
    return seq

def translate_dna_3frames(seq):
    proteins = [[] for x in xrange(3)]
    for n in range(0,len(seq),3):
        for frame in xrange(3):
            codon = seq[n+frame:n+frame+3]
            proteins[frame].append(GENETIC_CODE.get(codon, '.'))
    for frame in xrange(3):
        proteins[frame] = ''.join(proteins[frame]).rstrip('.')
    return proteins

def find_orfs(seq, start_codon='M', stop_codon='*'):
    '''
    start is inclusive and end is exclusive
    includes stop codon
    '''
    # find first start codon
    start = seq.find(start_codon)
    while (start >= 0):
        end = seq.find(stop_codon, start)
        if end == -1:
            yield start, len(seq), seq[start:]
            return
        yield start, end+1, seq[start:end+1]
        # find next start codon
        start = seq.find(start_codon, end+1)

def transcript_to_genome_pos(t, pos, reverse=False):
    reverse = (t.strand == NEG_STRAND) or reverse
    if reverse:
        exon_iter = xrange(len(t.exons)-1, -1, -1)
    else:
        exon_iter = xrange(0, len(t.exons))
    offset = 0
    for i in exon_iter:
        exon = t.exons[i]
        exon_size = exon.end - exon.start
        if pos < offset + exon_size:
            if reverse:
                return exon.start + exon_size - (pos - offset) - 1
            else:
                return exon.start + (pos - offset)
        offset += exon_size
    if reverse:
        return t.start
    else:
        return (t.end-1)

def genome_interval_to_exons(start, end, t):
    assert start >= t.start
    assert end <= t.end
    assert start < end
    i = 0
    while t.exons[i].end < start:
        i += 1
    start_exon = i
    while t.exons[i].end < end:
        i += 1
    end_exon = i
    newexons = []
    if start_exon == end_exon:
        newexons.append(Exon(start,end))
    else:
        newexons.append(Exon(start, t.exons[start_exon].end))
        for i in xrange(start_exon+1, end_exon):
            newexons.append(t.exons[i])
        newexons.append(Exon(t.exons[end_exon].start, end))
    return newexons

def orf_to_genome(t, frame, aa_start, aa_end):
    t_orf_start = frame + (3 * aa_start)
    t_orf_end = frame + (3 * aa_end) - 1 # make end position inclusive
    reverse = ((t.strand == NEG_STRAND) or (frame >= 3))
    g_orf_start = transcript_to_genome_pos(t, t_orf_start, reverse)
    g_orf_end = transcript_to_genome_pos(t, t_orf_end, reverse)
    if reverse:
        g_orf_start, g_orf_end = g_orf_end, g_orf_start
    g_orf_end += 1 # end position is inclusive so adjust
    orf_exons = genome_interval_to_exons(g_orf_start, g_orf_end, t)
    return g_orf_start, g_orf_end, orf_exons

def translate_transcript(t, ref_fa):
    # get DNA sequence of transcript
    seq = get_transcript_dna_sequence(t, ref_fa)
    # translate on all reading frames
    if t.strand == NO_STRAND:
        aa_seqs = translate_dna_3frames(seq)
        aa_seqs.extend(translate_dna_3frames(DNA_reverse_complement(seq)))
    else:
        aa_seqs = translate_dna_3frames(seq)
    return aa_seqs

def get_all_transcript_orfs(t, ref_fa, min_orf_length):
    # get amino acid sequences in all reading frames
    aa_seqs = translate_transcript(t, ref_fa)
    # get ORFs
    for frame, aa_seq in enumerate(aa_seqs):
        for aa_start, aa_end, orf_seq in find_orfs(aa_seq):
            if len(orf_seq) < min_orf_length:
                continue
            orf_start, orf_end, orf_exons = \
                orf_to_genome(t, frame, aa_start, aa_end)
            orf = ORFInfo()
            orf.transcript_id = t.attrs['transcript_id']
            orf.gene_id = t.attrs['gene_id']
            orf.frame = frame
            orf.chrom = t.chrom
            orf.start = orf_start
            orf.end = orf_end
            orf.strand = strand_int_to_str(t.strand)
            orf.exons = orf_exons
            orf.seq = orf_seq
            yield orf

def get_first_transcript_orf(t, ref_fa):
    def first_orf_to_genome(t, start, end):
        reverse = (t.strand == NEG_STRAND)
        g_orf_start = transcript_to_genome_pos(t, start, reverse)
        g_orf_end = transcript_to_genome_pos(t, end, reverse)
        if reverse:
            g_orf_start, g_orf_end = g_orf_end, g_orf_start
        g_orf_end += 1 # end position is inclusive so adjust
        orf_exons = genome_interval_to_exons(g_orf_start, g_orf_end, t)
        return g_orf_start, g_orf_end, orf_exons

    def translate_orf(seq):
        aa_seq = []
        for n in range(0, len(seq), 3):
            codon = seq[n:n+3]
            aa = GENETIC_CODE.get(codon, '.')
            aa_seq.append(aa)
            if aa == '*':
                break
        aa_seq = ''.join(aa_seq).rstrip('.')
        return aa_seq

    def find_first_orf(t, ref_fa):
        orf = ORFInfo()
        orf.transcript_id = t.attrs['transcript_id']
        orf.gene_id = t.attrs['gene_id']
        orf.chrom = t.chrom
        # get transcript sequence
        seq = get_transcript_dna_sequence(t, ref_fa)
        # find first ATG in sequence
        start = seq.find('ATG')
        if start == -1:
            orf.start = t.start
            orf.end = t.start
            orf.strand = '.'
            orf.exons = []
            orf.seq = ''
            orf.frame = 0
        else:
            aa_seq = translate_orf(seq[start:])
            end = start + 3 * len(aa_seq) - 1
            orf_start, orf_end, orf_exons = \
                first_orf_to_genome(t, start, end)
            orf.start = orf_start
            orf.end = orf_end
            orf.strand = strand_int_to_str(t.strand)
            orf.exons = orf_exons
            orf.seq = aa_seq
            orf.frame = (start % 3)
        return orf

    if t.strand == NO_STRAND:
        t.strand = POS_STRAND
        orfpos = find_first_orf(t, ref_fa)
        t.strand = NEG_STRAND
        orfneg = find_first_orf(t, ref_fa)
        if len(orfpos.seq) >= len(orfneg.seq):
            return orfpos
        else:
            return orfneg
    else:
        return find_first_orf(t, ref_fa)

def group_unique_orfs(line_iter):
    cur_orfs = []
    num_unique_orfs = 1
    cur_orf_id = ORFInfo.make_orf_id(num_unique_orfs)
    for line in line_iter:
        orf = ORFInfo.from_table(line)
        if (len(cur_orfs) > 0) and (cur_orfs[0].seq != orf.seq):
            yield cur_orfs
            cur_orfs = []
            num_unique_orfs += 1
            cur_orf_id = ORFInfo.make_orf_id(num_unique_orfs)
        orf.orf_id = cur_orf_id
        cur_orfs.append(orf)
    if len(cur_orfs) > 0:
        yield cur_orfs

def run_signalp(fasta_files, output_file, tmp_dir):
    def _consumer(worker_index, worker_fasta_file, output_queue):
        # process with signalp
        logging.debug('[Worker %d] Running SignalP' % (worker_index))
        args = ['signalp', '-f', 'short', '-t', 'euk', worker_fasta_file]
        tmp_output_file = os.path.join(tmp_dir, 'worker%d.signalp.txt' % (worker_index))
        with open(tmp_output_file, 'w') as outfileh:
            retcode = subprocess.call(args, stdout=outfileh)
            if retcode != 0:
                logging.error('[Worker %d] Error running SignalP' % (worker_index))
        # put results into output queue
        with open(tmp_output_file) as f:
            for line in f:
                if not line:
                    continue
                line = line.strip()
                if not line:
                    continue
                if line.startswith('#'):
                    continue
                output_queue.put(line)
        output_queue.put(None)
        logging.debug('[Worker %d] Finished' % (worker_index))
    # create multiprocessing queue for passing data
    output_queue = Queue(maxsize=len(fasta_files)*4)
    # start consumer processes
    procs = []
    for i in xrange(len(fasta_files)):
        p = Process(target=_consumer, args=(i, fasta_files[i], output_queue))
        p.start()
        procs.append(p)
    # get results from consumers
    num_alive = len(procs)
    tmp_output_file = os.path.join(tmp_dir, 'signalp.merge.txt')
    with open(tmp_output_file, 'w') as f:
        while num_alive > 0:
            result = output_queue.get()
            if result is None:
                num_alive -= 1
                logging.debug("Main process detected worker finished, %d still alive" % (num_alive))
            else:
                print >>f, result
    logging.debug("Joining all processes")
    # wait for consumers to finish
    for p in procs:
        p.join()
    # sort
    logging.debug('Sorting signal peptide results')
    def sort_signalp(line):
        return line.split(None,1)[0]
    sort_tmp_dir = os.path.join(tmp_dir, 'sort_tmp')
    os.makedirs(sort_tmp_dir)
    batch_sort(input=tmp_output_file,
               output=output_file,
               key=sort_signalp,
               buffer_size=SORT_BUFFER_SIZE,
               tempdirs=[sort_tmp_dir])
    # cleanup temporary files
    shutil.rmtree(sort_tmp_dir)
    return 0

def run_pfam(fasta_files, pfam_dir, output_file, tmp_dir):
    def _consumer(worker_index, worker_fasta_file, output_queue):
        # process with pfam
        logging.debug('[Worker %d] Running pfam_scan.pl' % (worker_index))
        tmp_output_file = os.path.join(tmp_dir, 'worker%d.pfam.txt' % (worker_index))
        args = ['pfam_scan.pl',
                '-cpu', '1',
                '-pfamB',
                '-fasta', worker_fasta_file,
                '-dir', pfam_dir,
                '-outfile', tmp_output_file]
        retcode = subprocess.call(args)
        if retcode != 0:
            logging.error('[Worker %d] Error running pfam_scan.pl' % (worker_index))
        # put results into output queue
        with open(tmp_output_file) as f:
            for line in f:
                if not line:
                    continue
                line = line.strip()
                if not line:
                    continue
                if line.startswith('#'):
                    continue
                output_queue.put(line)
        output_queue.put(None)
        logging.debug('[Worker %d] Finished' % (worker_index))
    # create multiprocessing queues for passing data
    output_queue = Queue(maxsize=len(fasta_files)*4)
    # start consumer processes
    procs = []
    for i in xrange(len(fasta_files)):
        p = Process(target=_consumer, args=(i, fasta_files[i], output_queue))
        p.start()
        procs.append(p)
    # get results from consumers
    num_alive = len(procs)
    tmp_output_file = os.path.join(tmp_dir, 'pfam.merge.txt')
    with open(tmp_output_file, 'w') as f:
        while num_alive > 0:
            result = output_queue.get()
            if result is None:
                num_alive -= 1
                logging.debug("Main process detected worker finished, %d still alive" % (num_alive))
            else:
                print >>f, result
    logging.debug("Joining all processes")
    # wait for consumers to finish
    for p in procs:
        p.join()
    # sort
    logging.debug('Sorting pfam results')
    def sort_pfam(line):
        return line.split(None,1)[0]
    sort_tmp_dir = os.path.join(tmp_dir, 'sort_tmp')
    os.makedirs(sort_tmp_dir)
    batch_sort(input=tmp_output_file,
               output=output_file,
               key=sort_pfam,
               buffer_size=SORT_BUFFER_SIZE,
               tempdirs=[sort_tmp_dir])
    # cleanup temporary files
    shutil.rmtree(sort_tmp_dir)
    return 0

def merge_results(orf_table_file, signalp_file, pfam_file, output_file):
    class SortedFileParser:
        def __init__(self, fileh, key_col=0, key_cmp=ORFInfo.cmp_orf_id, sep=None):
            self.fileh = fileh
            self.key_col = key_col
            self.key_cmp = key_cmp
            self.fields = None
            self.sep = sep
        def seek(self, orf_id):
            try:
                while True:
                    if self.fields is None:
                        c = -1
                    else:
                        c = self.key_cmp(self.fields[self.key_col], orf_id)
                    if c == 0:
                        # found matching line
                        return self.fields
                    elif c == 1:
                        # larger than current id so no match
                        return None
                    # smaller than current id so keep seeking
                    self.fields = self.fileh.next().strip().split(self.sep)
            except StopIteration:
                pass
    with open(output_file, 'w') as outfile:
        with open(orf_table_file) as infile:
            pfam_parser = SortedFileParser(open(pfam_file), key_col=0)
            signalp_parser = SortedFileParser(open(signalp_file), key_col=0)
            for line in infile:
                orf = ORFInfo.from_table(line)
                signalp_fields = signalp_parser.seek(orf.orf_id)
                if signalp_fields is None:
                    signalp_fields = ['None'] * len(SIGNALP_HEADER)
                pfam_fields = pfam_parser.seek(orf.orf_id)
                if pfam_fields is None:
                    pfam_fields = ['None'] * len(PFAM_HEADER)
                fields = orf.to_table()
                fields.extend(signalp_fields)
                fields.extend(pfam_fields)
                print >>outfile, '\t'.join(fields)

def orf_analysis(gtf_file, genome_fasta_file, pfam_dir,
                 output_dir, min_orf_length, first_orf_only,
                 num_processes):
    #
    # extract transcript DNA sequences, translate to protein, and
    # search for ORFs
    #
    logging.debug('Finding ORFs in transcript sequences')
    # output files
    tmp_dir = os.path.join(output_dir, 'tmp')
    if not os.path.exists(tmp_dir):
        logging.info("Creating tmp directory '%s'" % (tmp_dir))
        os.makedirs(tmp_dir)
    orf_bed_file = os.path.join(output_dir, 'transcript_orfs.bed')
    unique_orf_file = os.path.join(output_dir, 'unique_orfs.txt')
    unique_orf_bed_file = os.path.join(output_dir, 'unique_orfs.bed')
    orf_file = os.path.join(tmp_dir, 'transcript_orfs.no_ids.txt')
    sorted_orf_file = os.path.join(tmp_dir, 'transcript_orfs.no_ids.sortbyorf.txt')
    sorted_orf_id_file = os.path.join(tmp_dir, 'transcript_orfs.sortbyorf.txt')
    signalp_file = os.path.join(output_dir, 'signalp.txt')
    pfam_file = os.path.join(output_dir, 'pfam.txt')
    merged_orf_id_file = os.path.join(output_dir, 'transcript_orfs.sortbyorf.merged.txt')
    sorted_merged_orf_id_file = os.path.join(output_dir, 'transcript_orfs.sortbytranscript.merged.txt')
    # open output files
    orf_fileh = open(orf_file, 'w')
    orf_bed_fileh = open(orf_bed_file, 'w')
    # open genome fasta file
    ref_fa = pysam.Fastafile(genome_fasta_file)
    num_finished = 1
    for locus_transcripts in parse_gtf(open(gtf_file)):
        for t in locus_transcripts:
            if first_orf_only:
                # get first ORF
                orf = get_first_transcript_orf(t, ref_fa)
                if len(orf.seq) >= min_orf_length:
                    print >>orf_fileh, '\t'.join(orf.to_table())
                    print >>orf_bed_fileh, '\t'.join(orf.to_bed())
            else:
                # get all ORFs
                for orf in get_all_transcript_orfs(t, ref_fa, min_orf_length):
                    print >>orf_fileh, '\t'.join(orf.to_table())
                    print >>orf_bed_fileh, '\t'.join(orf.to_bed())
            if (num_finished % 10000) == 0:
                logging.debug('Processed %d transcripts' % (num_finished))
            num_finished += 1
    # cleanup
    orf_fileh.close()
    orf_bed_fileh.close()
    #
    # sort ORF table by ORF amino acid sequence to group identical ORFs
    # together
    #
    logging.debug('Sorting ORFs by amino acid sequence')
    def sort_by_seq(line):
        '''comparison function for batch_sort'''
        fields = line.strip().split('\t')
        return fields[ORFInfo.SEQ_COL_NUM]
    sort_tmp_dir = os.path.join(tmp_dir, 'sort_tmp')
    os.makedirs(sort_tmp_dir)
    batch_sort(input=orf_file,
               output=sorted_orf_file,
               key=sort_by_seq,
               buffer_size=SORT_BUFFER_SIZE,
               tempdirs=[sort_tmp_dir])
    shutil.rmtree(sort_tmp_dir)
    #
    # assign each ORF a unique id and write to FASTA file
    #
    logging.debug('Determining unique ORFs')
    orf_fasta_prefix = os.path.join(tmp_dir, 'orf')
    orf_fasta_files = []
    orf_fasta_sizes = []
    for i in xrange(num_processes):
        orf_fasta_files.append(open('%s%d.fasta' % (orf_fasta_prefix, i), 'w'))
        orf_fasta_sizes.append(0)
    orf_file_index = 0
    outfileh = open(sorted_orf_id_file, 'w')
    unique_orf_fileh = open(unique_orf_file, 'w')
    print >>unique_orf_fileh, '\t'.join(['orf_id', 'orf_length', 'total_occurrences', 'unique_genomic_occurrences'])
    unique_orf_bed_fileh = open(unique_orf_bed_file, 'w')
    with open(sorted_orf_file) as infileh:
        for orfs in group_unique_orfs(infileh):
            # write to master transcript/ORF table
            for orf in orfs:
                print >>outfileh, '\t'.join(orf.to_table())
            # write ORF to fasta file
            lines = to_fasta(orfs[0].orf_id, orfs[0].seq.strip('*'))
            print >>orf_fasta_files[orf_file_index], lines
            orf_fasta_sizes[orf_file_index] += 1
            # advance to next fasta file
            orf_file_index = (orf_file_index + 1) % (num_processes)
            # group by genomic position and write ORFs to BED file
            unique_genome_orfs = {}
            for orf in orfs:
                k = (orf.chrom, orf.strand, tuple(orf.exons))
                if k in unique_genome_orfs:
                    continue
                unique_genome_orfs[k] = orf
            for orf in unique_genome_orfs.itervalues():
                print >>unique_orf_bed_fileh, '\t'.join(orf.to_bed(orf.orf_id))
            # write unique ORF to tab-delimited text file
            fields = [orfs[0].orf_id, len(orfs[0].seq), len(orfs), len(unique_genome_orfs)]
            print >>unique_orf_fileh, '\t'.join(map(str,fields))
    # cleanup
    unique_orf_bed_fileh.close()
    outfileh.close()
    # get fasta files with lines written
    orf_fasta_file_names = []
    for i in xrange(len(orf_fasta_files)):
        orf_fasta_files[i].close()
        if orf_fasta_sizes[i] > 0:
            orf_fasta_file_names.append(orf_fasta_files[i].name)
    #
    # search FASTA file against signalp
    #
    logging.debug('Searching for signal peptides')
    retcode = run_signalp(orf_fasta_file_names, signalp_file, tmp_dir)
    if retcode != 0:
        logging.error('Error searching for signal peptides')
        return 1
    #
    # search FASTA file against Pfam
    #
    logging.error('Scanning for Pfam domains')
    retcode = run_pfam(orf_fasta_file_names, pfam_dir, pfam_file, tmp_dir)
    if retcode != 0:
        logging.error('Error running pfam_scan.pl')
    #
    # merge results from Pfam and signalp
    #
    logging.debug('Merging SignalP and Pfam results')
    merge_results(sorted_orf_id_file, signalp_file, pfam_file, merged_orf_id_file)
    #
    # sort by transcript id
    #
    logging.debug('Sorting ORFs by transcript ID')
    def sort_by_transcript_id(line):
        return line.split('\t', 1)[0]
    sort_tmp_dir = os.path.join(tmp_dir, 'sort_tmp')
    os.makedirs(sort_tmp_dir)
    batch_sort(input=merged_orf_id_file,
               output=sorted_merged_orf_id_file,
               key=sort_by_transcript_id,
               buffer_size=SORT_BUFFER_SIZE,
               tempdirs=[sort_tmp_dir])
    shutil.rmtree(sort_tmp_dir)
    # cleanup
    ref_fa.close()
    return 0

def full_transcript_analysis(gtf_file, genome_fasta_file, pfam_dir,
                             output_dir, num_processes):
    # output files
    pfam_file = os.path.join(output_dir, 'full_length_pfam.txt')
    # open genome fasta file
    ref_fa = pysam.Fastafile(genome_fasta_file)
    # convert transcripts to amino acid sequences and write to fasta file
    logging.debug('Writing transcript amino acid sequences to FASTA file(s)')
    tmp_dir = os.path.join(output_dir, 'tmp')
    os.makedirs(tmp_dir)
    fasta_prefix = os.path.join(tmp_dir, 'full')
    fasta_files = []
    fasta_sizes = []
    for i in xrange(num_processes):
        fasta_files.append(open('%s%d.fasta' % (fasta_prefix, i), 'w'))
        fasta_sizes.append(0)
    num_finished = 1
    fasta_file_index = 0
    for locus_transcripts in parse_gtf(open(gtf_file)):
        for t in locus_transcripts:
            # get amino acid sequences in all reading frames
            aa_seqs = translate_transcript(t, ref_fa)
            for frame, aa_seq in enumerate(aa_seqs):
                lines = to_fasta('%s|frame=%d' % (t.attrs['transcript_id'], frame), aa_seq)
                print >>fasta_files[fasta_file_index], lines
            fasta_sizes[fasta_file_index] += 1
            fasta_file_index = (fasta_file_index + 1) % num_processes
            if (num_finished % 10000) == 0:
                logging.debug('Processed %d transcripts' % (num_finished))
            num_finished += 1
    # get fasta files with lines written
    fasta_file_names = []
    for i in xrange(len(fasta_files)):
        fasta_files[i].close()
        if fasta_sizes[i] > 0:
            fasta_file_names.append(fasta_files[i].name)
    # cleanup
    ref_fa.close()
    #
    # search FASTA file against Pfam
    #
    logging.debug('Scanning for Pfam domains')
    retcode = run_pfam(fasta_file_names, pfam_dir, pfam_file, tmp_dir)
    if retcode != 0:
        logging.error('Error running pfam_scan.pl')
        return retcode
    return 0

def main():
    logging.basicConfig(level=logging.DEBUG,
                        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
    parser = argparse.ArgumentParser()
    parser.add_argument('--pfam', dest='pfam_dir')
    parser.add_argument('--genome-fasta', dest='genome_fasta_file')
    parser.add_argument('--min-orf-length', dest='min_orf_length', type=int, default=30)
    parser.add_argument('-o', '--output-dir', dest='output_dir', default='out')
    parser.add_argument('-p', '--num-processes', dest='num_processes', type=int, default=1)
    parser.add_argument('--mode', dest='mode', choices=['orf', 'first_orf', 'full'], default='orf')
    parser.add_argument('gtf_file')
    args = parser.parse_args()
    # get args
    pfam_dir = args.pfam_dir
    genome_fasta_file = args.genome_fasta_file
    gtf_file = args.gtf_file
    min_orf_length = args.min_orf_length
    num_processes = args.num_processes
    output_dir = args.output_dir
    mode = args.mode
    # check command line parameters
    if which('pfam_scan.pl') is None:
        parser.error("'pfam_scan.pl' not found in PATH")
    if which('hmmscan') is None:
        parser.error("'hmmscan' not found in PATH")
    if which('signalp') is None:
        parser.error("'signalp' not found in PATH")
    if not check_pfam_dir(pfam_dir):
        parser.error("Required Pfam-A and Pfam-B files not found at '%s'" % (pfam_dir))
    if not os.path.exists(genome_fasta_file):
        parser.error("Genome FASTA file '%s' not found" % (genome_fasta_file))
    if not os.path.exists(gtf_file):
        parser.error("GTF file '%s' not found" % (gtf_file))
    #if os.path.exists(output_dir):
    #    parser.error("Output directory '%s' already exists" % (output_dir))
    # create output dir
    if not os.path.exists(output_dir):
        logging.info("Creating output directory '%s'" % (output_dir))
        os.makedirs(output_dir)
    if mode == 'full':
        return full_transcript_analysis(gtf_file, genome_fasta_file,
                                        pfam_dir, output_dir, num_processes)
    else:
        first_orf_only = (mode == 'first_orf')
        return orf_analysis(gtf_file, genome_fasta_file, pfam_dir,
                            output_dir, min_orf_length, first_orf_only,
                            num_processes)
    return 0

if __name__ == '__main__':
    sys.exit(main())


from PyQt5 import QtWidgets, QtCore
from pyqtgraph import PlotWidget, plot
import pyqtgraph as pg
import sys  # We need sys so that we can pass argv to QApplication

from influxdb import InfluxDBClient
import prog_files.vars as CN
import random

class MainWindow(QtWidgets.QMainWindow):
    def __init__(self, *args, **kwargs):
        super(MainWindow, self).__init__(*args, **kwargs)
        self._var_list = {}
        self._client_db = InfluxDBClient(CN.DB_ADDRESS, database=CN.DB_NAME)
        self.graphWidget = pg.PlotWidget()
        self.setCentralWidget(self.graphWidget)
        self.graphWidget.setBackground('w')
        self.graphWidget.addLegend()
        # plot data: x, y value
        self.graphWidget.showGrid(x=True, y=True)
        self.timer = QtCore.QTimer()
        self.timer.setInterval(1000) # 1s
        self.timer.timeout.connect(self._update_plot_data)
        self.timer.start()
        self._last_time_measure = 0

    def _update_plot_data(self):
        res = self._client_db.query("SELECT * FROM "+CN.MEASURE_NAME+" WHERE \"time\" > " + str(self._last_time_measure)
                                    + " GROUP BY \"time()\"",
                                    epoch="ns")
        time = self._last_time_measure
        for pts in res.get_points():
            time = pts.get("time")
            for variable in pts:
                if variable == "time":
                    continue
                if pts[variable] is None:
                    continue
                if self._var_list.get(variable, None) is None:
                    self._var_list[variable] = {
                        "time": [],
                        "value": [],
                        "data":  self.graphWidget.plot([], [], variable, pen=pg.mkPen(color=(
                            random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)
                        )))
                    }
                self._var_list[variable]["time"].append(time)
                self._var_list[variable]["value"].append(pts[variable])
        self._last_time_measure = time

        for variables in self._var_list:
            self._var_list[variables]["data"].setData(self._var_list[variables]["time"],
                                                       self._var_list[variables]["value"])

def main():
    app = QtWidgets.QApplication(sys.argv)
    main = MainWindow()
    main.show()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
# @Author: marcus
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
from util_func import N_LABELS, N_CLOUDS
from util_func import load_image, encode_tags
from keras.preprocessing.image import ImageDataGenerator

class BatchImgGen(object):
    """
    A custom batch image data / label generator.
    Can be called infinite number of times.
    """
    def __init__(self, data, batch_size=64, pixel=None,
                 aug_times=0, aug_params=None):
        """
        aug_times:
            num of batches of augmented images (in addition to original image)
            to generate. if aug_times = 2 and batch_size=64, then each batch
            would have 192 images (64 original, 128 augmented)
        """
        self.df = data  # a pd dataframe
        self.idx = np.arange(len(data))
        self.curr_idx = 0
        self.batch_size = batch_size
        self.pixel = pixel
        self.steps = self._cal_steps()
        self.aug_times = aug_times  # 0: no image augumentation
        if aug_times:
            self.aug_gen = self.get_aug_generator(aug_params)

    def __next__(self):
        return self._get_batch()

    def next(self):
        return self.__next__()

    def get_steps(self):
        return self.steps

    def _cal_steps(self):
        steps, remain = divmod(len(self.df), self.batch_size)
        return steps + (remain != 0) * 1

    def _get_batch(self):
        pix = self.pixel or 256
        n_imgs = self.batch_size * (1 + self.aug_times)
        X = np.zeros((n_imgs, pix, pix, 3), dtype=np.float32)
        Y = np.zeros((n_imgs, N_LABELS), dtype=np.uint8)
        for i in range(self.batch_size):
            index = self.idx[self.curr_idx]
            X[i] = self._get_img(index)
            Y[i] = self._get_label(index)
            self._update_curr_idx()

        if self.aug_times:   # image augmentation
            X_orig, Y_orig = X[:self.batch_size], Y[:self.batch_size]
            aug = self.aug_gen.flow(X_orig, Y_orig, batch_size=self.batch_size)
            for aug_batch in range(1, 1+self.aug_times):
                aug_x, aug_y = next(aug)
                lo = aug_batch * self.batch_size
                X[lo:lo+self.batch_size] = aug_x
                Y[lo:lo+self.batch_size] = aug_y

        return X, Y

    def get_aug_generator(self, aug_params):
        if aug_params is None:
            aug_params = {
                "data_format": "channels_last", "fill_mode": 'reflect',
                "horizontal_flip": True, "vertical_flip": True,
                "rotation_range": 25.0, "channel_shift_range": 0.02,
                "width_shift_range": 0.15, "height_shift_range": 0.15,
                "zoom_range": 0.1}

        return ImageDataGenerator(**aug_params)

    def _get_img(self, index):
        return load_image(self.df.image_name.values[index], self.pixel)

    def _get_label(self, index):
        return encode_tags(self.df.tags.values[index])

    def _update_curr_idx(self):
        self.curr_idx += 1
        if self.curr_idx == len(self.idx):
            self._shuffle_idx()
            self.curr_idx = 0
        return

    def _shuffle_idx(self):
        np.random.shuffle(self.idx)  # not the index in df
        return

class TestImgIter(BatchImgGen):
    def __iter__(self):
        return self

    def _get_batch(self):
        if self.curr_idx == len(self.df):
            raise StopIteration
        pix = self.pixel or 256
        n_samples = min(self.batch_size, len(self.df) - self.curr_idx)
        X = np.zeros((n_samples, pix, pix, 3), dtype=np.float32)
        for i in range(n_samples):
            X[i] = self._get_img(self.curr_idx)
            self.curr_idx += 1
        return X

class BatchImgGen2(BatchImgGen):
    """For PlanetAmazonCNN2 (two outputs, cloud and common labels)"""

    def __next__(self):
        X, Y = self._get_batch()
        return X, {"cloud_output": Y[:, :N_CLOUDS],
                   "common_output": Y[:, N_CLOUDS:]}