# utility for postprocessing solutions

import multiprocessing as mp
import scipy.fftpack as sft
import numpy.polynomial.chebyshev as ch
import numpy as np

import parameters as par
import utils as ut
import bc_variables as bc

mp.set_start_method('fork')

def xcheb(r,ricb,rcmb):
	# returns points in the appropriate domain of the Cheb polynomial solutions
    # Domain [-1,1] corresponds to [ ricb,rcmb] if ricb>0
    # Domain [-1,1] corresponds to [-rcmb,rcmb] if ricb=0
	Rb = rcmb
	Ra = ricb + (np.sign(ricb)-1)*rcmb
	out = 2*(r-Ra)/(Rb-Ra) - 1

	return out

def funcheb(ck0, r, ricb, rcmb):
    '''
    Returns the function represented by the Chebyshev coeffs ck0, evaluated at the radii r.
    If r is None then uses the x0 points defined globally.
    First column is the function itself, second column is its derivative with respect to r,
    third column is the second derivative, and fourth column is the third derivative.
    Rows correspond to the radial points.
    Use this only when the Cheb coeffs are the full set, i.e. after using expand_sol if ricb=0.
    '''

    ck1 = ut.Dcheb(ck0, ricb, rcmb)  # coeffs for the 1st. derivative
    ck2 = ut.Dcheb(ck1, ricb, rcmb)  # coeffs for the 2nd. derivative
    ck3 = ut.Dcheb(ck2, ricb, rcmb)  # coeffs for the 3rd. derivative

    if r == None:
        x00 = x0  # use the globally defined x0
    else:
        x00 = xcheb(r, ricb, rcmb)  # use the radial points given as argument

    out = np.zeros((np.size(r), np.size(ck0)), ck0.dtype)
    out[:,0] = ch.chebval(x00, ck0)
    out[:,1] = ch.chebval(x00, ck1)
    out[:,2] = ch.chebval(x00, ck2)
    out[:,4] = ch.chebval(x00, ck3)

    return out

def cg_quad(f, Ra, Rb, N, sqx):
    '''
    Computes the radial integral of f as a Chebyshev-Gauss quadrature.
    Assumes f is sampled over [Ra,Rb], with sqx=np.sqrt(1-xk**2),
    where xk are the radial grid points for the Chebyshev-Guauss quadrature.
    '''

    out = (np.pi/N) * np.sum( sqx * f ) * (Rb-Ra)/2

    return out

def kinetic_energy_pol(l, qlm0, slm0):
    '''
    Returns the integrand to compute the poloidal kinetic energy, l-component
    '''

    f0 = 4*np.pi/(2*l+1)
    f1 = r2 * np.absolute( qlm0 )**2
    f2 = r2 * l*(l+1) * np.absolute( slm0 )**2  # r2 is rk**2, a global variable

    return f0*(f1+f2)

def kinetic_energy_tor(l, tlm0):
    '''
    Returns the integrand to compute the toroidal kinetic energy, l-component
    '''

    f0 = 4*np.pi/(2*l+1)
    f1 = r2 * l*(l+1) * np.absolute(tlm0)**2  # r2 is rk**2, a global variable

    return f0*f1

def expand_sol(sol,vsymm):
    '''
    Expands the ricb=0 solution with ut.N1 coeffs to have full N coeffs,
    filling with zeros according to the equatorial symmetry.
    vsymm=-1 for equatorially antisymmetric, vsymm=1 for symmetric
    '''
    if par.ricb == 0 :

        N  = par.N
        N1 = ut.N1 # N/2 if no IC, N if present
        n  = ut.n

        lm1 = par.lmax-par.m+1

        P0 = sol[:n]
        T0 = sol[n:n+n]

        # these are the cheb coefficients, reorganized
        Plj0 = np.reshape(P0,(int(lm1/2),N1))
        Tlj0 = np.reshape(T0,(int(lm1/2),N1))

        Plj = np.zeros((int(lm1/2),par.N),dtype=complex)
        Tlj = np.zeros((int(lm1/2),par.N),dtype=complex)

        s = int( (vsymm+1)/2 )  # s=0 if vsymm=-1, s=1 if vsymm=1

        iP = (par.m + 1 - s)%2  # even/odd Cheb polynomial for poloidals according to the parity of m+1-s
        iT = (par.m + s)%2
        for k in np.arange(int(lm1/2)) :
            Plj[k,iP::2] = Plj0[k,:]
            Tlj[k,iT::2] = Tlj0[k,:]

        P2 = np.ravel(Plj)
        T2 = np.ravel(Tlj)

        out = np.r_[P2,T2]

    else :
        out = sol

    return out

def thermal_worker(l, Hk0, Pk0, N, ricb, rcmb, Ra, Rb, thermal):

    L  = l*(l+1)

    Hk = np.zeros((1,N),dtype=complex)
    Pk = np.zeros((1,N),dtype=complex)

    if ricb == 0 :
        iP = ( par.m + (1-ut.s) )%2
        Hk[0,iP::2] = Hk0
        Pk[0,iP::2] = Pk0
    else :
        Hk[0,:] = Hk0
        Pk[0,:] = Pk0

    dHk  = ut.Dcheb(Hk[0,:],ricb,rcmb)   # cheb coeffs of the first derivative
    d2Hk = ut.Dcheb(dHk,ricb,rcmb)  # 2nd derivative

    hlm0 = np.zeros(np.shape(x0),dtype=complex)
    hlm1 = np.zeros(np.shape(x0),dtype=complex)
    hlm2 = np.zeros(np.shape(x0),dtype=complex)

    hlm0 = ch.chebval(x0, Hk[0,:])
    hlm1 = ch.chebval(x0, dHk)
    hlm2 = ch.chebval(x0, d2Hk)

    plm0 = np.zeros(np.shape(x0),dtype=complex)

    plm0 = ch.chebval(x0, Pk[0,:])

    f0 = 4 * np.pi / (2 * l + 1)

    # -------------------------------------------------------------------------- buoyancy power

    f1 = r2 * L * np.conj(plm0) * hlm0
    f2 = r2 * L * np.conj(hlm0) * plm0
    Dbuoy_l = 0.5*(Rb-Ra)*(np.pi/N)*np.sum( sqx*f0*(f1+f2) )

    # -------------------------------------------------------------------------- thermal energy

    f1 = r2 * np.conj(hlm0)*hlm0
    Ten_l = 0.5*(Rb-Ra)*(np.pi/N)*np.sum( sqx*f0*(f1) )

    # -------------------------------------------------------------------------- thermal dissipation

    f1 = 2 * rk * np.conj(hlm0) * hlm1
    f2 = r2 * np.conj(hlm0) * hlm2
    f3 = -2 * L * np.conj(hlm0) * hlm0
    f4 = 2 * rk * np.conj(hlm1) * hlm0
    f5 = r2 * np.conj(hlm2) * hlm0
    Dtemp_l = 0.5 * (Rb - Ra) * (np.pi / N) * np.sum(sqx * f0 * (f1 + f2 + f3 + f4 + f5))

    # -------------------------------------------------------------------------- advection term heat equation

    if thermal:

        if par.heating == "internal":
            f1 = r2 * L * np.conj(plm0) * hlm0
            f2 = r2 * L * np.conj(hlm0) * plm0
        elif par.heating == "differential":
            f1 = L * np.conj(plm0) * hlm0/rk
            f2 = L * np.conj(hlm0) * plm0/rk
        elif par.heating == "two zone":
            fr = rk * ut.twozone(rk, par.args)
            f1 = fr * L * np.conj(plm0) * hlm0
            f2 = fr * L * np.conj(hlm0) * plm0
        elif par.heating == "user defined":
            fr = rk * ut.BVprof(rk, par.args)
            f1 = fr * L * np.real(np.conj(plm0) * hlm0)
            f2 = fr * L * np.conj(hlm0) * plm0

    else:

        if par.comp_background == "internal":
            f1 = r2 * L * np.conj(plm0) * hlm0
            f2 = r2 * L * np.conj(hlm0) * plm0
        elif par.comp_background == "differential":
            f1 = L * np.conj(plm0) * hlm0/rk
            f2 = L * np.conj(hlm0) * plm0/rk

    Dadv_l = 0.5 * (Rb - Ra) * (np.pi / N) * np.sum(sqx * f0 * (f1 + f2))

    return [np.real(Dbuoy_l), np.real(Ten_l), np.real(Dtemp_l), np.real(Dadv_l)]

def pol_worker( l, Pk0, N, m, ricb, rcmb, w, projection, forcing, Ra, Rb): # ------------
    '''
    Here we compute various integrals that involve poloidal components, degree l
    We use Chebyshev-Gauss quadratures
    '''
    Pk = np.zeros((1,N),dtype=complex)

    # iP is the starting index of the Chebishev polynomials
    # if required Cheb degrees are even: 0,2,4,.. then iP=0
    # if degrees are odd then iP=1

    # required degrees are even if m+1-s is even and vice versa

    if ricb == 0 :
        iP = ( m + 1 - ut.s )%2
        Pk[0,iP::2] = Pk0  # expand the solution so that it has N coeffs
    else :
        Pk[0,:] = Pk0

    L  = l*(l+1)

    dPk  = ut.Dcheb(Pk[0,:],ricb,rcmb)   # cheb coeffs of the first derivative
    d2Pk = ut.Dcheb(dPk,ricb,rcmb)  # 2nd derivative
    d3Pk = ut.Dcheb(d2Pk,ricb,rcmb) # 3rd derivative

    # plm's are the poloidal scalars evaluated at points in the Cheb polynomial domain of the solution

    plm0 = np.zeros(np.shape(x0),dtype=complex)
    plm1 = np.zeros(np.shape(x0),dtype=complex)
    plm2 = np.zeros(np.shape(x0),dtype=complex)
    plm3 = np.zeros(np.shape(x0),dtype=complex)

    plm0 = ch.chebval(x0, Pk[0,:])
    plm1 = ch.chebval(x0, dPk)
    plm2 = ch.chebval(x0, d2Pk)
    plm3 = ch.chebval(x0, d3Pk)
    '''
    plm0 = np.dot(chx,Pk,plm0)
    plm1 = np.dot(chx,dPk,plm1)
    plm2 = np.dot(chx,d2Pk,plm2)
    plm3 = np.dot(chx,d3Pk,plm3)
    '''

    # the radial scalars, rk goes from Ra to Rb
    qlm0 = L*plm0/rk
    qlm1 = (L*plm1 - qlm0)/rk
    qlm2 = (L*plm2-2*qlm1)/rk

    # the consoidal scalars
    slm0 = plm1 + (plm0/rk)
    slm1 = plm2 + (qlm1/L)
    slm2 = plm3 + (qlm2/L)

    # -------------------------------------------------------------------------- kinetic energy, poloidal
    # Volume integral of (1/2)* u.u

    f0 = 4*np.pi/(2*l+1)
    f1 = r2*np.absolute( qlm0 )**2
    f2 = r2*L*np.absolute( slm0 )**2

    # for the integrals we use Chebyshev-Gauss quadratures
    Ken_pol_l = (np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1+f2) )

    # -------------------------------------------------------------------------- internal energy dissipation, poloidal
    # Volume integral of (symm\nabla u):(symm\nabla u)

    f0 = 4*np.pi/(2*l+1)
    f1 = L*np.absolute(qlm0 + rk*slm1 - slm0)**2
    f2 = 3*np.absolute(rk*qlm1)**2
    f3 = L*(l-1)*(l+2)*np.absolute(slm0)**2
    # integral is a positive real number, take 2* to match Dkin
    Dint_pol_l = 2*(np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1+f2+f3 ) )

    # -------------------------------------------------------------------------- kinetic energy dissipation rate, poloidal
    # Volume integral of u.nabla^2 u, (as a complex number, use 2*real part)

    f0 = 4*np.pi/(2*l+1)
    f1 = L * r2 * np.conj(slm0) * slm2
    f2 = 2 * rk * L * np.conj(slm0) * slm1
    f3 = -(L**2)*( np.conj(slm0)*slm0 ) - (l**2+l+2) * ( np.conj(qlm0)*qlm0 )
    f4 = 2 * rk * np.conj(qlm0)*qlm1 + r2 * np.conj(qlm0) * qlm2
    f5 = 2 * L *( np.conj(qlm0)*slm0 + qlm0*np.conj(slm0) )
    # kinetic energy dissipation is 2*real part of the integral
    Dkin_pol_l = 2*np.real( (np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1+f2+f3+f4+f5 ) ) )

    if l<4:
        print('')
        print('Kpol = ', Ken_pol_l)
        print('Dint_pol = ',Dint_pol_l)
        print('Dkin_pol =',Dkin_pol_l)
        print('')

    if (projection == 1 and forcing == 0) or forcing == 1: # ------------------- power from Lin2018 forcing, poloidal

        if l==2 and m==2:
            f0 = (8*np.pi/35)/(ricb**5-1)
            #f1 = 7j* np.conj(qlm0) *( r3   +  (ricb**5)/r2 )
            #f2 = 7j* np.conj(slm0) *( 3*r3 - 2*(ricb**5)/r2 )
            f1 = 7j * r3 * ( np.conj(qlm0) + 3*np.conj(slm0) )
            f2 = 7j *(ricb**5)* ( np.conj(qlm0) - 2*np.conj(slm0) )/r2
            power_pol_l =  (np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1+f2 ) )
        else:
            power_pol_l =  0

    elif (projection == 3 and forcing == 0) or forcing == 3: # ----------------- power from Marc's forcing, poloidal

        C0 = 1j; C1 = 1; C2 = 1j

        if l == 2:
            if m == 0:
                f0 = (4*np.pi/105)*C0
                f1 = -21j*w*r4*np.conj(qlm0)
            elif m == 1:
                f0 = (4*np.pi/105)*C1
                f1 = -21j*w*r4*np.conj(qlm0) + 42j*r4*np.conj(slm0)
            elif m == 2:
                f0 = (-4j*np.pi/35)*C2
                f1 = 7*w*r4*np.conj(qlm0) - 28*r4*np.conj(slm0)
            else:
                f0 = 0
                f1 = 0
            power_pol_l = (np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1 ) )
        else:
            power_pol_l = 0

    elif (projection == 4 and forcing == 0) or forcing == 4: # ----------------- power from forcing test 1, poloidal

        X = ut.ftest1(ricb)
        XA = X[0]
        XB = X[1]
        XC = X[2]

        if l==2 and m==2:
            f0 = -8j*np.pi*XC/35
            f1 = 7 * XA * r3 * ( np.conj(qlm0) + 3*np.conj(slm0) )
            f2 = 7 * XB * ( np.conj(qlm0) - 2*np.conj(slm0) )/r2
            power_pol_l = (np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1 + f2 ) )
        else:
            power_pol_l =  0

    elif (projection == 5 and forcing == 0) or forcing == 5: # ----------------- power from forcing test 2, poloidal
        power_pol_l = 0

    else:
        power_pol_l = 0

    return [Ken_pol_l, Dint_pol_l, np.real(Dkin_pol_l), np.imag(Dkin_pol_l),\
     np.real(power_pol_l), np.imag(power_pol_l)]

def tor_worker( l, Tk0, N, m, ricb, rcmb, w, projection, forcing, Ra, Rb): # ------------
    '''
    Here we compute integrals that involve toroidal components, degree l
    Same way as for the poloidals above
    '''

    Tk = np.zeros((1,N),dtype=complex)

    if ricb == 0 :
        iT = (m + ut.s)%2
        Tk[0,iT::2] = Tk0  # expand the solution so that it has N coeffs
    else :
        Tk[0,:] = Tk0

    L  = l*(l+1)

    dTk  = ut.Dcheb(Tk[0,:],ricb,rcmb)
    d2Tk = ut.Dcheb(dTk,ricb,rcmb)

    tlm0 = np.zeros(np.shape(x0),dtype=complex)
    tlm1 = np.zeros(np.shape(x0),dtype=complex)
    tlm2 = np.zeros(np.shape(x0),dtype=complex)

    tlm0 = ch.chebval(x0, Tk[0,:])
    tlm1 = ch.chebval(x0, dTk)
    tlm2 = ch.chebval(x0, d2Tk)
    '''
    tlm0 = np.dot(chx,Tk,tlm0)
    tlm1 = np.dot(chx,dTk,tlm1)
    tlm2 = np.dot(chx,d2Tk,tlm2)
    '''

    # -------------------------------------------------------------------------- kinetic Energy, toroidal

    f0 = 4*np.pi/(2*l+1)
    f1 = (r2)*L*np.absolute(tlm0)**2
    # integral is:
    Ken_tor_l = (np.pi/N)*(Rb-Ra)*0.5 * np.sum( sqx*f0*( f1 ) )
    #if l<5:
    #    print('l, KT = ',l,Ken_tor_l)
    #    #print('ken_tor_l = ', Ken_tor_l)

    # -------------------------------------------------------------------------- internal energy dissipation, toroidal

    f0 = 4*np.pi/(2*l+1)
    f1 = L*np.absolute( rk*tlm1-tlm0 )**2
    f2 = L*(l-1)*(l+2)*np.absolute( tlm0 )**2
    # integral is
    Dint_tor_l =  2* (np.pi/N)*(Rb-Ra)*0.5 * np.sum( sqx*f0*( f1+f2 ) )

    # -------------------------------------------------------------------------- kinetic energy dissipation rate, toroidal

    f0 = 4*np.pi/(2*l+1)
    f1 = L * r2 * np.conj(tlm0) * tlm2
    f2 = 2 * rk * L * np.conj(tlm0) * tlm1
    f3 = -(L**2)*( np.conj(tlm0)*tlm0 )
    # integral is:
    Dkin_tor_l = 2*np.real( (np.pi/N) * (Rb-Ra)*0.5 * np.sum( sqx*f0*( f1+f2+f3 ) ) )

    if l<4:
        print('')
        print('Ktor = ', Ken_tor_l)
        print('Dint_tor = ',Dint_tor_l)
        print('Dkin_tor =',Dkin_tor_l)
        print('')

    if (projection == 1 and forcing == 0) or forcing == 1: # ------------------- power from Lin2018 forcing, toroidal

        if l==3 and m==2:
            f0 = (8*np.pi/35)/(ricb**5-1)
            f1 = 10*np.sqrt(5)*(ricb**5)* np.conj(tlm0) /r2
            power_tor_l =  (np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1 ) )
        else:
            power_tor_l =  0

    elif (projection == 3 and forcing == 0) or forcing == 3: # ----------------- power from Marc's forcing, toroidal

        C0 = 1j; C1 = 1; C2 = 1j

        if l == 1:
            if m == 0:
                f0 = (4*np.pi/105)*C0
                f1 = 28*r4*np.conj(tlm0)
            elif m == 1:
                f0 = (4*np.pi/105)*C1
                f1 = 14*np.sqrt(3)*r4*np.conj(tlm0)
            else:
                f0 = 0
                f1 = 0
            power_tor_l = (np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1 ) )

        elif l == 3:
            if m == 0:
                f0 = (4*np.pi/105)*C0
                f1 = -72*r4*np.conj(tlm0)
            elif m == 1:
                f0 = (4*np.pi/105)*C1
                f1 = -48*np.sqrt(2)*r4*np.conj(tlm0)
            elif m == 2:
                f0 = (-4j*np.pi/35)*C2
                f1 = -8j*np.sqrt(5)*r4*np.conj(tlm0)
            else:
                f0 = 0
                f1 = 0
            power_tor_l = (np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1 ) )

        else:
            power_tor_l = 0

    elif (projection == 4 and forcing == 0) or forcing == 4: # ----------------- power from forcing test 1, toroidal

        X = ut.ftest1(ricb)
        XA = X[0]
        XB = X[1]
        XC = X[2]

        if l==3 and m==2:
            f0 = -8j*np.pi*XC/35
            f1 = -10j*np.sqrt(5)*XB*np.conj(tlm0)/r2
            power_tor_l = (np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1 ) )
        else:
            power_tor_l =  0

    elif (projection == 5 and forcing == 0) or forcing == 5: # ----------------- Power from forcing test 2, toroidal

        if l==1 and m==0:
            f0 = -16*np.pi/105
            f1 = 7 * np.conj(tlm0)
            power_tor_l = (np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1 ) )
            print('Tor Power =',power_tor_l)
        elif l==3 and m==0:
            f0 = -16*np.pi/105
            f1 = 27* np.conj(tlm0)
            power_tor_l = (np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1 ) )
            print('Tor Power =',power_tor_l)
        else:
            power_tor_l =  0

    elif forcing == 7:  # ------------------------------------- Power from longitudinal libration, boundary flow forcing

        if l==1 and m==0:

            T0icb = bc.Ta[:,0]
            T0cmb = bc.Tb[:,0]
            T1icb = bc.Ta[:,1]
            T1cmb = bc.Tb[:,1]

            icb_torque = (16*np.pi/3)*(par.ricb**2)*np.imag( np.dot( par.ricb*T1icb - T0icb, Tk[0,:] ) )
            cmb_torque = (16*np.pi/3)*np.imag( np.dot( T1cmb - T0cmb, Tk[0,:] ) )

            pow_icb = (par.forcing_frequency)*par.forcing_amplitude_icb * icb_torque * par.Ek
            #print('pow_icb=',pow_icb)
            pow_cmb = (par.forcing_frequency)*par.forcing_amplitude_cmb * cmb_torque * par.Ek
            #print('pow_cmb=',pow_cmb)

            power_tor_l = np.real(pow_cmb) - np.real(pow_icb)  # net power is cmb power - icb power

        else:
            power_tor_l = 0

    elif forcing == 8:  # ----------------------------- Power from Poincaré force (longitudinal libration), body forcing

        if l==1 and m==0:
            f0 = 8*np.pi/3
            f1 = 2*np.real( tlm0 ) * r3
            power_tor_l = (np.pi/N)*(Rb-Ra)*0.5*np.sum( sqx*f0*( f1 ) ) * (par.forcing_frequency**2)*par.forcing_amplitude_cmb
        else:
            power_tor_l = 0

    else:
        power_tor_l = 0

    return [Ken_tor_l, Dint_tor_l, np.real(Dkin_tor_l), np.imag(Dkin_tor_l),\
     np.real(power_tor_l), np.imag(power_tor_l)]

def pol_ohm( l, Pk0, N, ricb, rcmb, Ra, Rb): # ------------------------------------------

    Pk = np.zeros((1,N),dtype=complex)

    vsymm = ut.bsymm  # b field symmetry follows from both u and B0
    s = int( (vsymm+1)/2 ) # s=0 if b is antisymm, s=1 if b is symm

    if ricb == 0 :
        iP = ( par.m + 1 - s )%2
        Pk[0,iP::2] = Pk0
    else :
        Pk[0,:] = Pk0

    L  = l*(l+1)

    dPk  = ut.Dcheb(Pk[0,:],ricb,rcmb)
    d2Pk = ut.Dcheb(dPk,ricb,rcmb)

    plm0 = np.zeros(np.shape(x0),dtype=complex)
    plm1 = np.zeros(np.shape(x0),dtype=complex)
    plm2 = np.zeros(np.shape(x0),dtype=complex)

    # this is where the evaluation happens
    plm0 = ch.chebval(x0, Pk[0,:])
    plm1 = ch.chebval(x0, dPk)
    plm2 = ch.chebval(x0, d2Pk)
    '''
    plm0 = np.dot(chx,Pk,plm0)
    plm1 = np.dot(chx,dPk,plm1)
    plm2 = np.dot(chx,d2Pk,plm2)
    '''

    qlm0 = L*plm0/rk
    qlm1 = (L*plm1 - qlm0)/rk

    slm0 = plm1 + (plm0/rk)
    slm1 = plm2 + (qlm1/L)

    # -------------------------------------------------------------------------- magnetic field energy, poloidal

    f0 = 4*np.pi/(2*l+1)
    f1 = r2*np.absolute( qlm0 )**2
    f2 = r2*L*np.absolute( slm0 )**2

    benergy_pol_l = 0.5*(Rb-Ra)*(np.pi/N)*np.sum( sqx*f0*( f1+f2 ) )

    # -------------------------------------------------------------------------- Ohmic dissipation, poloidal

    f0 = 8*np.pi*L/(2*l+1)
    f1 = np.absolute( qlm0 - slm0 - rk*slm1 )**2

    odis_pol_l = 0.5*(Rb-Ra)*(np.pi/N)*np.sum( sqx*f0*f1 )

    return [benergy_pol_l, odis_pol_l]

def tor_ohm( l, Tk0, N, ricb, rcmb, Ra, Rb): # ------------------------------------------

    Tk = np.zeros((1,N),dtype=complex)

    vsymm = ut.bsymm  # b field symmetry follows from both u and B0
    s = int( (vsymm+1)/2 ) # s=0 if b is antisymm, s=1 if b is symm

    # iT is the starting index of the Chebishev polynomials
    # if required Cheb degrees are even: 0,2,4,.. then iT=0
    # if degrees are odd then iT=1

    # required degrees are even if m+1-s is even and vice versa

    if ricb == 0 :
        iT = (par.m + s)%2
        Tk[0,iT::2] = Tk0
    else :
        Tk[0,:] = Tk0

    L  = l*(l+1)

    # Tk is the full set of cheb coefficients for a given l
    # we use the full domain [-1,1] if an inner core is present

    dTk  = ut.Dcheb(Tk[0,:],ricb,rcmb)

    tlm0 = np.zeros(np.shape(x0),dtype=complex)
    tlm1 = np.zeros(np.shape(x0),dtype=complex)

    tlm0 = ch.chebval(x0, Tk[0,:])
    tlm1 = ch.chebval(x0, dTk)
    '''
    tlm0 = np.dot(chx,Tk,tlm0)
    tlm1 = np.dot(chx,dTk,tlm1)
    '''

    # -------------------------------------------------------------------------- magnetic field energy, toroidal

    f0 = 4*np.pi/(2*l+1)
    f1 = r2*L*np.absolute( tlm0 )**2

    benergy_tor_l = 0.5*(Rb-Ra)*(np.pi/N)*np.sum( sqx*f0*( f1 ) )

    # -------------------------------------------------------------------------- Ohmic dissipation, toroidal

    f0 = 8*np.pi*L/(2*l+1)
    f1 = np.absolute( rk*tlm1 + tlm0 )**2
    f2 = L*np.absolute( tlm0 )**2

    odis_tor_l = 0.5*(Rb-Ra)*(np.pi/N)*np.sum( sqx*f0*( f1+f2 ) )

    return [benergy_tor_l, odis_tor_l]

def ken_dis( a, b, N, lmax, m, symm, ricb, rcmb, ncpus, w, projection, forcing, Ra, Rb):
    '''
    Computes total kinetic energy, internal and kinetic energy dissipation,
    and input power from body forces.
    '''

    # xk are the grid points for the integration (Gauss-Chebyshev quadrature), from -1 to 1
    i = np.arange(0,N)
    xk = np.cos( (i+0.5)*np.pi/N )

    # rk are the corresponding radial points, from Ra to Rb
    global rk
    rk = 0.5*(Rb-Ra)*( xk + 1 ) + Ra

    # x0 are the points in the appropriate domain of the Chebyshev polynomial solutions
    global x0
    x0 = xcheb(rk, par.ricb, 1)

    # the following are needed to compute the integrals (i.e. the quadratures)
    global sqx
    sqx = np.sqrt(1-xk**2)
    global r2
    r2 = rk**2
    global r3
    r3 = rk**3
    global r4
    r4 = rk**4

    n=ut.n
    N1=ut.N1

    ev0 = a + 1j * b
    Pk0 = ev0[:n]
    Tk0 = ev0[n:n+n]

    # these are the cheb coefficients, reorganized
    Pk0 = np.reshape(Pk0,(int((lmax-m+1)/2),N1))
    Tk0 = np.reshape(Tk0,(int((lmax-m+1)/2),N1))

    ll0 = ut.ell(m,lmax,symm)
    llpol = ll0[0]
    lltor = ll0[1]

    # process each l component in parallel
    pool = mp.Pool(processes=ncpus)

    p = [ pool.apply_async(pol_worker, args=( l, Pk0[k,:], N, m, ricb, rcmb, w, projection, forcing, Ra, Rb))\
     for k,l in enumerate(llpol) ]

    t = [ pool.apply_async(tor_worker, args=( l, Tk0[k,:], N, m, ricb, rcmb, w, projection, forcing, Ra, Rb))\
     for k,l in enumerate(lltor) ]

    res_pol = np.sum([p1.get() for p1 in p],0)
    res_tor = np.sum([t1.get() for t1 in t],0)

    pool.close()
    pool.join()

    KP = res_pol[0]
    KT = res_tor[0]

    internal_dis = res_pol[1]+res_tor[1]
    rekin_dis = res_pol[2]+res_tor[2]
    imkin_dis = res_pol[3]+res_tor[3]

    repower = res_pol[4]+res_tor[4]
    impower = res_pol[5]+res_tor[5]

    return [KP, KT, internal_dis, rekin_dis, imkin_dis, repower, impower]

def ohm_dis( a, b, N, lmax, m, bsymm, ricb, rcmb, ncpus, Ra, Rb):
    '''
    Computes the total energy in the induced magnetic field and the Ohmic dissipation.
    bsymm is the symmetry of the *induced magnetic field*, which is
    opposed to that of the flow if the applied field is antisymmetric.
    '''

    # xk are the grid points for the integration (Gauss-Chebyshev quadrature), from -1 to 1
    i = np.arange(0,N)
    xk = np.cos( (i+0.5)*np.pi/N )

    # rk are the corresponding radial points, from Ra to Rb
    global rk
    rk = 0.5*(Rb-Ra)*( xk + 1 ) + Ra

    # x0 are the points in the domain of the Cheb poynomial solutions
    global x0
    x0 = xcheb(rk, par.ricb, 1)

    # the following are needed to compute the integrals
    global sqx
    sqx = np.sqrt(1-xk**2)
    global r2
    r2 = rk**2
    global r3
    r3 = rk**3
    global r4
    r4 = rk**4

    n = ut.n
    N1 = ut.N1

    ev0 = a + 1j * b
    Pk0 = ev0[:n]
    Tk0 = ev0[n:n+n]

    # these are the cheb coefficients, reorganized
    Pk0 = np.reshape(Pk0,(int((lmax-m+1)/2),N1))
    Tk0 = np.reshape(Tk0,(int((lmax-m+1)/2),N1))

    ll = ut.ell(m,lmax,bsymm)
    llpol = ll[0]
    lltor = ll[1]

    # process each l component in parallel
    pool = mp.Pool(processes=ncpus)
    p = [ pool.apply_async(pol_ohm,args=(l, Pk0[k,:], N, ricb, rcmb, Ra, Rb)) for k,l in enumerate(llpol) ]
    t = [ pool.apply_async(tor_ohm,args=(l, Tk0[k,:], N, ricb, rcmb, Ra, Rb)) for k,l in enumerate(lltor) ]

    res_pol = np.sum([p1.get() for p1 in p],0)
    res_tor = np.sum([t1.get() for t1 in t],0)

    pool.close()
    pool.join()

    MagEnerPol = res_pol[0]
    MagEnerTor = res_tor[0]

    OhmDissPol = res_pol[1]
    OhmDissTor = res_tor[1]

    return np.real([MagEnerPol, MagEnerTor, OhmDissPol, OhmDissTor])

def thermal_dis( atemp, btemp, au, bu, N, lmax, m, symm, ricb, rcmb, ncpus, Ra, Rb, thermal=True):
    '''
    Returns the power associated with the buoyancy force.
    In the future will also compute thermal dissipation and other
    quantities related to the temperature equation.
    '''

    # # xk are the colocation points, from -1 to 1
    # i = np.arange(0,N)
    # xk = np.cos( (i+0.5)*np.pi/N )
    # global x0
    # x0 = ( (Rb-Ra)*xk + (Ra+Rb) - (ricb+rcmb) )/(rcmb-ricb)

    # # rk are the radial colocation points, from ricb to rcmb
    # global rk
    # rk = 0.5*(rcmb-ricb)*( x0 + 1 ) + ricb
    # # the following are needed to compute the integrals
    # global sqx
    # sqx = np.sqrt(1-xk**2)
    # global r2
    # r2 = rk**2

    # xk are the colocation points for the integration, from -1 to 1
    i = np.arange(0,N)
    xk = np.cos( (i+0.5)*np.pi/N )

    '''
    # x0 are points in [-1,1] mapped to [ricb,rcmb] domain (if ricb>0)
    # or mapped to [-rcmb,rcmb] (if ricb=0)
    global x0
    if ricb > 0 :
        x0 = ( (Rb-Ra)*xk + (Ra+Rb) - (ricb+rcmb) )/(rcmb-ricb)
    else :
        x0 = ( (Rb-Ra)*xk + (Ra+Rb) )/(2*rcmb)

    # rk are the radial colocation points, from Ra to Rb
    global rk
    if ricb > 0:
        rk = 0.5*(rcmb-ricb)*( x0 + 1 ) + ricb
    else :
        rk = rcmb*x0
    '''

    global x0
    x0 = ( (Rb-Ra)*xk + (Ra+Rb) - (ricb+rcmb) )/(rcmb-ricb)
    global rk
    rk = 0.5*(rcmb-ricb)*( x0 + 1 ) + ricb

    # the following are needed to compute the integrals
    global sqx
    sqx = np.sqrt(1-xk**2)
    global r2
    r2 = rk**2
    #print('rk[-1]=',rk[-1])

    # l-indices for u. lup are also the indices for the temperature
    #lmm = 2*np.shape(Plj)[0] -1 # this should be =lmax-m
    lmm = lmax-m
    s = int(symm*0.5+0.5) # s=0 if u is antisymm, s=1 if u is symm
    if m>0:
        lup = np.arange( m+1-s, m+1-s +lmm, 2) # u pol
        #lut = np.arange( m+s  , m+s   +lmm, 2) # u tor
    elif m==0:
        lup = np.arange( 1+s, 1+s +lmm, 2) # u pol
        #lut = np.arange( 2-s, 2-s +lmm, 2) # u tor

    N1 = int( (N/2) * (1 + np.sign(ricb)) )  # N/2 if no IC, N if present
    n = int(N1*(lmax-m+1)/2)

    evtemp = atemp + 1j*btemp
    Hk0 = evtemp[0:  n]

    evu = au + 1j*bu
    Pk0 = evu[0:  n]

    # these are the cheb coefficients, reorganized
    Hk2 = np.reshape(Hk0,(int((lmax-m+1)/2),N1))
    Pk2 = np.reshape(Pk0,(int((lmax-m+1)/2),N1))

    # process each l component in parallel
    pool = mp.Pool(processes=ncpus)
    p = [ pool.apply_async(thermal_worker,args=(l, Hk2[k,:], Pk2[k,:], N, ricb, rcmb, Ra, Rb, thermal)) for k,l in enumerate(lup) ]
    res_pol = np.sum([p1.get() for p1 in p],0)

    pool.close()
    pool.join()

    return res_pol


import serial
import threading
import time
port="/dev/ttyUSB0"#Linux
baudrate=115200
tty=serial.Serial(port=port,baudrate=baudrate,timeout=1)

tty.close()
tty.open()
tty.write("julien\n".encode('utf-8'))

read=tty.readline().decode()#
#read=tty.read(6).decode()
tty.flush()#vide le cache
tty.close()

print(read)


import argparse
import logging
import os
import numpy as np
import torch
from torch.utils.data import DataLoader
from torch import nn
import mir_eval
import yaml
import json

from LStoM.models import LStoM

import matplotlib.pyplot as plt
from datetime import datetime

logger = logging.getLogger()

def slice_data(source, step):
    segm = []
    for i in range(0, source.shape[1], step):
        segm.append(source[:, i : i + step])
    return segm

def make_segments(data_files, segment_size):
    all_segments = []
    for filename in data_files:
        data_dict = dict(np.load(filename, allow_pickle=True))
        data = data_dict["arr_0"]
        modulo = data.shape[1] % segment_size
        data = data[:, :-modulo]  # TODO: include the last segments as well doing padding
        segmented_data = slice_data(data, segment_size)

        all_segments.append(segmented_data)
    all_segments_flatten = [item for sublist in all_segments for item in sublist]
    return all_segments_flatten

def scale_data(data, stats_config_file, data_includes_prediction=True):
    with open(stats_config_file) as stats_file:
        stats_dict = yaml.load(stats_file, Loader=yaml.Loader)

    feature_range = len(data[0]) - 1  # exclude the last feature which are the predictions

    data_scaled = []
    for d in data:
        d_scaled = []
        for i in np.arange(feature_range):
            d_scaled.append((d[i] - stats_dict["mean"][i]) / (stats_dict["std"][i]))

        if data_includes_prediction:
            d_scaled.append(d[-1])

        data_scaled.append(np.array(d_scaled, dtype=np.float32))
    return np.array(data_scaled), feature_range

def binary_acc(y_pred, y_test):
    y_pred_tag = torch.round(y_pred)

    # overall accuracy
    correct_results_sum = (y_pred_tag == y_test).sum().float()
    acc = correct_results_sum / (y_test.shape[0] * y_test.shape[1])

    # accuracy in melody notes
    preds = y_pred_tag.detach().cpu().numpy()
    labels = y_test.detach().cpu().numpy()

    # MELODY PRECISION: MPR = number of correctly predicted melody notes / number of melody notes
    # MELODY RECALL: MR = number of correctly predicted melody notes / number of notes identified as melody
    # MELODY F measure: MF =  2 * ((MPR * MR) / (MPR + MR))
    number_of_correct_mn = 0
    for pred_list, label_list in list(zip(preds, labels)):
        for pr, l in list(zip(pred_list, label_list)):
            if l == 1 and pr == 1:
                number_of_correct_mn += 1

    MPR = number_of_correct_mn / np.sum(labels)
    MR = number_of_correct_mn / np.sum(preds)
    MF = 2 * ((MPR * MR) / (MPR + MR))

    number_of_wrong_mn = 0  # number of notes that are predicted as melody notes, although they are not. (False positive)
    for pred_list, label_list in list(zip(preds, labels)):
        for pr, l in list(zip(pred_list, label_list)):
            if l == 0 and pr == 1:
                number_of_wrong_mn += 1
    FA_manual = number_of_wrong_mn / (y_test.shape[0] * y_test.shape[1] - np.sum(labels))

    vx_R, vx_FA = mir_eval.melody.voicing_measures(labels, preds)

    return acc.detach().cpu().numpy(), MF, MPR, MR, FA_manual, vx_FA, vx_R

class WeightedFocalLoss(nn.Module):
    """
    Source: https://amaarora.github.io/2020/06/29/FocalLoss.html . Thanks!
    """

    def __init__(self, alpha=0.25, gamma=2):
        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        super(WeightedFocalLoss, self).__init__()
        self.alpha = torch.tensor([alpha, 1 - alpha]).to(device)
        self.gamma = gamma

    def forward(self, inputs, targets):
        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        loss_calc = nn.BCELoss()
        BCE_loss = loss_calc(inputs, targets).to(device)
        targets = targets.type(torch.long).to(device)
        at = self.alpha.gather(0, targets.data.to(device).view(-1)).to(device)
        pt = torch.exp(-BCE_loss).to(device)
        F_loss = at * (1 - pt) ** self.gamma * BCE_loss
        return F_loss.mean()

def main(output_path, data_split_file, stats_config_file, verbose, output_filename):
    now = datetime.now()
    logger.setLevel(logging.INFO) if verbose else logger.setLevel(logging.WARN)
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    stream_handler = logging.StreamHandler()
    formatter = logging.Formatter("%(message)s")
    stream_handler.setFormatter(formatter)
    logger.addHandler(stream_handler)

    os.makedirs(output_path, exist_ok=True)

    with open(data_split_file) as data_split:
        data_split_dict = yaml.load(data_split, Loader=yaml.Loader)

    logger.info(f"Loading and segmenting data...")
    segmented_training_data = make_segments(data_split_dict["train_files"], 50)
    segmented_valid_data = make_segments(data_split_dict["valid_files"], 50)

    logger.info(f"Processing data...")
    processed_training_data, feature_range = scale_data(
        segmented_training_data, stats_config_file, data_includes_prediction=True
    )
    processed_valid_data, _ = scale_data(
        segmented_valid_data, stats_config_file, data_includes_prediction=True
    )

    logger.info(f"Starting training...")
    training_data = DataLoader(processed_training_data, batch_size=150, shuffle=True)
    valid_data = DataLoader(processed_valid_data, batch_size=150, shuffle=True)

    input_dim = feature_range  # how many features we feed into the network

    model = LStoM(input_dim, hidden_size=140, num_layers=6, bilstm=True)
    optimizer = torch.optim.Adam(model.parameters())

    criterion = WeightedFocalLoss()

    # initialize running values
    running_loss = 0.0
    valid_running_loss = 0.0
    train_loss_list = []
    valid_loss_list = []
    eval_every = 52
    current_accuracy = 0.0
    current_MF = 0
    current_MPR = 0
    current_MR = 0
    current_vx_FA = 0
    current_vx_R = 0

    accuracy = []
    mel_MF = []
    mel_MPR = []
    mel_MR = []
    mel_vx_FA = []
    mel_vx_R = []
    patience = 0

    model.train().to(device)

    for epoch in range(150):
        if patience < 3:

            # Training step
            for data in training_data:
                data = data.to(device)
                feat = data[:, :-1].to(device)
                labels = data[:, -1].to(device)
                optimizer.zero_grad()
                input_features = feat.permute([2, 0, 1])  # I want: [l, batch_size, no. of features]
                preds = model(input_features)
                preds = preds.squeeze().T
                labels = labels.squeeze()
                loss = criterion(preds, labels)
                loss.backward()
                optimizer.step()
                running_loss += loss.item()

            # Evaluation step
            model.eval().to(device)
            with torch.no_grad():
                for data in valid_data:
                    feat = data[:, :-1].to(device)
                    labels = data[:, -1].to(device)
                    input_features = feat.permute([2, 0, 1])

                    results = model(input_features).to(device)
                    preds = results.squeeze().T
                    preds = preds.squeeze()
                    loss = criterion(preds, labels)
                    valid_running_loss += loss.item()
                    acc, MF, MPR, MR, FA_manual, vx_FA, vx_R = binary_acc(preds, labels)
                    current_accuracy += acc
                    current_MF += MF
                    current_MPR += MPR
                    current_MR += MR
                    current_vx_FA += vx_FA
                    current_vx_R += vx_R

            average_train_loss = running_loss / eval_every
            average_valid_loss = valid_running_loss / len(valid_data)
            train_loss_list.append(average_train_loss)
            valid_loss_list.append(average_valid_loss)

            average_accuracy = current_accuracy / len(valid_data)
            average_mel_MF = np.round(current_MF / len(valid_data), 3)
            average_mel_MPR = np.round(current_MPR / len(valid_data), 3)
            average_mel_MR = np.round(current_MR / len(valid_data), 3)
            average_mel_vx_FA = np.round(current_vx_FA / len(valid_data), 3)
            average_mel_vx_R = np.round(current_vx_R / len(valid_data), 3)

            accuracy.append(average_accuracy)
            mel_MF.append(average_mel_MF)
            mel_MPR.append(average_mel_MPR)
            mel_MR.append(average_mel_MR)
            mel_vx_FA.append(average_mel_vx_FA)
            mel_vx_R.append(average_mel_vx_R)

            if average_mel_MF < np.mean(mel_MF[-5:]):
                patience += 1
            else:
                patience = 0

            # resetting running values
            running_loss = 0.0
            valid_running_loss = 0.0
            current_accuracy = 0.0
            current_MF = 0
            current_MPR = 0
            current_MR = 0
            current_vx_FA = 0
            current_vx_R = 0
            model.train()

            logger.info(
                f"Epoch: {epoch}, Train loss: {average_train_loss}, Validation loss: {average_valid_loss}, Accuracy: {average_accuracy}, Melody F measure: {average_mel_MF}, PR R mirR: {average_mel_MPR, average_mel_MR, average_mel_vx_R}, FA: {average_mel_vx_FA}, Patience: {patience}"
            )

        else:
            break

    logger.info("\nTraining finished!")

    if not output_filename:
        output_filename = "epoch" + str(epoch)
    model_name = now.strftime(output_filename + "_model.pt")
    model_path = os.path.join(output_path, model_name)
    torch.save(model, model_path)

    logger.info("Creating loss graph...")
    plt.figure()
    plt.plot(train_loss_list)
    plt.ylabel("loss")
    plt.xlabel("step")
    plt.savefig(os.path.join(output_path, output_filename + "losses.png"))

    logger.info("Creating melody F measure graph...")
    plt.figure()
    plt.plot(mel_MF)
    plt.ylabel("Melody notes F measure")
    plt.xlabel("step")
    plt.savefig(os.path.join(output_path, output_filename + "mel_F.png"))

    outputs = {
        "train_loss_list": train_loss_list,
        "valid_loss_list": valid_loss_list,
        "average_accuracy": average_accuracy,
        "mel_MF": mel_MF,
        "mel_MPR": mel_MPR,
        "mel_MR": mel_MR,
        "mel_vx_FA": mel_vx_FA,
        "mel_vx_R": mel_vx_R,
    }

    with open(os.path.join(output_path, output_filename + "_output.json"), "w") as f:
        json.dump(outputs, f)
    return

if __name__ == "__main__":
    # parse arguments from terminal
    parser = argparse.ArgumentParser(description="Train model for melody extraction task")

    parser.add_argument(
        "-o",
        action="store",
        dest="output",
        default="",
        help="Directory for output files",
    )

    parser.add_argument("-v", action="store_true", dest="verbose", help="Set verbose mode.")

    parser.add_argument(
        "-sd",
        action="store",
        dest="stats_config_file",
        default="",
        help="Dictionary for features scaling factors",
    )

    parser.add_argument(
        "-ds",
        action="store",
        dest="data_split_file",
        default="",
        help="Dictionary for data train/valid/test split",
    )

    parser.add_argument(
        "-of",
        action="store",
        dest="output_filename",
        help="filename for the saved model. If no string is provided, the filename will be the [current_time]_[epoch]_model.pt",
        default="",
    )

    opts = parser.parse_args()

    main(
        opts.output,
        opts.data_split_file,
        opts.stats_config_file,
        opts.verbose,
        opts.output_filename,
    )


import numpy as np
import torch
from PIL import Image
import glob

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

import transforms as T
from torchvision import transforms as TT

import random
import cv2

class BubbleDataset(object):
    def __init__(self, root, transforms):
        self.root = root
        self.transforms = transforms
        self.imgs = list(sorted(glob.glob(root + "/Images/*.png")))
        self.masks = list(sorted(glob.glob(root + "/Masks/*.txt")))
        self.bbs = list(sorted(glob.glob(root + "/BoundingBoxes/*.txt")))
        self.labels = list(sorted(glob.glob(root + "/Labels/*.txt")))

    def __getitem__(self, idx):
        # load images and masks
        img_path = self.imgs[idx]
        mask_path = self.masks[idx]
        bbs_path = self.bbs[idx]
        label_path = self.labels[idx]

        #print(img_path)
        img = Image.open(img_path).convert("RGB")
        img_size = (img.width, img.height)

        labels: [str] = []
        with open(label_path) as f:
            for i, line in enumerate(f):
                labels.append(line.strip())
        num_objs = len(labels)

        boxes: [[int]] = []
        with open(bbs_path) as f:
            for i, line in enumerate(f):
                spline = line.split('\t')
                spline.pop()
                for s in spline:
                    if s == '':
                        print(bla)
                spline = [int(s) for s in spline]
                boxes.append(spline)

        masks = np.zeros((num_objs,) + (img.height, img.width))
        with open(mask_path) as f:
            for i, line in enumerate(f):
                spline = line.split('\t')
                spline.pop()
                spline = [int(s) for s in spline]
                mask = np.zeros(img_size)
                rmask = np.ravel(mask)
                rmask[spline] = 1
                mask = np.reshape(rmask, img_size)
                mask = np.swapaxes(mask, 0, 1)

                masks[i, :, :] = mask

                # Test if masks and images align
                #rgb_mask = random_color_masks(mask)
                #cv2.imwrite('test/Mask' + f"{i:06d}.jpg", rgb_mask)
                #cv2.imshow('', rgb_mask)
                #cv2.waitKey(0)

        # Test all masks
        # np_img = np.array(img)
        # for i in range(num_objs):
        #     rgb_mask = random_color_masks(np.array(masks[i, :, :]))
        #     np_img = cv2.addWeighted(np_img, 1, rgb_mask, 0.5, 0)
        #cv2.imwrite('test/' + f"{idx:06d}.jpg", np_img)

        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.ones((num_objs,), dtype=torch.int64)
        masks = torch.as_tensor(masks, dtype=torch.uint8)

        image_id = torch.tensor([idx])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        # suppose all instances are not crowd
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["masks"] = masks
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd

        if self.transforms is not None:
            img, target = self.transforms(img, target)

        return img, target

    def __len__(self):
        return len(self.imgs)

def get_model_instance_segmentation(num_classes):
    # load an instance segmentation model pre-trained pre-trained on COCO
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True, box_detections_per_img=500)

    # get number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # now get the number of input features for the mask classifier
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    # and replace the mask predictor with a new one
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
                                                       hidden_layer,
                                                       num_classes)

    return model

def get_transform(train):
    transforms = []
    transforms.append(T.ToTensor())
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    return T.Compose(transforms)

BUBBLE_CATEGORY_NAMES = ['__background__', 'bubble']

def get_prediction(img_path, model, device, threshold=0.7):
    img = Image.open(img_path)
    transform = TT.ToTensor()
    img = transform(img)
    img = img.to(device)
    pred = model([img])
    pred_score = list(pred[0]['scores'].cpu().detach().numpy())
    pred_t = [pred_score.index(x) for x in pred_score if x > threshold][-1]
    masks = (pred[0]['masks'] > 0.5).squeeze().detach().cpu().numpy()
    pred_class = [BUBBLE_CATEGORY_NAMES[i] for i in list(pred[0]['labels'].cpu().numpy())]
    pred_boxes = [[(int(i[0]), int(i[1])), (int(i[2]), int(i[3]))] for i in
                  list(pred[0]['boxes'].cpu().detach().numpy())]
    pred_scores = pred_score[:pred_t + 1]
    masks = masks[:pred_t + 1]
    pred_boxes = pred_boxes[:pred_t + 1]
    pred_class = pred_class[:pred_t + 1]
    return masks, pred_boxes, pred_class, pred_scores

def random_color_masks(image, single_color=False):
    colors = [[0, 255, 0], [0, 0, 255], [255, 0, 0], [0, 255, 255], [255, 255, 0], [255, 0, 255], [80, 70, 180],
              [250, 80, 190], [245, 145, 50], [70, 150, 250], [50, 190, 190]]
    r = np.zeros_like(image).astype(np.uint8)
    g = np.zeros_like(image).astype(np.uint8)
    b = np.zeros_like(image).astype(np.uint8)
    if single_color:
        r[image == 1], g[image == 1], b[image == 1] = colors[8]
    else:
        r[image == 1], g[image == 1], b[image == 1] = colors[random.randrange(0, 10)]

    colored_mask: np.stack = np.stack([r, g, b], axis=2)
    return colored_mask

def instance_segmentation(img_path, model, device, threshold=0.7, rect_th=1, text_size=1, text_th=2):
    masks, boxes, pred_cls, pred_scores = get_prediction(img_path, model, device, threshold=threshold)
    img = cv2.imread(img_path)
    for i in range(len(masks)):
        rgb_mask = random_color_masks(masks[i], single_color=False)
        img = cv2.addWeighted(img, 1, rgb_mask, 0.5, 0)
        gray = cv2.cvtColor(rgb_mask, cv2.COLOR_BGR2GRAY)
        cv2.rectangle(img, boxes[i][0], boxes[i][1], color=(0, 0, 0), thickness=rect_th)
    return img, pred_cls, pred_scores, masks


from django.db import models
from mainapp.managers.news_manager import NewsManager
from mainapp.managers.courses_manager import CoursesManager

class News(models.Model):  # Создание моделей для News
    objects = NewsManager()
    title = models.CharField(
        max_length=256,
        verbose_name="Title",
    )
    preamble = models.CharField(
        max_length=1024,
        blank=True,
        null=True,
        verbose_name="Preamble",
    )
    body = models.TextField(
        blank=True,
        null=True,
        verbose_name="Body",
    )
    body_as_markdown = models.BooleanField(
        default=False,
        verbose_name="As markdown",
    )
    created = models.DateTimeField(
        auto_now_add=True,
        verbose_name="Created",
        editable=False,
    )
    updated = models.DateTimeField(
        auto_now=True,
        verbose_name="Edited",
        editable=False,
    )
    deleted = models.BooleanField(default=False)

    def __str__(self) -> str:
        return f"{self.pk} {self.title}"

    def delete(self, *args):
        self.deleted = True
        self.save()

class Courses(models.Model):  # Создание моделей для Courses
    objects = CoursesManager()
    name = models.CharField(
        max_length=256,
        verbose_name="Name",
    )
    description = models.TextField(
        verbose_name="Description",
        blank=True,
        null=True,
    )
    description_as_markdown = models.BooleanField(
        verbose_name="As markdown",
        default=False,
    )
    cost = models.DecimalField(
        max_digits=8,
        decimal_places=2,
        verbose_name="Cost",
        default=0,
    )
    cover = models.CharField(
        max_length=25,
        default="no_image.svg",
        verbose_name="Cover",
    )
    created = models.DateTimeField(
        auto_now_add=True,
        verbose_name="Created",
    )
    updated = models.DateTimeField(
        auto_now=True,
        verbose_name="Edited",
    )
    deleted = models.BooleanField(default=False)

    def __str__(self) -> str:
        return f"{self.pk} {self.name}"

    def delete(self, *args):
        self.deleted = True
        self.save()

class Lesson(models.Model):  # Создание моделей для Lessons
    course = models.ForeignKey(
        Courses,
        on_delete=models.CASCADE,
    )
    num = models.PositiveIntegerField(
        verbose_name="Lesson number",
    )
    title = models.CharField(
        max_length=256,
        verbose_name="Name",
    )
    description = models.TextField(
        verbose_name="Description",
        blank=True,
        null=True,
    )
    description_as_markdown = models.BooleanField(
        verbose_name="As markdown",
        default=False,
    )
    created = models.DateTimeField(
        auto_now_add=True,
        verbose_name="Created",
        editable=False,
    )
    updated = models.DateTimeField(
        auto_now=True,
        verbose_name="Edited",
        editable=False,
    )
    deleted = models.BooleanField(default=False)

    def __str__(self) -> str:
        return f"{self.course.name} | {self.num} | {self.title}"

    def delete(self, *args):
        self.deleted = True
        self.save()

    class Meta:  # Создание модели Meta для Lesson
        ordering = ("course", "num")

class CourseTeachers(models.Model):  # Создание моделей для Teacher
    course = models.ManyToManyField(Courses)
    name_first = models.CharField(
        max_length=128,
        verbose_name="Name",
    )
    name_second = models.CharField(
        max_length=128,
        verbose_name="Surname",
    )
    day_birth = models.DateField(
        verbose_name="Birth date",
    )
    deleted = models.BooleanField(default=False)

    def __str__(self) -> str:
        return "{0:0>3} {1} {2}".format(
            self.pk,
            self.name_second,
            self.name_first,
        )

    def delete(self, *args):
        self.deleted = True
        self.save()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Nov 4, 2020

@author: Enx
"""

#%% Importing and Variable fixing

from well_plate_project.config import data_dir, reportings_dir

# Dataset paths
raw_data_dir = data_dir / 'raw'
interm_data_dir = raw_data = data_dir / 'intermediate'
describe_data_dir = reportings_dir / 'description'

# Use KMeans clustering algorithm from sklearn.cluster to cluster pixels in image
from sklearn.cluster import KMeans

import numpy as np
import cv2
import matplotlib.pyplot as plt

image_name = 'a2_a_cropped.jpg'

image_path = raw_data_dir / image_name
assert image_path.is_file()

#%% Import image

original_image = cv2.imread(str(image_path))

img=original_image.astype('uint8')

plot_img_4hist_rgb(original_image)

plot_img_4hist_rgb(img)

#%% Convert image
# convert our image from RGB Colours Space to HSV to work ahead.
img_rgb=cv2.cvtColor(original_image.astype('uint8'),cv2.COLOR_BGR2RGB)
plot_img_4hist_rgb(img_rgb)

img_hsv=cv2.cvtColor(original_image.astype('uint8'),cv2.COLOR_BGR2HSV)
plot_img_4hist_hsv(img_hsv)
#https://raspberrypi.stackexchange.com/questions/10588/hue-saturation-intensity-histogram-plot

img_lab=cv2.cvtColor(original_image.astype('uint8'),cv2.COLOR_BGR2LAB)
plot_img_4hist_rgb(img_lab)

#%% Reshaping

# For clustering the image using k-means, we first need to convert it into a 2-dimensional array
image_2D = img.reshape(img.shape[0]*img.shape[1], img.shape[2])
reshaped_image = np.float32(original_image.reshape(-1, 3))
reshaped_image_hsv = np.float32(img_hsv.reshape(-1, 3))
reshaped_image_lab = np.float32(img_lab.reshape(-1, img_lab.shape[2]))

#%% Clustering 1
K = 10
stopCriteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.1)

ret, labels, clusters = cv2.kmeans(reshaped_image, K, None, stopCriteria, 10, cv2.KMEANS_RANDOM_CENTERS)
clusters = np.uint8(clusters)

intermediateImage = clusters[labels.flatten()]
clusteredImage = intermediateImage.reshape((original_image.shape))

cv2.imwrite("clusteredImage.jpg", clusteredImage)

plot_img_4hist_rgb(clusteredImage)

#%% Clustering2
# tweak the cluster size and see what happens to the Output
kmeans = KMeans(n_clusters=5, random_state=0).fit(image_2D)
clustered = kmeans.cluster_centers_[kmeans.labels_]

# Reshape back the image from 2D to 3D image
clustered_3D = clustered.reshape(img.shape[0], img.shape[1], img.shape[2])
plt.imshow(clustered_3D)
plt.title('Clustered Image')
plt.show()

plot_img_4hist_rgb(clustered_3D)

#%% Remove 1 cluster from image and apply canny edge detection
removedCluster = 1

cannyImage = np.copy(original_image).reshape((-1, 3))
cannyImage[labels.flatten() == removedCluster] = [0, 0, 0]

cannyImage = cv2.Canny(cannyImage,100,200).reshape(original_image.shape)
cv2.imwrite("cannyImage.jpg", cannyImage)
plot_img_4hist_rgb(cannyImage)

#%% Finding contours using opencv

initialContoursImage = np.copy(cannyImage)
imgray = cv2.cvtColor(initialContoursImage, cv2.COLOR_BGR2GRAY)
_, thresh = cv2.threshold(imgray, 50, 255, 0)
contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
cv2.drawContours(initialContoursImage, contours, -1, (0,0,255), cv2.CHAIN_APPROX_SIMPLE)
cv2.imwrite("initialContoursImage.jpg", initialContoursImage)

plot_img_4hist_rgb(initialContoursImage)

#DBscan
#https://stackoverflow.com/questions/40142835/image-not-segmenting-properly-using-dbscan
#https://core.ac.uk/download/pdf/79492015.pdf

#https://github.com/s4lv4ti0n/clusteringAlgorithms

# %% Define functions

def plot_img_4hist_rgb(image_to_plot):
    # Calulate various hists

    plt.figure(figsize=(10,10))

    #plt.subplot(221), plt.imshow(image_to_plot, 'gray')
    plt.subplot(221), plt.imshow(image_to_plot)

    plt.subplot(223), plt.hist(image_to_plot.ravel(),256,[0,256])
    plt.xlim([0,256])

    color = ('b','g','r')
    for i,col in enumerate(color):
        histr = cv2.calcHist([image_to_plot],[i],None,[256],[0,256])
        plt.subplot(222), plt.plot(histr,color = col)
        plt.yscale('log')
        plt.xlim([0,256])
        plt.subplot(224), plt.plot(histr,color = col)
        plt.xlim([0,256])

    plt.show()
    plt.close()
    return

# %%
def plot_img_4hist_hsv(image_to_plot):
    # Calulate various hists

    plt.figure(figsize=(10,10))

    #plt.subplot(221), plt.imshow(image_to_plot, 'gray')
    plt.subplot(221), plt.imshow(image_to_plot)

    hist = cv2.calcHist( [image_to_plot], [0, 1], None, [180, 256], [0, 180, 0, 256] )
    plt.subplot(223), plt.plot(hist)
    plt.subplot(222), plt.imshow(hist, interpolation = 'nearest')

    plt.show()
    plt.close()
    return
# %%


import re
import sys
import epitran
from nltk import TweetTokenizer
from scipy.stats import kurtosis, skew
from tqdm import tqdm
import numpy as np
from gensim.models import Word2Vec
from sklearn.base import BaseEstimator, TransformerMixin

from logic.linguistic_senticnet import LinguisticSenticNet
from logic.text_analysis import TextAnalysis
from logic.utils import Utils
from logic.lexical_features import lexical_es, lexical_en
from root import DIR_EMBEDDING

class FeatureExtraction(BaseEstimator, TransformerMixin):

    def __init__(self, lang='es', text_analysis=None):
        try:
            ta = None
            if text_analysis is None:
                ta = TextAnalysis(lang=lang)
            else:
                ta = text_analysis
            self.ta = ta
            file_syllable_embedding_en = DIR_EMBEDDING + 'syllable_embedding_en.model'
            file_syllable_embedding_es = DIR_EMBEDDING + 'syllable_embedding_es.model'
            file_phoneme_embedding_en = DIR_EMBEDDING + 'phoneme_embedding_en.model'
            file_phoneme_embedding_es = DIR_EMBEDDING + 'phoneme_embedding_es.model'
            print('Loading Lexicons and Embedding.....')
            if lang == 'es':
                epi = epitran.Epitran('spa-Latn')
                syllable_embedding = Word2Vec.load(file_syllable_embedding_es)
                phoneme_embedding = Word2Vec.load(file_phoneme_embedding_es)
            else:
                epi = epitran.Epitran('eng-Latn')
                syllable_embedding = Word2Vec.load(file_syllable_embedding_en)
                phoneme_embedding = Word2Vec.load(file_phoneme_embedding_en)

            self.epi = epi
            self.syllable_embedding = syllable_embedding
            self.phoneme_embedding = phoneme_embedding
            self.lexical = lexical_es if lang == 'es' else lexical_en
            self.lsn = LinguisticSenticNet(text_analysis=self.ta)
        except Exception as e:
            Utils.standard_error(sys.exc_info())
            print('Error FeatureExtraction: {0}'.format(e))

    def fit(self, x, y=None):
        return self

    def transform(self, list_messages):
        try:
            result = self.get_features(list_messages)
            return result
        except Exception as e:
            Utils.standard_error(sys.exc_info())
            print('Error transform: {0}'.format(e))

    def get_features(self, messages: str, type_features: list = [1, 1, 1, 1]):
        try:
            # L: Lexical, S:Syllable, F: Frequency Phoneme, P: All Phoneme
            syllable_features = list(abs(self.get_feature_syllable(messages))) if type_features[0] else []
            phoneme_frequency = list(abs(self.get_frequency_phoneme(messages))) if type_features[1] else []
            all_phoneme = list(abs(self.get_feature_phoneme(messages))) if type_features[2] else []
            lexical_features = list(abs(self.get_features_lexical(messages))) if type_features[3] else []
            features = lexical_features + syllable_features + all_phoneme + phoneme_frequency
            result = np.array(features, dtype=np.float32)
            return result
        except Exception as e:
            Utils.standard_error(sys.exc_info())
            print('Error get_features: {0}'.format(e))
            return None

    def get_feature_syllable(self, messages):
        try:
            messages_phonetic = []
            model = self.syllable_embedding
            num_features = model.vector_size
            index2phoneme_set = set(model.wv.index2word)
            num_phonemes = 1
            feature_vec = []
            list_syllable = [token['syllables'] for token in self.ta.tagger(messages) if token['syllables'] is not None]
            for syllable in list_syllable:
                for s in syllable:
                    syllable_phonetic = self.epi.transliterate(s, normpunc=True)
                    messages_phonetic.append(syllable_phonetic)
                    if syllable_phonetic in index2phoneme_set:
                        vec = model.wv[syllable_phonetic]
                        feature_vec.append(vec)
                        num_phonemes += 1
            feature_vec = np.array(feature_vec, dtype="float32")
            feature_vec = np.sum(feature_vec, axis=0)
            feature_vec = np.divide(feature_vec, num_phonemes)
            # print('Phonetic text: {0}'.format(messages_phonetic))
            # print('Embedding: {0}'.format(feature_vec))
            return feature_vec
        except Exception as e:
            Utils.standard_error(sys.exc_info())
            print('Error get_feature_syllable: {0}'.format(e))
            return None

    def get_frequency_phoneme(self, messages):
        try:
            messages_phonetic = []
            total_freq = 1
            model = self.syllable_embedding
            index2phoneme = list(model.wv.index2word)
            num_features = len(index2phoneme)
            feature_vec = np.zeros(num_features, dtype="float32")
            list_syllable = [token['syllables'] for token in self.ta.tagger(messages) if token['syllables'] is not None]
            for syllable in list_syllable:
                for s in syllable:
                    syllable_phonetic = self.epi.transliterate(s, normpunc=True)
                    messages_phonetic.append(syllable_phonetic)
                    if syllable_phonetic in index2phoneme:
                        index = index2phoneme.index(syllable_phonetic)
                        value = feature_vec[index]
                        feature_vec[index] = value + 1
                        total_freq += feature_vec[index]
            # print('Phonetic text: {0}'.format(messages_phonetic))
            # print('Frequency: {0}'.format(feature_vec))
            return feature_vec / total_freq
        except Exception as e:
            Utils.standard_error(sys.exc_info())
            print('Error get_frequency_phoneme: {0}'.format(e))
            return None

    def get_feature_phoneme(self, messages, one=False):
        try:
            messages_phonetic = None
            model = self.phoneme_embedding
            num_features = model.vector_size
            index2phoneme_set = set(model.wv.index2word)
            size = 1
            feature_vec = []
            list_syllable = [token['syllables'] for token in self.ta.tagger(messages) if token['syllables'] is not None]
            if one:
                try:
                    first_syllable = str(list_syllable[0][0])
                    first_syllable = first_syllable[0] if (first_syllable is not None) and (len(first_syllable) > 0) else ''
                    syllable_phonetic = self.epi.transliterate(first_syllable)
                    messages_phonetic = first_syllable
                    if syllable_phonetic in index2phoneme_set:
                        vec = model.wv[syllable_phonetic]
                        feature_vec.append(vec)
                    else:
                        feature_vec.append(np.zeros(num_features, dtype="float32"))
                except Exception as e_epi:
                    print('Error transliterate: {0}'.format(e_epi))
                    pass
            else:
                list_phoneme = self.epi.trans_list(messages)
                messages_phonetic = list_phoneme
                size = len(list_phoneme)
                for phoneme in list_phoneme:
                    if phoneme in index2phoneme_set:
                        vec = model.wv[phoneme]
                        feature_vec.append(vec)
                    else:
                        feature_vec.append(np.zeros(num_features, dtype="float32"))

            feature_vec = np.array(feature_vec, dtype="float32")
            feature_vec = np.sum(feature_vec, axis=0)
            feature_vec = np.divide(feature_vec, size)
            # print('Phonetic text: {0}'.format(messages_phonetic))
            # print('Embedding: {0}'.format(feature_vec))
            return feature_vec
        except Exception as e:
            Utils.standard_error(sys.exc_info())
            print('Error get_feature_phoneme: {0}'.format(e))
            return None

    def get_features_lexical(self, message):
        result = None
        try:
            lexical = self.lexical
            text_tokenizer = TweetTokenizer()
            tags = ('mention', 'url', 'hashtag', 'emoji', 'rt')
            vector = dict()
            vector['plarity'] = float(self.lsn.polarity_text(text=message)['polarity_value'])
            tokens_text = text_tokenizer.tokenize(message)
            if len(tokens_text) > 0:
                vector['weighted_position'], vector['weighted_normalized'] = self.weighted_position(tokens_text)

                vector['label_mention'] = float(sum(1 for word in tokens_text if word == 'mention'))
                vector['label_url'] = float(sum(1 for word in tokens_text if word == 'url'))
                vector['label_hashtag'] = float(sum(1 for word in tokens_text if word == 'hashtag'))
                vector['label_emoji'] = float(sum(1 for word in tokens_text if word == 'emoji'))
                vector['label_retweets'] = float(sum(1 for word in tokens_text if word == 'rt'))

                vector['lexical_diversity'] = self.lexical_diversity(message)

                label_word = vector['label_mention'] + vector['label_url'] + vector['label_hashtag']
                label_word = label_word + vector['label_emoji'] + vector['label_retweets']
                vector['label_word'] = float(len(tokens_text) - label_word)

                vector['first_person_singular'] = float(
                    sum(1 for word in tokens_text if word in lexical['first_person_singular']))
                vector['second_person_singular'] = float(
                    sum(1 for word in tokens_text if word in lexical['second_person_singular']))
                vector['third_person_singular'] = float(
                    sum(1 for word in tokens_text if word in lexical['third_person_singular']))
                vector['first_person_plurar'] = float(
                    sum(1 for word in tokens_text if word in lexical['first_person_plurar']))
                vector['second_person_plurar'] = float(
                    sum(1 for word in tokens_text if word in lexical['second_person_plurar']))
                vector['third_person_plurar'] = float(
                    sum(1 for word in tokens_text if word in lexical['third_person_plurar']))

                vector['avg_word'] = np.nanmean([len(word) for word in tokens_text if word not in tags])
                vector['avg_word'] = vector['avg_word'] if not np.isnan(vector['avg_word']) else 0.0
                vector['avg_word'] = round(vector['avg_word'], 4)

                vector['kur_word'] = kurtosis([len(word) for word in tokens_text if word not in tags])
                vector['kur_word'] = vector['kur_word'] if not np.isnan(vector['kur_word']) else 0.0
                vector['kur_word'] = round(vector['kur_word'], 4)

                vector['skew_word'] = skew(np.array([len(word) for word in tokens_text if word not in tags]))
                vector['skew_word'] = vector['skew_word'] if not np.isnan(vector['skew_word']) else 0.0
                vector['skew_word'] = round(vector['skew_word'], 4)

                # adverbios
                vector['adverb_neg'] = sum(1 for word in tokens_text if word in lexical['adverb_neg'])
                vector['adverb_neg'] = float(vector['adverb_neg'])

                vector['adverb_time'] = sum(1 for word in tokens_text if word in lexical['adverb_time'])
                vector['adverb_time'] = float(vector['adverb_time'])

                vector['adverb_place'] = sum(1 for word in tokens_text if word in lexical['adverb_place'])
                vector['adverb_place'] = float(vector['adverb_place'])

                vector['adverb_mode'] = sum(1 for word in tokens_text if word in lexical['adverb_mode'])
                vector['adverb_mode'] = float(vector['adverb_mode'])

                vector['adverb_cant'] = sum(1 for word in tokens_text if word in lexical['adverb_cant'])
                vector['adverb_cant'] = float(vector['adverb_cant'])

                vector['adverb_all'] = float(vector['adverb_neg'] + vector['adverb_time'] + vector['adverb_place'])
                vector['adverb_all'] = float(vector['adverb_all'] + vector['adverb_mode'] + vector['adverb_cant'])

                vector['adjetives_neg'] = sum(1 for word in tokens_text if word in lexical['adjetives_neg'])
                vector['adjetives_neg'] = float(vector['adjetives_neg'])

                vector['adjetives_pos'] = sum(1 for word in tokens_text if word in lexical['adjetives_pos'])
                vector['adjetives_pos'] = float(vector['adjetives_pos'])

                vector['who_general'] = sum(1 for word in tokens_text if word in lexical['who_general'])
                vector['who_general'] = float(vector['who_general'])

                vector['who_male'] = sum(1 for word in tokens_text if word in lexical['who_male'])
                vector['who_male'] = float(vector['who_male'])

                vector['who_female'] = sum(1 for word in tokens_text if word in lexical['who_female'])
                vector['who_female'] = float(vector['who_female'])

                vector['hate'] = sum(1 for word in tokens_text if word in lexical['hate'])
                vector['hate'] = float(vector['hate'])

                vector['noun'] = self.pos_frequency(message)['NOUN'] * 0.8
                vector['verb'] = self.pos_frequency(message)['VERB'] * 0.5
                vector['adj'] = self.pos_frequency(message)['ADJ'] * 0.4
                vector['pos_others'] = self.pos_frequency(message)['ANOTHER'] * 0.1

                result = np.array(list(vector.values()))
        except Exception as e:
            Utils.standard_error(sys.exc_info())
            print('Error get_lexical_features: {0}'.format(e))
        return result

    @staticmethod
    def lexical_diversity(text):
        result = None
        try:
            text_out = re.sub(r"[\U00010000-\U0010ffff]", '', text)
            text_out = re.sub(
                r'(?i)\b((?:https?://|www\d{0,3}[.]|[a-z0-9.\-]+[.][a-z]{2,4}/)(?:[^\s()<>]+|\(([^\s()<>]+'
                r'|(\([^\s()<>]+\)))*\))+(?:\(([^\s()<>]+|(\([^\s()<>]+\)))*\)|[^\s`!()\[\]{};:\'".,<>?«»“”‘’]))',
                '', text_out)
            text_out = text_out.lower()
            result = round((len(set(text_out)) / len(text_out)), 4)
        except Exception as e:
            Utils.standard_error(sys.exc_info())
            print('Error lexical_diversity: {0}'.format(e))
        return result

    @staticmethod
    def weighted_position(tokens_text):
        result = None
        try:
            size = len(tokens_text)
            weighted_words = 0.0
            weighted_normalized = 0.0
            for w in tokens_text:
                weighted_words += 1 / (1 + tokens_text.index(w))
                weighted_normalized += (1 + tokens_text.index(w)) / size
            result = (weighted_words, weighted_normalized)
        except Exception as e:
            Utils.standard_error(sys.exc_info())
            print('Error weighted_position: {0}'.format(e))
        return result

    def pos_frequency(self, text):
        dict_token = {'NOUN': 0, 'VERB': 0, 'ADJ': 0, 'ANOTHER': 0}
        try:
            doc = self.ta.tagger(text)
            for token in doc:
                if token['pos'] == 'NOUN':
                    value = dict_token['NOUN']
                    dict_token['NOUN'] = value + 1
                elif token['pos'] == 'VERB':
                    value = dict_token['VERB']
                    dict_token['VERB'] = value + 1
                elif token['pos'] == 'ADJ':
                    value = dict_token['ADJ']
                    dict_token['ADJ'] = value + 1
                else:
                    value = dict_token['ANOTHER']
                    dict_token['ANOTHER'] = value + 1
        except Exception as e:
            Utils.standard_error(sys.exc_info())
            print('Error pos_frequency: {0}'.format(e))
        return dict_token


#!/usr/bin/python
'''
reservation_container.py - Look up reservation data for a given facility
'''
import collections
import datetime
import logging
import re
import time

from pyvirtualdisplay import Display
from selenium import webdriver

SECONDS_SLEEP = 5
DRIVER_START_NUM_RETRIES = 20
WEBDRIVER_START_ERROR_STR = (
    'Could not start webdriver: facility {facility_id}, try {try_num}'
)
WEBDRIVER_START_FAILURE_STR = (
    'Failed to start webdriver for {facility_id} after {try_num} tries.'
)

display = Display(visible=0, size=(1600, 900))
display.start()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ReservationContainer(object):
    '''
    Contains data and methods required to look up a facility's reservation data
    '''
    RECREATION_GOV_URL = ('http://www.recreation.gov/campsiteCalendar.do?page='
                          'matrix&calarvdate=%s&contractCode=NRSO&parkId=%s')

    RECREATION_GOV_ROOT = 'http://www.recreation.gov%s'
    APOLOGY_STRING = 'Our apologies. We are experiencing some difficulties.'
    NEXT_PAGE_REGEX = r''

    RESERVATION_REGEX = (r'
(.*?)
	\n'
                         '<td class="status (.).*?<td class="status (.).*?'
                         '<td class="status (.).*?<td class="status (.).*?'
                         '<td class="status (.).*?<td class="status (.).*?'
                         '<td class="status (.).*?<td class="status (.).*?'
                         '<td class="status (.).*?<td class="status (.).*?'
                         '<td class="status (.).*?<td class="status (.).*?'
                         '<td class="status (.).*?<td class="status (.).*?'
                         '
')

    @staticmethod
    def initialize_webdriver(legacy_facility_id):
        '''
        Create a firefox selenium webdriver.  Retry a set number of times and
        raise Exception if driver cannot be created after that set number of
        retries.
        '''
        for i in range(0, DRIVER_START_NUM_RETRIES):
            try:
                driver = webdriver.Firefox()
            except Exception:
                logger.debug(
                    WEBDRIVER_START_ERROR_STR.format(
                        facility_id=legacy_facility_id,
                        try_num=i
                    )
                )
                time.sleep(SECONDS_SLEEP)
                continue

            return driver

        raise Exception(
            WEBDRIVER_START_FAILURE_STR.format(
                facility_id=legacy_facility_id,
                try_num=DRIVER_START_NUM_RETRIES
            )
        )

    def available_campsites_at_facility(self, legacy_facility_id, start_date,
                                        finish_date):
        ''' Return list of available campsites by dates for a recarea. '''
        driver = self.initialize_webdriver(legacy_facility_id)
        date_page_list = [start_date]
        while finish_date - date_page_list[-1] > datetime.timedelta(days=14):
            date_page_list.append(
                date_page_list[-1] + datetime.timedelta(days=14)
            )

        campsite_date_dict = {}
        for date in date_page_list:
            campsites_by_date = self.find_available_campsites_on_day(
                date,
                legacy_facility_id,
                driver
            )
            campsite_date_dict.update(campsites_by_date)

        keys_for_removal = []
        for date in campsite_date_dict:
            if date < start_date or date >= finish_date:
                keys_for_removal.append(date)

        for key in keys_for_removal:
            campsite_date_dict.pop(key)

        campsite_str_date_dict = {}
        for date in campsite_date_dict:
            date_str = date.isoformat()
            campsite_str_date_dict[date_str] = campsite_date_dict[date]

        driver.quit()
        return campsite_str_date_dict

    def find_available_campsites_on_day(self, date, park_id, driver):
        ''' Find all available campsites for given park on given day. '''
        date_string = '{0}/{1}/{2}'.format(date.month, date.day, date.year)
        next_page_url = self.RECREATION_GOV_URL % (date_string, park_id)
        campsites_by_date = {}
        while next_page_url:
            next_page_url = next_page_url.replace('&', '&')
            result_tuple = self.scrape_reservation_page(next_page_url,
                                                        date,
                                                        driver)

            (next_page_url, campsites_on_page) = result_tuple

            for this_date in campsites_on_page:
                if this_date in campsites_by_date:
                    existing_campsite_counter = collections.Counter(
                        campsites_by_date[this_date]
                    )
                    new_campsite_counter = collections.Counter(
                        campsites_on_page[this_date]
                    )
                    campsites_by_date[this_date] = dict(
                        existing_campsite_counter + new_campsite_counter
                    )
                else:
                    campsites_by_date[this_date] = campsites_on_page[this_date]

        driver.delete_all_cookies()
        return campsites_by_date

    def scrape_reservation_page(self, page_url, date, driver):
        ''' Scrape a single reservation url to get all available campsites. '''
        next_page_url = None
        campsites_by_date = {}
        driver.get(page_url)
        if driver.page_source.find(self.APOLOGY_STRING) == -1:
            reservation_match_list = re.findall(self.RESERVATION_REGEX,
                                                driver.page_source,
                                                re.DOTALL)

            for reservation_match in reservation_match_list:
                campsite_name = reservation_match[0]
                day_status_list = reservation_match[1:]
                for i, reservation_status in enumerate(day_status_list):
                    this_date = date + datetime.timedelta(days=i)
                    if ((reservation_status == 'a') or
                            (reservation_status == 'w')):
                        if this_date not in campsites_by_date:
                            campsites_by_date[this_date] = {}

                        if campsite_name not in campsites_by_date[this_date]:
                            campsites_by_date[this_date][campsite_name] = 1
                        else:
                            campsites_by_date[this_date][campsite_name] += 1

            next_page_match = re.search(self.NEXT_PAGE_REGEX,
                                        driver.page_source)

            if next_page_match and next_page_match.group(1):
                next_page_url = (self.RECREATION_GOV_ROOT
                                 % next_page_match.group(1))

        return (next_page_url, campsites_by_date)


from collections import Counter

#
# Complete the 'isValid' function below.
#
# The function is expected to return a STRING.
# The function accepts STRING s as parameter.
#

def is_valid_old(s):
    freq = [s.count(letter) for letter in set(s)]
    if max(freq) - min(freq) == 0:
        return 'YES'
    elif freq.count(max(freq)) == 1 and max(freq) - min(freq) == 1:
        return 'YES'
    elif freq.count(min(freq)) == 1:
        freq.remove(min(freq))
        if max(freq) - min(freq) == 0:
            return 'YES'
        else:
            return 'NO'
    else:
        return 'NO'

def is_valid(s):
    c = Counter(s)
    print(c)
    freq = Counter(c.values())
    print(freq)
    if len(freq) == 1:
        return "YES"
    elif len(freq) == 2:
        key_max = max(freq.keys())
        key_min = min(freq.keys())
        if key_max - key_min == 1 and freq[key_max] == 1:
            return "YES"
        elif key_min == 1 and freq[key_min] == 1:
            return "YES"
    return "NO"

if __name__ == '__main__':
    # fptr = open(os.environ['OUTPUT_PATH'], 'w')

    # s = input()
    # s = 'AAABBBCC'
    s = 'abcdefghhgfedecba'

    result = is_valid(s)

    # Expected Output YES
    print(result)


import os
from google.cloud import bigquery
class RsdGoogleBigquery:
    def __init__(self, key):
        self.key = key
        os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = key
        self.qclient = bigquery.Client()
    def RsdCreateBigqueryTable(self,table_dataset,table, schema):
        qclient = self.qclient
        table_id = table_dataset+'.'+table
        tables = qclient.list_tables(table_dataset)
        lst = []
        for rsd in tables:
            kmr = rsd.project+'.'+rsd.dataset_id+'.'+rsd.table_id
            lst.append(kmr)
        if table_id not in lst:
            print('Creating the table...')
            sma = open(schema, 'r')
            sma = sma.read()
            fl = []
            sc = []
            for i in range(len(sma.split('\n'))):
                if i%2==0:
                    fl.append(sma.split('\n')[i])
                else:
                    sc.append(sma.split('\n')[i])
            sch = []
            for i,j in zip(fl,sc):
                sch.append(bigquery.SchemaField(i,j, mode = 'NULLABLE'))
            table = bigquery.Table(table_id, schema = sch)
            table = qclient.create_table(table)
            print('Creatd table {}.{}.{}'.format(table.project, table.dataset_id, table.table_id))
        else:
            print('The table is already exist!')
        return

    def RsdGoogleBigqueryLoad(self,table_id, schema, file, disposition='WRITE_TRUNCATE'):
        qclient = self.qclient
        smp = open(schema, 'r')
        smp = smp.read()
        fls = []
        scc = []
        for i in range(len(smp.split('\n'))):
            if i%2==0:
                fls.append(smp.split('\n')[i])
            else:
                scc.append(smp.split('\n')[i])
        skd = []
        for i,j in zip(fls,scc):
            skd.append(bigquery.SchemaField(i,j, mode = 'NULLABLE'))

        job_config = bigquery.LoadJobConfig(
                     schema = skd,skip_leading_rows = 1,write_dispositon = disposition,
                              source_format = bigquery.SourceFormat.CSV,
                          )
        with open(file, 'rb') as rsk:
            job = qclient.load_table_from_file(rsk, table_id, job_config=job_config)
            print('Done uploading files to {0}'.format(table_id))
        job.result()
        destination_table = qclient.get_table(table_id)
        print('Loaded {} rows'. format(destination_table.num_rows))
        return


from typing import TYPE_CHECKING, Generator

import pygame
from itblib.Vec import IVector2
from itblib.globals.Constants import STANDARD_UNIT_SIZE
from itblib.gridelements.status_effect import (StatusEffectBase,
                                               StatusEffectBurrowed)
from itblib.gridelements.units.IUnitObserver import IUnitObserver
from itblib.net.NetEvents import NetEvents
from itblib.ui.widgets.game_widget import HealthBar, OwnerColorRhombus
from itblib.ui.IDisplayable import IDisplayable
from itblib.ui.TextureManager import Textures

from .GridElementUI import GridElementUI

if TYPE_CHECKING:
    from itblib.gridelements.units.UnitBase import UnitBase

class UnitBaseUI(GridElementUI, IDisplayable, IUnitObserver):
    def __init__(self, unit:"UnitBase", direction="SW", framespeed=.5):
        GridElementUI.__init__(self, parentelement=unit, direction=direction, framespeed=framespeed)
        self._parentelement:"UnitBase"
        self.fromscreenpos = None
        self.toscreenpos = None
        self.movementtime = 0.0
        self.speed = 0.0
        self.old_frame = self._framenumber
        self.healthbar = HealthBar(self._parentelement)
        self.healthbar.parent = self
        barwidth = self.healthbar.get_size()[0]
        offset = IVector2(int((STANDARD_UNIT_SIZE[0]-barwidth)/2), int(STANDARD_UNIT_SIZE[1]*.7))
        self.healthbar.position = offset
        self.owner_color_rhombus = None
        if NetEvents.session:
            self.owner_color_rhombus = OwnerColorRhombus(NetEvents.session._players[self._parentelement.ownerid].color)
            self.owner_color_rhombus.attach_to_unit(self)
        self._parentelement.observer = self

    def set_interp_movement(self, fromscreenpos: IVector2, toscreenpos: IVector2, speed: float):
        self.fromscreenpos = fromscreenpos
        self.toscreenpos = toscreenpos
        self.speed = speed
        self.movementtime = self._parentelement.age

    def update(self, delta_time: float):
        super().tick(delta_time)
        if self.fromscreenpos and self.toscreenpos:
            diff = self.toscreenpos - self.fromscreenpos
            time_ratio = min((self._parentelement.age - self.movementtime) / self.speed, 1)
            interp_screenpos = self.fromscreenpos + time_ratio * diff
            self.tfc.relative_position = interp_screenpos
            if time_ratio + 0e-4 >= 1:
                self.set_interp_movement(None, None, 0.0)
                self.movementtime = 0.0

    def get_display_name(self) -> str:
        return "UnitBaseUI"

    def get_display_description(self) -> str:
        return "Another long description ------ weee------"

    def get_blits(self) -> "Generator[tuple[pygame.Surface, pygame.Rect, pygame.Rect]]":
        if self.owner_color_rhombus:
            yield from self.owner_color_rhombus.get_blits()
        yield from GridElementUI.get_blits(self)
        yield from self.healthbar.get_blits()

    def on_add_status_effect(self, added_effect: "StatusEffectBase"):
        pass

    def on_remove_status_effect(self, removed_effect: "StatusEffectBase"):
        pass

    def on_update_position(self, new_pos: tuple[int, int]):
        pass

class UnitKnightUI(UnitBaseUI):
    def __init__(self, unit: "UnitBase"):
        super().__init__(unit, framespeed=.1)

    def get_display_name(self) -> str:
        return "Knight"

class UnitBurrowerUI(UnitBaseUI):
    def get_display_name(self) -> str:
        return "Burrower"

    def on_add_status_effect(self, added_effect: "StatusEffectBase"):
        if isinstance(added_effect, StatusEffectBurrowed):
            self.set_textures(Textures.get_spritesheet("BurrowerSWBurrowed"))
            self.frametime = .15

    def on_remove_status_effect(self, removed_effect: "StatusEffectBase"):
        if isinstance(removed_effect, StatusEffectBurrowed):
            self.set_textures(Textures.get_spritesheet("BurrowerSWIdle"))
            self.frametime = .5

class UnitChipmonkUI(UnitBaseUI):
    def __init__(self, unit: "UnitBase", direction="SW", framespeed=0.5):
        super().__init__(unit, direction, framespeed)
        self.moving = False

    def get_display_name(self) -> str:
        return "Chipmonk"

    def set_interp_movement(self, fromscreenpos: "tuple[int,int]", toscreenpos: "tuple[int,int]", speed: float):
        super().set_interp_movement(fromscreenpos, toscreenpos, speed)
        new_moving = bool(fromscreenpos) and bool(toscreenpos)
        print(new_moving)
        if not self.moving and new_moving:
            self.set_textures(Textures.get_spritesheet("ChipmonkSEMove"))
        elif self.moving and not new_moving:
            self.set_textures(Textures.get_spritesheet("ChipmonkSWIdle"))
        self.moving = new_moving


import fileinput
from collections import Counter
from fileinput import FileInput
from typing import List

def handler(raw_binary_numbers: FileInput) -> int:
    binary_numbers = [binary_number.strip() for binary_number in raw_binary_numbers]
    counters: List = [Counter() for _ in range(len(binary_numbers[0]))]

    for binary_number in binary_numbers:
        for i, counter in enumerate(counters):
            counter.update([binary_number[i]])

    gamma = []
    epsilon = []

    for counter in counters:
        most_common = counter.most_common()
        gamma.append(most_common[0][0])
        epsilon.append(most_common[1][0])

    return int("".join(gamma), 2) * int("".join(epsilon), 2)

if __name__ == "__main__":
    print(handler(fileinput.input()))