# 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 csv
import math
import logging
from pprint import pformat
from datetime import datetime
from cStringIO import StringIO
from mako.template import Template

import ssbench
import ssbench.util
from ssbench.ordered_dict import OrderedDict

REPORT_TIME_FORMAT = '%F %T UTC'

if hasattr(csv.DictWriter, "writeheader"):
    DictWriter = csv.DictWriter
else:
    class DictWriter(csv.DictWriter):

        def writeheader(self):
            header = dict(zip(self.fieldnames, self.fieldnames))
            self.writerow(header)

class Reporter(object):
    def __init__(self, run_results):
        self.run_results = run_results

    def read_results(self, nth_pctile=95, format_numbers=True):
        self.scenario, self.unpacker = self.run_results.read_results()
        self.stats = self.calculate_scenario_stats(nth_pctile, format_numbers)

    def write_rps_histogram(self, target_file):
        target_file.write('"Seconds Since Start","Requests Completed"\n')
        for i, req_count in enumerate(self.stats['time_series']['data'], 1):
            target_file.write('%d,%d\n' % (i, req_count))

    def scenario_template(self):
        return """
${scenario.name}  (generated with ssbench version ${scenario.version})
Worker count: ${'%3d' % agg_stats['worker_count']}   Concurrency: ${'%3d' % scenario.user_count}  Ran ${start_time} to ${stop_time} (${'%.0f' % round(duration)}s)
Object expiration (X-Delete-After): ${scenario.delete_after} (sec)

%% Ops    C   R   U   D       Size Range       Size Name
% for size_datum in size_data:
${size_datum['pct_total_ops']}   % ${size_datum['crud_pcts']}      ${size_datum['size_range']}  ${size_datum['size_name']}
% endfor
---------------------------------------------------------------------
        ${'%3.0f' % weighted_c} ${'%3.0f' % weighted_r} ${'%3.0f' % weighted_u} ${'%3.0f' % weighted_d}      CRUD weighted average

% for label, stats, sstats in stat_list:
% if stats['req_count']:
${label}
       Count: ${'%5d' % stats['req_count']} (${'%5d' % stats['errors']} error; ${'%5d' % stats['retries']} retries: ${'%5.2f' % stats['retry_rate']}%)  Average requests per second: ${'%5.1f' % stats['avg_req_per_sec']}
                            min       max      avg      std_dev  ${'%02d' % nth_pctile}%-ile  ${'%15s' % ''}  Worst latency TX ID
       First-byte latency: ${stats['first_byte_latency']['min']} - ${stats['first_byte_latency']['max']}  ${stats['first_byte_latency']['avg']}  (${stats['first_byte_latency']['std_dev']})  ${stats['first_byte_latency']['pctile']}  (all obj sizes)  ${stats['worst_first_byte_latency'][1] if 'worst_first_byte_latency' in stats else ''}
       Last-byte  latency: ${stats['last_byte_latency']['min']} - ${stats['last_byte_latency']['max']}  ${stats['last_byte_latency']['avg']}  (${stats['last_byte_latency']['std_dev']})  ${stats['last_byte_latency']['pctile']}  (all obj sizes)  ${stats['worst_last_byte_latency'][1] if 'worst_last_byte_latency' in stats else ''}
% for size_str, per_size_stats in sstats.iteritems():
% if per_size_stats:
       First-byte latency: ${per_size_stats['first_byte_latency']['min']} - ${per_size_stats['first_byte_latency']['max']}  ${per_size_stats['first_byte_latency']['avg']}  (${per_size_stats['first_byte_latency']['std_dev']})  ${per_size_stats['first_byte_latency']['pctile']}  ${'(%8s objs)' % size_str}  ${per_size_stats['worst_first_byte_latency'][1] if 'worst_first_byte_latency' in per_size_stats else ''}
       Last-byte  latency: ${per_size_stats['last_byte_latency']['min']} - ${per_size_stats['last_byte_latency']['max']}  ${per_size_stats['last_byte_latency']['avg']}  (${per_size_stats['last_byte_latency']['std_dev']})  ${per_size_stats['last_byte_latency']['pctile']}  ${'(%8s objs)' % size_str}  ${per_size_stats['worst_last_byte_latency'][1] if 'worst_last_byte_latency' in per_size_stats else ''}
% endif
% endfor

% endif
% endfor
Distribution of requests per worker-ID: ${jobs_per_worker_stats['min']} - ${jobs_per_worker_stats['max']} (avg: ${jobs_per_worker_stats['avg']}; stddev: ${jobs_per_worker_stats['std_dev']})
"""

    def generate_default_report(self, output_csv=False):
        """Format a default summary report based on calculated statistics for
        an executed scenario.

        :returns: A report (string) suitable for printing, emailing, etc.
        """

        stats = self.stats
        template = Template(self.scenario_template())
        tmpl_vars = {
            'size_data': [],
            'stat_list': [
                ('TOTAL', stats['agg_stats'], stats['size_stats']),
                ('CREATE', stats['op_stats'][ssbench.CREATE_OBJECT],
                 stats['op_stats'][ssbench.CREATE_OBJECT]['size_stats']),
                ('READ', stats['op_stats'][ssbench.READ_OBJECT],
                 stats['op_stats'][ssbench.READ_OBJECT]['size_stats']),
                ('UPDATE', stats['op_stats'][ssbench.UPDATE_OBJECT],
                 stats['op_stats'][ssbench.UPDATE_OBJECT]['size_stats']),
                ('DELETE', stats['op_stats'][ssbench.DELETE_OBJECT],
                 stats['op_stats'][ssbench.DELETE_OBJECT]['size_stats']),
            ],
            'agg_stats': stats['agg_stats'],
            'nth_pctile': stats['nth_pctile'],
            'start_time': datetime.utcfromtimestamp(
                stats['time_series']['start_time']
            ).strftime(REPORT_TIME_FORMAT),
            'stop_time': datetime.utcfromtimestamp(
                stats['time_series']['stop']).strftime(REPORT_TIME_FORMAT),
            'duration': stats['time_series']['stop']
            - stats['time_series']['start_time'],
            'jobs_per_worker_stats': stats['jobs_per_worker_stats'],
            'weighted_c': 0.0,
            'weighted_r': 0.0,
            'weighted_u': 0.0,
            'weighted_d': 0.0,
        }
        for size_data in self.scenario.sizes_by_name.values():
            if size_data['size_min'] == size_data['size_max']:
                size_range = '%-15s' % (
                    self._format_bytes(size_data['size_min']),)
            else:
                size_range = '%s - %s' % (
                    self._format_bytes(size_data['size_min']),
                    self._format_bytes(size_data['size_max']))
            initial_files = self.scenario._scenario_data['initial_files']
            initial_total = sum(initial_files.values())
            pct_total = (initial_files.get(size_data['name'], 0)
                         / float(initial_total) * 100.0)
            tmpl_vars['size_data'].append({
                'crud_pcts': '  '.join(map(lambda p: '%2.0f' % p,
                                           size_data['crud_pcts'])),
                'size_range': size_range,
                'size_name': size_data['name'],
                'pct_total_ops': '%3.0f%%' % pct_total,
            })
            tmpl_vars['weighted_c'] += \
                pct_total * size_data['crud_pcts'][0] / 100.0
            tmpl_vars['weighted_r'] += \
                pct_total * size_data['crud_pcts'][1] / 100.0
            tmpl_vars['weighted_u'] += \
                pct_total * size_data['crud_pcts'][2] / 100.0
            tmpl_vars['weighted_d'] += \
                pct_total * size_data['crud_pcts'][3] / 100.0
        if output_csv:
            csv_fields = [
                'scenario_name', 'ssbench_version', 'worker_count',
                'concurrency', 'start_time', 'stop_time', 'duration',
                'delete_after']
            csv_data = {
                'scenario_name': self.scenario.name,
                'ssbench_version': self.scenario.version,
                'worker_count': tmpl_vars['agg_stats']['worker_count'],
                'concurrency': self.scenario.user_count,
                'start_time': tmpl_vars['start_time'],
                'stop_time': tmpl_vars['stop_time'],
                'duration': tmpl_vars['duration'],
                'delete_after': str(self.scenario.delete_after),
            }
            for label, stats, sstats in tmpl_vars['stat_list']:
                label_lc = label.lower()
                if stats.get('req_count', 0):
                    self._add_csv_kv(csv_fields, csv_data,
                                     '%s_count' % label_lc, stats['req_count'])
                    self._add_csv_kv(csv_fields, csv_data,
                                     '%s_errors' % label_lc,
                                     stats['errors'])
                    self._add_csv_kv(csv_fields, csv_data,
                                     '%s_retries' % label_lc,
                                     stats['retries'])
                    self._add_csv_kv(csv_fields, csv_data,
                                     '%s_retry_rate' % label_lc,
                                     '%f' % stats['retry_rate'])
                    self._add_csv_kv(csv_fields, csv_data,
                                     '%s_avg_req_per_s' % label_lc,
                                     stats['avg_req_per_sec'])
                    self._add_stats_for(csv_fields, csv_data, label, 'all',
                                        stats, tmpl_vars['nth_pctile'])
                    for size_str, per_size_stats in sstats.iteritems():
                        if per_size_stats:
                            self._add_stats_for(csv_fields, csv_data, label,
                                                size_str, per_size_stats,
                                                tmpl_vars['nth_pctile'])
            csv_file = StringIO()
            csv_writer = DictWriter(csv_file, csv_fields,
                                    lineterminator='\n',
                                    quoting=csv.QUOTE_NONNUMERIC)
            csv_writer.writeheader()
            csv_writer.writerow(csv_data)
            return csv_file.getvalue()
        else:
            return template.render(scenario=self.scenario, **tmpl_vars)

    def _add_csv_kv(self, csv_fields, csv_data, key, value):
        csv_fields.append(key)
        csv_data[key] = value

    def _add_stats_for(self, csv_fields, csv_data, label, size_str, stats,
                       nth_pctile):
        for latency_type in ('first', 'last'):
            latency_stats = stats['%s_byte_latency' % latency_type]
            key_base = '%s_%s_%s_' % (label.lower(), latency_type, size_str)
            self._add_csv_kv(csv_fields, csv_data, key_base + 'min',
                             latency_stats['min'])
            self._add_csv_kv(csv_fields, csv_data, key_base + 'max',
                             latency_stats['max'])
            self._add_csv_kv(csv_fields, csv_data, key_base + 'avg',
                             latency_stats['avg'])
            self._add_csv_kv(csv_fields, csv_data, key_base + 'std_dev',
                             latency_stats['std_dev'])
            self._add_csv_kv(csv_fields, csv_data,
                             key_base + '%d_pctile' % nth_pctile,
                             latency_stats['pctile'])
            worst_key = 'worst_%s_byte_latency' % latency_type
            self._add_csv_kv(
                csv_fields, csv_data, key_base + 'worst_txid',
                stats[worst_key][1] if worst_key in stats else '')

    def _format_bytes(self, byte_count):
        units = [' B', 'kB', 'MB', 'GB']
        i = 0
        while round(byte_count / 1000.0, 3) >= 1.0:
            byte_count = byte_count / 1000.0
            i += 1
        return '%3.0f %s' % (round(byte_count), units[i])

    def calculate_scenario_stats(self, nth_pctile=95, format_numbers=True):
        """Compute various statistics from worker job result dicts.

        :param nth_pctile: Use this percentile when calculating the stats
        :param format_numbers: Should various floating-point numbers be
        formatted as strings or left full-precision floats
        :returns: A stats python dict which looks something like:
            SERIES_STATS = {
                'min': 1.1,
                'max': 1.1,
                'avg': 1.1,
                'std_dev': 1.1,
                'median': 1.1,
            }
            {
                'agg_stats': {
                    'worker_count': 1,
                    'start': 1.1,
                    'stop': 1.1,
                    'req_count': 1,
                    'retries': 0,
                    'errors' : 0,
                    'avg_req_per_sec': 1.1, # req_count / (stop - start)?
                    'retry_rate': 0.0,
                    'first_byte_latency': SERIES_STATS,
                    'last_byte_latency': SERIES_STATS,
                },
                'worker_stats': {
                    1: {  # keys are worker_ids
                        'start': 1.1,
                        'stop': 1.1,
                        'req_count': 1,
                        'retries': 0,
                        'retry_rate': 0.0,
                        'errors': 0,
                        'avg_req_per_sec': 1.1, # req_count / (stop - start)?
                        'first_byte_latency': SERIES_STATS,
                        'last_byte_latency': SERIES_STATS,
                    },
                    # ...
                },
                'op_stats': {
                    CREATE_OBJECT: { # keys are CRUD constants: CREATE_OBJECT, READ_OBJECT, etc.
                        'req_count': 1, # num requests of this CRUD type
                        'avg_req_per_sec': 1.1, # total_requests / sum(last_byte_latencies)
                        'first_byte_latency': SERIES_STATS,
                        'last_byte_latency': SERIES_STATS,
                        'size_stats': {
                            'small': { # keys are size_str values
                                'req_count': 1, # num requests of this type and size
                                'retries': 0, # num of retries
                                'avg_req_per_sec': 1.1, # total_requests / sum(last_byte_latencies)
                                'errors': 0,
                                'retry_rate': 0.0,
                                'first_byte_latency': SERIES_STATS,
                                'last_byte_latency': SERIES_STATS,
                            },
                            # ...
                        },
                    },
                    # ...
                },
                'size_stats': {
                    'small': { # keys are size_str values
                        'req_count': 1, # num requests of this size (for all CRUD types)
                        'retries': 0, # num of retries
                        'acutual_request_count': 1, # num requests includes retries
                        'avg_req_per_sec': 1.1, # total_requests / sum(last_byte_latencies)
                        'errors': 0,
                        'retry_rate': 0.0,
                        'first_byte_latency': SERIES_STATS,
                        'last_byte_latency': SERIES_STATS,
                    },
                    # ...
                },
                'time_series': {
                    'start': 1, # epoch time of first data point
                    'data': [
                        1, # number of requests finishing during this second
                        # ...
                    ],
                },
            }
        """
        # Each result looks like:
        # {
        #   'worker_id': 1,
        #   'type': 'get_object',
        #   'size': 4900000,
        #   'size_str': 'large',
        #   'first_byte_latency': 0.9137639999389648,
        #   'last_byte_latency': 0.913769006729126,
        #   'retries': 1
        #   'completed_at': 1324372892.360802,
        # }
        # OR
        # {
        #   'worker_id': 1,
        #   'type': 'get_object',
        #   'size_str': 'large'
        #   'completed_at': 1324372892.360802,
        #   'retries': 1
        #   'exception': '...',
        # }
        logging.info('Calculating statistics...')
        agg_stats = dict(start=2 ** 32, stop=0, req_count=0)
        op_stats = {}
        for crud_type in [ssbench.CREATE_OBJECT, ssbench.READ_OBJECT,
                          ssbench.UPDATE_OBJECT, ssbench.DELETE_OBJECT]:
            op_stats[crud_type] = dict(
                req_count=0, avg_req_per_sec=0,
                size_stats=OrderedDict.fromkeys(
                    self.scenario.sizes_by_name.keys()))

        req_completion_seconds = {}
        start_time = 0
        completion_time_max = 0
        completion_time_min = 2 ** 32
        stats = dict(
            nth_pctile=nth_pctile,
            agg_stats=agg_stats,
            worker_stats={},
            op_stats=op_stats,
            size_stats=OrderedDict.fromkeys(
                self.scenario.sizes_by_name.keys()))
        for results in self.unpacker:
            skipped = 0
            for result in results:
                try:
                    res_completed_at = result['completed_at']
                    res_completion_time = int(res_completed_at)
                    res_worker_id = result['worker_id']
                    res_type = result['type']
                    res_size_str = result['size_str']
                except KeyError as err:
                    logging.info('Skipped result with missing keys (%r): %r',
                                 err, result)
                    skipped += 1
                    continue

                try:
                    res_exception = result['exception']
                except KeyError:
                    try:
                        res_last_byte_latency = result['last_byte_latency']
                    except KeyError:
                        logging.info('Skipped result with missing'
                                     ' last_byte_latency key: %r',
                                     result)
                        skipped += 1
                        continue
                    if res_completion_time < completion_time_min:
                        completion_time_min = res_completion_time
                        start_time = (
                            res_completion_time - res_last_byte_latency)
                    if res_completion_time > completion_time_max:
                        completion_time_max = res_completion_time
                    req_completion_seconds[res_completion_time] = \
                        1 + req_completion_seconds.get(res_completion_time, 0)
                    result['start'] = res_completed_at - res_last_byte_latency
                else:
                    # report log exceptions
                    logging.warn('calculate_scenario_stats: exception from '
                                 'worker %d: %s',
                                 res_worker_id, res_exception)
                    try:
                        res_traceback = result['traceback']
                    except KeyError:
                        logging.warn('traceback missing')
                    else:
                        logging.info(res_traceback)

                # Stats per-worker
                if res_worker_id not in stats['worker_stats']:
                    stats['worker_stats'][res_worker_id] = {}
                self._add_result_to(stats['worker_stats'][res_worker_id],
                                    result)

                # Stats per-file-size
                try:
                    val = stats['size_stats'][res_size_str]
                except KeyError:
                    stats['size_stats'][res_size_str] = {}
                else:
                    if not val:
                        stats['size_stats'][res_size_str] = {}
                self._add_result_to(stats['size_stats'][res_size_str],
                                    result)

                self._add_result_to(agg_stats, result)

                type_stats = op_stats[res_type]
                self._add_result_to(type_stats, result)

                # Stats per-operation-per-file-size
                try:
                    val = type_stats['size_stats'][res_size_str]
                except KeyError:
                    type_stats['size_stats'][res_size_str] = {}
                else:
                    if not val:
                        type_stats['size_stats'][res_size_str] = {}
                self._add_result_to(
                    type_stats['size_stats'][res_size_str], result)
            if skipped > 0:
                logging.warn("Total number of results skipped: %d", skipped)

        agg_stats['worker_count'] = len(stats['worker_stats'].keys())
        self._compute_req_per_sec(agg_stats)
        self._compute_retry_rate(agg_stats)
        self._compute_latency_stats(agg_stats, nth_pctile, format_numbers)

        jobs_per_worker = []
        for worker_stats in stats['worker_stats'].values():
            jobs_per_worker.append(worker_stats['req_count'])
            self._compute_req_per_sec(worker_stats)
            self._compute_retry_rate(worker_stats)
            self._compute_latency_stats(worker_stats, nth_pctile,
                                        format_numbers)
        stats['jobs_per_worker_stats'] = self._series_stats(jobs_per_worker,
                                                            nth_pctile,
                                                            format_numbers)
        logging.debug('Jobs per worker stats:\n' +
                      pformat(stats['jobs_per_worker_stats']))

        for op_stats_dict in op_stats.itervalues():
            if op_stats_dict['req_count']:
                self._compute_req_per_sec(op_stats_dict)
                self._compute_retry_rate(op_stats_dict)
                self._compute_latency_stats(op_stats_dict, nth_pctile,
                                            format_numbers)
                for size_str, size_stats in \
                        op_stats_dict['size_stats'].iteritems():
                    if size_stats:
                        self._compute_req_per_sec(size_stats)
                        self._compute_retry_rate(size_stats)
                        self._compute_latency_stats(size_stats, nth_pctile,
                                                    format_numbers)
                    else:
                        op_stats_dict['size_stats'].pop(size_str)
        for size_str, size_stats in stats['size_stats'].iteritems():
            if size_stats:
                self._compute_req_per_sec(size_stats)
                self._compute_retry_rate(size_stats)
                self._compute_latency_stats(size_stats, nth_pctile,
                                            format_numbers)
            else:
                stats['size_stats'].pop(size_str)
        time_series_data = [req_completion_seconds.get(t, 0)
                            for t in range(completion_time_min,
                                           completion_time_max + 1)]
        stats['time_series'] = dict(start=completion_time_min,
                                    start_time=start_time,
                                    stop=completion_time_max,
                                    data=time_series_data)

        return stats

    def _compute_latency_stats(self, stat_dict, nth_pctile, format_numbers):
        try:
            for latency_type in ('first_byte_latency', 'last_byte_latency'):
                stat_dict[latency_type] = self._series_stats(
                    stat_dict.get(latency_type, []), nth_pctile,
                    format_numbers)
        except KeyError:
            logging.exception('stat_dict: %r', stat_dict)
            raise

    def _compute_req_per_sec(self, stat_dict):
        try:
            sd_start = stat_dict['start']
        except KeyError:
            stat_dict['avg_req_per_sec'] = 0.0
        else:
            delta_t = stat_dict['stop'] - sd_start
            if delta_t != 0:
                stat_dict['avg_req_per_sec'] = round(
                    stat_dict['req_count'] / delta_t,
                    6)
            else:
                stat_dict['avg_req_per_sec'] = float("inf")

    def _compute_retry_rate(self, stat_dict):
        stat_dict['retry_rate'] = round((float(stat_dict['retries']) /
                                         stat_dict['req_count']) * 100, 6)

    def _add_result_to(self, stat_dict, result):
        if 'errors' not in stat_dict:
            stat_dict['errors'] = 0
        try:
            res_start = result['start']
        except KeyError:
            pass
        else:
            try:
                sd_start = stat_dict['start']
            except KeyError:
                stat_dict['start'] = res_start
            else:
                if res_start < sd_start:
                    stat_dict['start'] = res_start
        try:
            sd_stop = stat_dict['stop']
        except KeyError:
            stat_dict['stop'] = result['completed_at']
        else:
            if result['completed_at'] > sd_stop:
                stat_dict['stop'] = result['completed_at']
        stat_dict['retries'] = \
            stat_dict.get('retries', 0) + int(result['retries'])
        if 'exception' not in result:
            stat_dict['req_count'] = stat_dict.get('req_count', 0) + 1
            self._rec_latency(stat_dict, result)
        else:
            stat_dict['errors'] += 1

    def _series_stats(self, sequence, nth_pctile, format_numbers):
        sequence = filter(None, sequence)
        if not sequence:
            # No data available
            return dict(min=' N/A  ', max='  N/A  ', avg='  N/A  ',
                        pctile='  N/A  ', std_dev='  N/A  ', median='  N/A  ')
        sequence.sort()
        try:
            minval = min(sequence)
            maxval = max(sequence)
            mean = ssbench.util.mean(sequence)
        except ZeroDivisionError:
            # Handle the case of a single-element sequence (population standard
            # deviation divides by N-1)
            minval = sequence[0]
            maxval = sequence[0]
            mean = sequence[0]
        if format_numbers:
            return dict(
                min='%6.3f' % minval,
                max='%7.3f' % maxval,
                avg='%7.3f' % mean,
                pctile='%7.3f' % self.pctile(sequence, nth_pctile),
                std_dev='%7.3f' % ssbench.util.uncorrected_stdev(sequence),
                median='%7.3f' % ssbench.util.median(sequence))
        else:
            return dict(
                min=round(minval, 6),
                max=round(maxval, 6),
                avg=round(mean, 6),
                pctile=round(self.pctile(sequence, nth_pctile), 6),
                std_dev=round(ssbench.util.uncorrected_stdev(sequence), 6),
                median=round(ssbench.util.median(sequence), 6))

    def pctile(self, sequence, nth_pctile):
        seq_len = len(sequence)
        rank = seq_len * nth_pctile / 100.0
        if float(int(rank)) == rank:
            # integer rank means we interpolate between two values
            rank = int(rank)
            return (sequence[rank - 1] + sequence[rank]) / 2.0
        else:
            return sequence[int(math.ceil(rank)) - 1]

    def _rec_latency(self, stats_dict, result):
        for latency_type in ('first_byte_latency', 'last_byte_latency'):
            if latency_type in stats_dict:
                stats_dict[latency_type].append(result[latency_type])
            else:
                stats_dict[latency_type] = [result[latency_type]]
            if result[latency_type] is not None:
                worst_key = 'worst_%s' % latency_type
                if worst_key not in stats_dict \
                        or result[latency_type] > stats_dict[worst_key][0]:
                    stats_dict[worst_key] = (round(result[latency_type], 6),
                                             result['trans_id'])


import torch
from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration

# Init is ran on server startup
# Load your model to GPU as a global variable here using the variable name "model"
def init():
    global model
    global tokenizer

    tokenizer = BlenderbotTokenizer.from_pretrained("facebook/blenderbot-3B")
    model = BlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-3B").to("cuda")

# Inference is ran for every server call
# Reference your preloaded global model variable here.
def inference(model_inputs:dict) -> dict:
    global model
    global tokenizer

    # Parse out your arguments
    prompt = model_inputs.get('prompt', None)
    if prompt == None:
        return {'message': "No prompt provided"}

    # Run the model
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    reply_ids = model.generate(**inputs)

    result = tokenizer.batch_decode(reply_ids)

    # Return the results as a dictionary
    return result


# ! /usr/env/bin python3
# -*-encoding utf-8-*-

"""
--------------------------------------------------------------------------------
    Stacked histogram
--------------------------------------------------------------------------------
Author: Simone Santoni, simone.santoni.1@city.ac.uk;
        credits to https://github.com/selva86
Dates: - created Thu Oct 24 11:06:17 UTC 2019
       - last change Thu Oct 24 11:06:17 UTC 2019
Notes: This script produces a histogram that shows the distribution of a
       focal variable conditionally on one or multiple categorical variables
"""

# %% libraries
import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rc
import pandas as pd

# %%  matplotlib viz setup
plt.style.use('seaborn-bright')
rc('font',**{'family':'sans-serif','sans-serif':['Avant Garde']})
rc('text', usetex=True)

# %% fake data
'''
Simulate three random variables drawn from different distributions
'''
# contoinuous vars
x = np.random.normal(loc=-1, scale=0.75, size=1000)
y = np.random.normal(loc=0, scale=1, size=1000)
z = np.random.normal(loc=1, scale=0.5, size=1000)

# qualitative categories
categories = ['x', 'y', 'z']

np.shape(x)

# %% create figure

# figure
fig = plt.figure(figsize=(6, 4))

# add figure
ax = fig.add_subplot(1, 1, 1)

# plot data
ax.hist([x, y, z], bins=50, histtype="barstacked", density=True,
         label=categories, color=["darkorange", "darkorchid", "lime"],
         alpha=0.6)

# axes
ax.set_xlabel("Values")
ax.set_ylabel("Density")

# legend
ax.legend(loc="best")

# save plot
out_f = os.path.join(os.getcwd(), 'stacked_histogram.pdf')
plt.savefig(out_f,
            transparent=True,
            bbox_inches='tight',
            pad_inches=0)

# show plot
plt.show()


from django.urls import path
from . import views

urlpatterns = [
    path('index', views.index_view, name='index_page'),
    path('course/<str:title>', views.topic_view, name='topic_page'),
    path('search', views.search_view, name='search_page'),
    path('new', views.new_page_view, name='new_page'),
    path('edit/<str:title>', views.edit_page_view, name='edit_page'),
    path('random', views.random_page_view, name='random_page'),
]


import unittest

import numpy as np

import ftlookup

class FtlookupTest(unittest.TestCase):
    def setUp(self):
        self.sut = ftlookup.FastTextWrapper()
        self.sut.loadModel('data/alice.bin')

    def test(self):
        with open('data/vectors.w2vt') as f:
            for line in f:
                word, *entries = line.split()

                expected_embedding = np.array(entries).astype(np.float)
                actual_embedding = self.sut.getVector(word)

                msg = 'Vector for [{0}] did not match! Expected: {1}, was {2}'.format(word, expected_embedding, actual_embedding)
                assert np.allclose(expected_embedding, actual_embedding, atol=0.000001), msg


import os
import pandas as pd
import csv
import requests as r
import datetime as dt
import re

# All Events Data

url="https://nuforc.org/webreports/ndxevent.html"

resp=r.get(url).text
html=resp.split("
')[0].strip()
    row['record_key']=each_record.split('>')[0].split('.html')[0].strip()
    row['record_date']=each_record.split('>')[1].strip('')[4].strip('</TD')
    records.append(row)

# df=pd.DataFrame(records)
# df

eventkeys=row.keys()

# All Events Data - truth checker
with open('all_events.csv','w',newline='') as recordsfile:
    dict_writer=csv.DictWriter(recordsfile, eventkeys)
    dict_writer.writeheader()
    dict_writer.writerows(records)

# All States

states = r.get(url).text
states_records=states.split("
<A HREF=")[1:]

# Alternative approach
# data = pd.read_html('https://nuforc.org/webreports/ndxe202303.html')
# data[0]
# :D


from django.shortcuts import render
from web.models import *
from django.shortcuts import render, get_object_or_404
from django.core.paginator import Paginator, EmptyPage

# Create your views here.
def home(request):
    stores = Store.objects.all()
    category= Category.objects.all()

    # Create a dictionary to store the first associated coupon for each store
    store_coupons = {}

    for store in stores:
        # Get the first coupon associated with the store, if any
        coupon = Coupon.objects.filter(store=store).first()

        # Add the store and its associated coupon (if any) to the dictionary
        store_coupons[store] = coupon

    context = {
        "stores": stores,
        "store_coupons": store_coupons,
        "categories": category
    }
    return render(request, 'index.html', context)

def stores(request):
    # Get all stores from the database and order them by name
    stores = Store.objects.order_by('name')

    # Create a list to store the grouped stores
    store_groups = []

    # Initialize variables to keep track of the current letter and stores for that letter
    current_letter = None
    current_stores = []

    for store in stores:
        first_letter = store.name[0].upper()

        if first_letter != current_letter:
            # If the letter changes, add the current stores to the list
            if current_letter is not None:
                store_groups.append((current_letter, current_stores))
            current_letter = first_letter
            current_stores = []

        current_stores.append(store)

    # Add the last group of stores
    if current_stores:
        store_groups.append((current_letter, current_stores))

    context = {
        'store_groups': store_groups
    }

    return render(request, 'stores.html', context)

def category(request):
    category= Category.objects.all()
    context = {
        'categories': category
    }
    return render(request, 'category.html', context=context)

def coupans(request, page=1):
    page = int(page)
    items_per_page = 15

    # Query all coupons
    all_coupons = Coupon.objects.filter(is_deal=False)

    # Get the count of code (where is_deal is False) and deal (where is_deal is True) coupons
    code_count = Coupon.objects.filter(is_deal=False).count()
    deal_count = Coupon.objects.filter(is_deal=True).count()

    # Get the total count of all coupons

    total_count = Coupon.objects.all().count()

    # Paginate all coupons
    paginator = Paginator(all_coupons, items_per_page)

    # Get the current page
    coupons = paginator.page(page)

    # If the current page is not the first page, get the coupons from the previous page
    if page > 1:
        previous_page = paginator.page(page - 1)
        previous_coupons = previous_page.object_list
    else:
        previous_coupons = []

    return render(request, 'coupans.html', {
        'coupons': coupons,
        'previous_coupons': previous_coupons,
        'code_count': code_count,
        'deal_count': deal_count,
        'total_count': total_count  # Pass the total count of all coupons to the template
    })

def dealCoupans(request, page=1):
    page = int(page)
    items_per_page = 15

    # Query all coupons
    all_coupons = Coupon.objects.filter(is_deal=True)

    # Get the count of code (where is_deal is False) and deal (where is_deal is True) coupons
    code_count = Coupon.objects.filter(is_deal=False).count()
    deal_count = Coupon.objects.filter(is_deal=True).count()

    # Get the total count of all coupons

    total_count = Coupon.objects.all().count()

    # Paginate all coupons
    paginator = Paginator(all_coupons, items_per_page)

    # Get the current page
    coupons = paginator.page(page)

    # If the current page is not the first page, get the coupons from the previous page
    if page > 1:
        previous_page = paginator.page(page - 1)
        previous_coupons = previous_page.object_list
    else:
        previous_coupons = []

    return render(request, 'dealCoupans.html', {
        'coupons': coupons,
        'previous_coupons': previous_coupons,
        'code_count': code_count,
        'deal_count': deal_count,
        'total_count': total_count  # Pass the total count of all coupons to the template
    })

def blog(request):
    return render(request, 'blog.html')

def about_us(request):
    return render(request, 'about.html')

def privacy_policy(request):
    return render(request, 'privacy-policy.html')

def contact_us(request):
    return render(request, 'contact_us.html')

def base(request):
    stores= Store.objects.all()

    context = {
        'stores': stores
    }
    return render(request, 'base.html', context=context)

def store_details(request,pk, page=1):
    page = int(page)
    items_per_page = 15

    # Query all coupons
    all_coupons = Coupon.objects.filter(store__id=pk)
    store = Store.objects.filter(id=pk)

    # Get the count of code (where is_deal is False) and deal (where is_deal is True) coupons
    code_count = all_coupons.filter(is_deal=False).count()
    deal_count = all_coupons.filter(is_deal=True).count()

    # Get the total count of all coupons

    total_count = all_coupons.count()

    # Paginate all coupons
    paginator = Paginator(all_coupons, items_per_page)

    # Get the current page
    coupons = paginator.page(page)
    categories = Category.objects.all()

    # If the current page is not the first page, get the coupons from the previous page
    if page > 1:
        previous_page = paginator.page(page - 1)
        previous_coupons = previous_page.object_list
    else:
        previous_coupons = []

    return render(request, 'store_details.html', {
        'coupons': coupons,
        'previous_coupons': previous_coupons,
        'code_count': code_count,
        'deal_count': deal_count,
        'total_count': total_count,
        'stores': store,
        'categories': categories
    })

def coupan_detail(request, pk):
    # Retrieve the coupon using the provided primary key
    coupon = get_object_or_404(Coupon, pk=pk)

    # Access the associated store's data
    store = coupon.store

    return render(request, 'coupan_detail.html', {'coupons': coupon, 'stores': store})

def deal_detail(request, pk):
    # Retrieve the coupon using the provided primary key
    coupon = get_object_or_404(Coupon, pk=pk)

    # Access the associated store's data
    store = coupon.store

    return render(request, 'deal_details.html', {'coupons': coupon, 'stores': store})


"""
Develop Model Evaluation Functions: model_eval.py
    Here are the helper functions for model evaluation
"""
##### Libraries #######################################################
import pandas as pd
import numpy as np
import sys
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import accuracy_score, precision_score, recall_score
from statistics import mean

# extend pirectory to get helper functions
sys.path.insert(0, "./helper_functions")

from reporting import dtype_selector

# pandas config
pd.options.mode.chained_assignment = None # stops slice warning

##### Functions and Classes to be called ###########################################
def cv_mse_stats(name, model, X, y):
    """
    Calculates mean accuracy, percision, and recall

    Parameters:
    ===========
    name: String name of the trial
    model: Sklearn algorithm
    X: Dataframe or Array of features
    y: Dataframe or Array of target values

    Returns:
    ========
    scoring: array[name, mean(accuracy), mean(percision), mean(recall)]
    """
    # container
    scoring = [name]

    for scoring_metric in ["accuracy", "precision", "recall"]:
        measurement = cross_val_score(model,
                                      X,
                                      y,
                                      scoring = scoring_metric,
                                      cv = 5,
                                      verbose = 0,
                                      n_jobs = 4)

        scoring.append(mean(measurement))

    return scoring

def stratified_cv(features, labels, cv = 5, threshold = 0.5, model = None):
    """
    Performs stratified cv testing

    Parameters:
    ===========
    features: DataFrame of features (Required)
    labels: Series of target labels (Required)
    cv: int
        the number of cross validations desired
    threshold: float
        the probability threshold for class prediction
    model: estimator
        Able to take a pretrained model, or use a standard logistic regression

    Returns:
    ========
    results: Dataframe that contains AUC metrics
    """
    # initialize the splitting
    sss = StratifiedKFold(n_splits = cv)

    # create containers
    accuracy = []
    precision = []
    recall = []

    # begin cv
    for train_index, test_index in sss.split(features, labels):
        # initialize logistic regressor
        if model == None:
            model_cl = LogisticRegression(max_iter = 400, random_state = 42)

        if model != None:
            model_cl = model

        # split data
        train_features = features.iloc[train_index, :]
        test_features = features.iloc[test_index, :]
        train_labels = labels.iloc[train_index]
        test_labels = labels.iloc[test_index]

        # train model
        model_cl.fit(train_features, train_labels)

        # predict
        pred_prob = model_cl.predict_proba(test_features)
        pred = [int(i > threshold) for i in pred_prob[:, 1]]

        # scoring
        accuracy.append(accuracy_score(test_labels, pred))
        precision.append(precision_score(test_labels, pred))
        recall.append(recall_score(test_labels, pred))

    results =  pd.DataFrame(data = {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall
    })

    return results

def threshold_testing(features, labels, model, model_name):
    """
    Tests the model across mulitple thresholds

    Parameters:
    ===========
    features: Dataframe or Array
    labels: Series, array, or list
    model: estimator
    model_name: String

    Returns:
    ========
    threshold_performance_df: Dataframe
        contains
    """
    threshold_performance = []
    for threshold_value in [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]:
        # get AUC measures for theshold
        results = stratified_cv(features,
                                labels,
                                threshold = threshold_value,
                                model = model)

        # organize results
        mean_results = results.mean(axis = 0)
        full_metric = [model_name,
                       threshold_value,
                       round(mean_results[0], 4),
                       round(mean_results[1], 4),
                       round(mean_results[2], 4)]

        # add to container
        threshold_performance = np.concatenate(
            (threshold_performance, full_metric),
            axis = 0
        )

    threshold_performance_df = pd.DataFrame(
        data = threshold_performance.reshape((8, 5)),
        columns = ["Model", "Threshold", "accuracy", "precision", "recall"]
    )

    return threshold_performance_df
def temp():
    """
    TEMP

    Parameters:
    ===========

    Returns:
    ========
    """
    return 0


#Thomas Pelowitz for CS368 on 2/14/23

#PROT
CODON = {
    "UUU": "F",
    "UUC": "F",
    "CUU": "L",
    "CUC": "L",
    "AUU": "I",
    "AUC": "I",
    "GUU": "V",
    "GUC": "V"
}

def prot(s: str) -> str:
    """
    Translates an RNA sequence into protein (single-letter notation).
    """

    current = ""
    output = ""

    for char in s:
        current += char

        if len(current) == 3:
            output += CODON[current]
            current = ""

    return output

def prot2 (s: str) -> str:
    number_of_codons = len(s) // 3

    protein = ""
    for i in range(number_of_codons):
        protein += CODON[s[i*3: (i*3)+3]]

    return protein


import sys, os
from standard_csv_reader2 import TheCSVReader
from video_segment_maker import MakeMovies
from cmd_creator import MakeCommand
import sys, os
from os import listdir
from os.path import isfile, join

VIDEO_PATH = 'FOLDER_NAME'
VIDEO_CLIPS_DIR = 'SUB_FOLDER_NAME'

clip_path = os.path.join(VIDEO_PATH, VIDEO_CLIPS_DIR)

def get_thumbnail(thumb_path=''):
    # GET FULL MOVIE

    # print("def get_movie : GET MOVIE FOLDER: ", self.cur_dir,  the_folder)
    for thumb in os.listdir(VIDEO_PATH):
        if thumb.endswith(".jpg"):
            the_thumbnail = thumb
            thumb_path = VIDEO_PATH + "\\" + the_thumbnail
            return thumb_path

def check_dir_files():
    print("CLIP PATH: ", clip_path)
    movie_clips = [f for f in listdir(clip_path) if isfile(join(clip_path, f))]
    return movie_clips

def write_com_file(self, the_ep='', the_file_name='', the_com=''):
    the_text_file = the_ep + "_cmds.txt"
    with open(the_text_file, 'w') as movie_cmds:
        movie_cmds.write(the_com)
        movie_cmds.write('\n')
        movie_cmds.write('\n')

if __name__ == "__main__":
    pass


from airflow import DAG
from airflow.operators.python import PythonOperator

from abc import abstractmethod
from datetime import datetime, timedelta

import pandas as pd
import numpy as np
from typing import Dict
import ticker_crawler
import pytz
from miner import Miner

class RSMiner(Miner):
    def __init__(self, input_tables: Dict[str, int], indicator_name: str, recursive_range: int):
        super().__init__(input_tables, indicator_name, recursive_range)

    def get_inputs(self, current_date: str, dict_input_dfs: Dict[str, pd.DataFrame], output_list: list):
        """
        :param output_list: list containing temp values of calculated indicator
        :param current_date: the date of current time in ISO Format; example: "2023-12-31T00:00"
        :param dict_input_dfs: dict of input dataframes with keys as input_table + "_data"
        :return: input variables for the formula
        """

        if not current_date.endswith(".000Z"):
            current_date += ".000Z"
        ticker_data = dict_input_dfs["ticker_data"]
        ticker_data = ticker_data.reset_index(drop=True)

        row_index_list = ticker_data.index[ticker_data['timestamp'] == current_date].tolist()
        if len(row_index_list) == 0:
            print("invalid timestamp")
            return None

        current_row_index = row_index_list[0]
        if current_row_index < 13:
            print(f"current_row_index is {current_row_index}, less than 13")
            return None

        close_list = (ticker_data.loc[current_row_index - 13:current_row_index, ["close"]])["close"].to_list()
        close_list = [int(close) for close in close_list]
        gain_list = []
        loss_list = []

        for i in range(1, len(close_list) - 1):
            diff = close_list[i] - close_list[i - 1]
            if diff >= 0:
                gain_list.append(diff)
            else:
                loss_list.append(diff)

        # time = ticker_data.loc[current_row_index, ["timestamp"]]
        #
        # if len(gain_list) == 0:
        #     print(f"gain_list of {time}: {len(gain_list)}")

        avg_gain = sum(gain_list) / len(gain_list) if len(gain_list) != 0 else 0
        avg_loss = sum(loss_list) / len(loss_list) if len(loss_list) != 0 else 0
        return avg_gain, avg_loss

    def formula(self, current_date: str, dict_input_dfs: Dict[str, pd.DataFrame], output_list: list):
        inputs = self.get_inputs(current_date, dict_input_dfs, output_list)
        if inputs is None:
            rs = 0
        else:
            avg_gain, avg_loss = inputs
            rs = avg_gain / np.abs(avg_loss) if avg_loss != 0 else 0

        output_list.append(rs)
        return output_list, rs

class RSIMiner(Miner):
    def __init__(self, input_tables: Dict[str, int], indicator_name: str, recursive_range: int):
        super().__init__(input_tables, indicator_name, recursive_range)

    def get_inputs(self, current_date: str, dict_input_dfs: Dict[str, pd.DataFrame], output_list: list):
        """
        :param output_list: list containing temp values of calculated indicator
        :param current_date: the date of current time in ISO Format; example: "2023-12-31T00:00"
        :param dict_input_dfs: dict of input dataframes with keys as input_table + "_data"
        :return: input variables for the formula
        """

        if not current_date.endswith(".000Z"):
            current_date += ".000Z"
        rs_data = dict_input_dfs["rs_data"]
        rs_data = rs_data.reset_index(drop=True)

        row_index_list = rs_data.index[rs_data['timestamp'] == current_date].tolist()
        if len(row_index_list) == 0:
            print("invalid timestamp")
            return None

        current_row_index = row_index_list[0]
        rs = rs_data.loc[current_row_index]["rs_data"]

        return rs

    def formula(self, current_date: str, dict_input_dfs: Dict[str, pd.DataFrame], output_list: list):
        rs = self.get_inputs(current_date, dict_input_dfs, output_list)
        if rs is None:
            rsi = 0
        else:
            rsi = 100 / (100 + rs)

        output_list.append(rs)
        return output_list, rsi

class OBVMiner(Miner):
    def __init__(self, input_tables, indicator_name, recursive_range):
        super().__init__(input_tables, indicator_name, recursive_range)

    def get_inputs(self, current_date: str, dict_input_dfs: Dict[str, pd.DataFrame], output_list: list):
        if not current_date.endswith(".000Z"):
            current_date += ".000Z"
        ticker_data = dict_input_dfs["ticker_data"]
        ticker_data = ticker_data.reset_index(drop=True)

        row_index_list = ticker_data.index[ticker_data['timestamp'] == current_date].tolist()
        if len(row_index_list) == 0:
            return None

        current_row_index = row_index_list[0]
        cur_close = float(ticker_data["close"][current_row_index])
        cur_volume = float(ticker_data["close"][current_row_index])
        prev_close = float(ticker_data["close"][current_row_index - 1]) if current_row_index > 0 else 0

        prev_1_obv = output_list[-1] if len(output_list) > 0 else None

        return cur_close, cur_volume, prev_close, prev_1_obv

    def formula(self, current_date: str, dict_input_dfs: Dict[str, pd.DataFrame], output_list: list):
        inputs = self.get_inputs(current_date, dict_input_dfs, output_list)
        if inputs is None:
            cur_obv = None
        else:
            cur_close, cur_volume, prev_close, prev_1_obv = inputs
            if prev_1_obv is None:
                cur_obv = 0
            else:
                cur_obv = prev_1_obv + np.sign(cur_close - prev_close) * cur_volume
        output_list.append(cur_obv)
        return output_list, cur_obv
# SECTION
tz = pytz.timezone('Asia/Ho_Chi_Minh')

def localize_utc_tz(d):
    return tz.fromutc(d)

indicator_dag = DAG(
    dag_id='indicator_dag',
    default_args={
        'retries': 1
    },
    description="DAG for stock indicators computation",
    schedule_interval='15 15 * * 1-5',
    start_date=datetime.today().replace(tzinfo=tz),
    tags=['indicator'],
    user_defined_filters={
        'local_tz': localize_utc_tz,
    }
)

crawler_task = PythonOperator(task_id="crawler", python_callable=ticker_crawler.main, dag=indicator_dag)

def compute_rs():
    print("RS Miner executing...")
    # rs_miner = RSMiner(input_tables={"ticker_data": 20}, indicator_name="rs", recursive_range=0)
    # rs_miner.execute()
    return "RS Miner finish"

rs_task = PythonOperator(task_id="rs", python_callable=compute_rs, dag=indicator_dag)

def compute_rsi():
    print("RSI Miner executing")
    # rsi_miner = RSIMiner(input_tables={"rs_data": 0}, indicator_name="rsi", recursive_range=0)
    # rsi_miner.execute()
    return "RSI Miner finish"

rsi_task = PythonOperator(task_id="rsi", python_callable=compute_rsi, dag=indicator_dag)

def compute_obv():
    obv_miner = OBVMiner(input_tables={"ticker_data":0}, indicator_name="obv",
                                       recursive_range=1)
    obv_miner.execute()
    print("OBV Miner executing...")
    return "OBV_MINER"

obv_task = PythonOperator(task_id="obv", python_callable=compute_obv, dag=indicator_dag)
# SECTION
crawler_task >> (rs_task)
rs_task >> rsi_task

obv_task


#!/usr/bin/env python3
import pandas as pd
from nltk.tokenize import word_tokenize
from nltk.stem import PorterStemmer
import nltk

nltk.download('averaged_perceptron_tagger')

porter = PorterStemmer()

bad_words = []
with open('../data/bad-words.txt') as f:
    for word in f:
        word = word.strip().lower()
        bad_words.append(word)
        bad_words.append(porter.stem(word))
bad_words = list(set(bad_words))

df = pd.read_csv('../data/labeled_data.csv')
hate_content = df.loc[df['label'] == 0].tweet.values
hate_content = ' '.join(hate_content)

hate_content_list = word_tokenize(hate_content)
hate_content_list = list(set(hate_content_list))
hate_content_list_stem = [porter.stem(word) for word in hate_content_list] + hate_content_list

new_bad_words = []
for word in bad_words:
    if word in hate_content_list_stem:
        new_bad_words.append(word)
        new_bad_words.append(porter.stem(word))
new_bad_words = list(set(new_bad_words))
print(new_bad_words)

words_tag = nltk.pos_tag(new_bad_words)
word_tag = {word: tag for word, tag in words_tag}
tags = set(word_tag.values())
tags = ['<SWEAR-' + tag + '>' for tag in tags]
print(word_tag)

with open ('bad_words_in_hate.txt','a') as file:
    for word in new_bad_words:
        file.write(word+'\n')


def isValid(s):
    counts = {}
    for c in s:
        if c not in counts:
            counts[c] = 0
        counts[c] = counts[c] + 1

    uniqueCounts = {}
    for char, count in counts.items():
        if count not in uniqueCounts:
            uniqueCounts[count] = 0
        uniqueCounts[count] = uniqueCounts[count] + 1

    lengthOfUniqueCounts = len(uniqueCounts)
    if lengthOfUniqueCounts == 1:
        return "YES"
    if lengthOfUniqueCounts > 2:
        return "NO"

    firstNumber = None
    firstNumberCount = None
    secondNumber = None
    secondNumberCount = None
    for count, numberOfCount in uniqueCounts.items():
        if not firstNumber:
            firstNumber = count
            firstNumberCount = numberOfCount
        else:
            secondNumber = count
            secondNumberCount = numberOfCount

    biggerNumber = None
    biggerNumberCount = None
    smallerNumber = None
    smallerNumberCount = None

    if firstNumber > secondNumber:
        biggerNumber = firstNumber
        biggerNumberCount = firstNumberCount
        smallerNumber = secondNumber
        smallerNumberCount = secondNumberCount
    else:
        biggerNumber = secondNumber
        biggerNumberCount = secondNumberCount
        smallerNumber = firstNumber
        smallerNumberCount = firstNumberCount

    if biggerNumberCount > 1:
        if smallerNumber == 1 and smallerNumberCount == 1:
            return "YES"
        return "NO"

    if biggerNumber - 1 > smallerNumber:
        return "NO"

    return "YES"

def test0():
    print(isValid("aabbcd"))

def test1():
    print(isValid("aaaabbcc"))

def test2():
    print(isValid("aaaabbbcc"))

def test3():
    print(isValid("aaaabbbcccc"))

def test4():
    print(isValid("aacbb"))

test0()
test1()
test2()
test3()
test4()


import tensorflow as tf
import numpy as np
import config.config_param as cfg

anchors = cfg.anchors
anchor_masks = cfg.anchor_masks

def yolo_head(y_pred, anchors, calc_loss=False):
    """
    yolo_head split raw prediction from yolo_body into x,y,w,h and bounding box then process with normalization
    and create grid of three sizes
    :param y_pred: raw prediction from yolo_body
    :param anchors: three different anchor box from cfg
    :param calc_loss: True when train, False when predict
    :return: bounding box from prediction and grid
    """
    grid_size = tf.shape(y_pred)[1]

    # reshape_pred: (batch_size, grid, grid, anchors, (x, y, w, h, obj, ...classes))

    box_xy, box_wh, confidence, class_probs = tf.split(y_pred, (2, 2, 1, cfg.num_classes), axis=-1)
    # i.e. box_xy.shape = batch_size, grid, grid, num_bbox, xy
    # (8, 13, 13, 3, 2)
    # (8, 26, 26, 3, 2)
    # (8, 52, 52, 3, 2)

    # sigmoid let result lie between [0,1]
    box_xy = tf.sigmoid(box_xy)
    confidence = tf.sigmoid(confidence)
    class_probs = tf.sigmoid(class_probs)

    # create grid

    grid_y = tf.tile(tf.reshape(tf.range(grid_size), [-1, 1, 1, 1]), [1, grid_size, 1, 1])
    grid_x = tf.tile(tf.reshape(tf.range(grid_size), [1, -1, 1, 1]), [grid_size, 1, 1, 1])
    grid = tf.concat([grid_x, grid_y], axis=-1)
    grid = tf.cast(grid, tf.float32)

    # -----------------------------normalization x,y,w,h-------------------------------------------------------------
    # cell is offset from the top left corner of the image by (c x , c y ) and the bounding box prior has width and
    # height p w , p h
    # bx = σ(tx ) + cx
    # by = σ(ty ) + cy
    # bw = pwe^tw
    # bh = phe^th
    # ----------------------------------------------------------------------------------------------------------------
    box_xy = (box_xy + grid) / tf.cast(grid_size, tf.float32)
    box_wh = tf.exp(box_wh) * anchors / cfg.input_shape

    pred_xywh = tf.concat([box_xy, box_wh], axis=-1)  # original xywh for loss

    if calc_loss:
        return pred_xywh, grid

    return box_xy, box_wh, confidence, class_probs


import torch.nn as nn
import torch.nn.functional as F
import resnet as resnet_mod
import pytorch_cifar.models.resnet as resnet
import classifier_models.preact_resnet as preact_resnet
from classifier_models import PreActResNet18, ResNet18
import collections

class BasicBlockNoReLU(nn.Module):
    expansion = 1

    def __init__(self, module):
        super().__init__()
        self.conv1 = module.conv1
        self.bn1 = module.bn1
        self.conv2 = module.conv2
        self.bn2 = module.bn2
        self.shortcut = module.shortcut

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        return out

class SequentialImageNetwork(nn.Sequential):
    def __init__(self, net=resnet.ResNet18()):
        if isinstance(net, collections.OrderedDict):
            super().__init__(net)
            return

        self.net_holder = (net,)
        i = 1
        layers = []
        name = f"layer{i}"
        while hasattr(net, name):
            print(name)
            layers.extend(list(getattr(net, name)))
            i += 1
            name = f"layer{i}"

        layers2 = []
        for layer in layers:
            if isinstance(layer, resnet.BasicBlock):
                layers2.append(BasicBlockNoReLU(layer))
                layers2.append(nn.ReLU())
            else:
                layers2.append(layer)

        super().__init__(
            net.conv1,
            net.bn1,
            nn.ReLU(),
            net.maxpool,
            *layers2,
            net.avgpool,
            nn.Flatten(),
            net.fc
        )

    @property
    def net(self):
        return self.net_holder[0]

class BasicBlockNoReLUpre(nn.Module):
    expansion = 1

    def __init__(self, module):
        super().__init__()
        self.conv1 = module.conv1
        self.bn1 = module.bn1
        self.conv2 = module.conv2
        self.bn2 = module.bn2
        if hasattr(module, "shortcut"):
            self.shortcut = module.shortcut

    def forward(self, x):
        out = F.relu(self.bn1(x))
        shortcut = self.shortcut(out) if hasattr(self, "shortcut") else x
        out = self.conv1(out)
        out = self.conv2(F.relu(self.bn2(out)))
        out += shortcut
        return out

class SequentialImageNetwork_pre(nn.Sequential):
    def __init__(self, net=PreActResNet18(10)):
        if isinstance(net, collections.OrderedDict):
            super().__init__(net)
            return

        self.net_holder = (net,)
        i = 1
        layers = []
        name = f"layer{i}"
        while hasattr(net, name):
            layers.extend(list(getattr(net, name)))
            i += 1
            name = f"layer{i}"

        layers2 = []
        for layer in layers:
            if isinstance(layer, preact_resnet.PreActBlock):
                layers2.append(BasicBlockNoReLUpre(layer))
            else:
                layers2.append(layer)

        super().__init__(
            net.conv1,
            *layers2,
            nn.AvgPool2d(8 if net.in_planes == 64 else 4),
            nn.Flatten(),
            net.linear,
        )

    @property
    def net(self):
        return self.net_holder[0]

class SequentialImageNetwork_imgnet(nn.Sequential):
    def __init__(self, net=PreActResNet18(10)):
        if isinstance(net, collections.OrderedDict):
            super().__init__(net)
            return

        self.net_holder = (net,)
        i = 1
        layers = []
        name = f"layer{i}"
        while hasattr(net, name):
            layers.extend(list(getattr(net, name)))
            i += 1
            name = f"layer{i}"

        layers2 = []
        for layer in layers:
            if isinstance(layer, preact_resnet.PreActBlock):
                layers2.append(BasicBlockNoReLUpre(layer))
            else:
                layers2.append(layer)

        super().__init__(
            net.conv1,
            net.bn1,
            net.relu,
            net.maxpool,
            *layers2,
            net.avgpool,
            nn.Flatten(),
            net.fc,
        )

    @property
    def net(self):
        return self.net_holder[0]

class BasicBlockSplitter(nn.Module):
    def __init__(self, block: resnet.BasicBlock, step="add"):
        super().__init__()
        self.block = block
        self.step = step

    def forward(self, x):
        if self.step == "identity":
            return x
        shortcut = self.block.shortcut(x)
        if self.step == "shortcut":
            return shortcut
        x = self.block.conv1(x)
        if self.step == "conv1":
            return x
        x = self.block.bn1(x)
        if self.step == "bn1":
            return x
        x = F.relu(x)
        if self.step == "relu1":
            return x
        x = self.block.conv2(x)
        if self.step == "conv2":
            return x
        x = self.block.bn2(x)
        if self.step == "bn2":
            return x
        x += shortcut
        if self.step == "add":
            return x
        x = F.relu(x)
        if self.step == "relu2":
            return x
        return x

class SequentialImageNetworkMod(nn.Sequential):
    def __init__(self, net=resnet_mod.resnet32()):
        if isinstance(net, collections.OrderedDict):
            super().__init__(net)
            return

        self.net_holder = (net,)
        i = 1
        layers = []
        name = f"layer{i}"
        while hasattr(net, name):
            layers.extend(list(getattr(net, name)))
            i += 1
            name = f"layer{i}"

        super().__init__(
            net.conv1,
            *layers,
            net.final_bn,
            nn.LeakyReLU(0.1),
            nn.AvgPool2d(8),
            nn.Flatten(),
            net.linear,
        )

    @property
    def net(self):
        return self.net_holder[0]


import networkx as nx
from productions.decorators import p13_isomorphism
from typing import Dict

class P13:
    left = nx.Graph()
    left.add_node(1, label='E')
    left.add_node(2, label='i')
    left.add_node(3, label='i')
    left.add_node(4, label='I')
    left.add_node(5, label='I')
    left.add_node(6, label='I')
    left.add_node(7, label='E')
    left.add_node(8, label='E')
    left.add_node(9, label='E')
    left.add_node(10, label='E')
    left.add_node(11, label='E')
    left.add_edge(1, 2)
    left.add_edge(1, 3)
    left.add_edge(2, 4)
    left.add_edge(2, 5)
    left.add_edge(3, 6)
    left.add_edge(4, 7)
    left.add_edge(4, 8)
    left.add_edge(5, 8)
    left.add_edge(5, 9)
    left.add_edge(6, 10)
    left.add_edge(6, 11)
    left.add_edge(7, 8)
    left.add_edge(8, 9)
    left.add_edge(10, 11)

    @staticmethod
    @p13_isomorphism(left, all_isomorphisms=True)
    def apply(G: nx.Graph, isomorphisms: Dict = None, options: Dict = {"apply": None}):
        if not isomorphisms or len(isomorphisms) == 0:
            return False

        isomorphism = None
        if 'apply' in options and options['apply'] is not None:
            for i in isomorphisms:
                if options['apply'](i):
                    isomorphism = i
            if isomorphism is None:
                return False
        else:
            isomorphism = isomorphisms[0]

        nodes_in_G = list(isomorphism.keys())
        E_nodes = []
        E_coeff = []

        for node in nodes_in_G:
            if G.nodes[node]['label'] == 'E':
                adjacency = G.adj[node]
                is_connected_to_i = False
                for neighbour in adjacency:
                    if G.nodes[neighbour]['label'] == 'i' and neighbour in nodes_in_G:
                        is_connected_to_i = True
                        break
                if not is_connected_to_i:
                    if tuple(G.nodes[node]['pos']) not in E_coeff:
                        E_coeff.append(tuple(G.nodes[node]['pos']))
                    E_nodes.append(node)

        E_coeff = sorted(E_coeff, key=lambda x: (x[0], x[1]))

        to_remove_1 = None
        to_remove_2 = None

        for node in E_nodes:
            adjacency = G.adj[node]
            pos = tuple(G.nodes[node]['pos'])
            if pos == E_coeff[0] or pos == E_coeff[2]:
                to_remove = True
                for neighbour in adjacency:
                    if G.nodes[neighbour]['label'] == 'E' and tuple(G.nodes[neighbour]['pos']) == E_coeff[1] and neighbour in E_nodes:
                        to_remove = False
                if to_remove:
                    to_remove_1 = (node, tuple(G.nodes[node]['pos']))
                    for neighbour in adjacency:
                        if G.nodes[neighbour]['label'] == 'E' and neighbour in nodes_in_G:
                            to_remove_2 = (neighbour, tuple(G.nodes[neighbour]['pos']))
                    break

        to_remove_1_neighbours = G.adj[to_remove_1[0]]
        to_remove_2_neighbours = G.adj[to_remove_2[0]]

        E_nodes.remove(to_remove_1[0])
        E_nodes.remove(to_remove_2[0])
        G.remove_node(to_remove_1[0])
        G.remove_node(to_remove_2[0])

        for node in E_nodes:
            pos = tuple(G.nodes[node]['pos'])
            if pos == to_remove_1[1]:
                for neighbour in to_remove_1_neighbours:
                    if neighbour != to_remove_2[0]:
                        G.add_edge(node, neighbour)
            elif pos == to_remove_2[1]:
                for neighbour in to_remove_2_neighbours:
                    if neighbour != to_remove_1[0]:
                        G.add_edge(node, neighbour)

        return True


import socket
import textwrap
from tempfile import NamedTemporaryFile

import requests
from matplotlib.image import imread
from PIL import Image, ImageDraw, ImageFont

def is_connected(url="http://www.google.com/", timeout=3):
    try:
        requests.head(url, timeout=timeout)
        return True
    except requests.ConnectionError as ex:
        print(ex)
        return False

def get_ip():
    """
    Retrieve primary IP address on the local box
    see https://stackoverflow.com/a/28950776/91468
    :return:
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        # doesn't even have to be reachable
        s.connect(("8.8.8.8", 1))
        ip = s.getsockname()[0]
    except Exception:
        ip = "127.0.0.1"
    finally:
        s.close()
    return ip

def get_hostname():
    return socket.gethostname()

def number_to_string(number):
    """
    Transform a number to a string. If number is smaller than 1000, displays the last 5 significant figures.
    :param number: number to convert to string
    :return: string representation of the number
    """
    number_string = ""
    if number > 100:
        number_string = format(int(number), ",")
    else:
        number_string = str(float("%.5g" % number))
    return number_string

def get_percentage_diff(current, previous):
    if current == previous:
        return 0
    try:
        return (float(current - previous) / previous) * 100.0
    except ZeroDivisionError:
        return float("inf")

# see https://kavigupta.org/2019/05/18/Setting-the-size-of-figures-in-matplotlib/
def get_size(fig, dpi=100):
    with NamedTemporaryFile(suffix=".png") as f:
        fig.savefig(f.name, bbox_inches="tight", dpi=dpi)
        height, width, _channels = imread(f.name).shape
        return width / dpi, height / dpi

def set_size(fig, size, dpi=100, eps=1e-2, give_up=2, min_size_px=10):
    target_width, target_height = size
    set_width, set_height = target_width, target_height  # reasonable starting point
    deltas = []  # how far we have
    while True:
        fig.set_size_inches([set_width, set_height])
        actual_width, actual_height = get_size(fig, dpi=dpi)
        set_width *= target_width / actual_width
        set_height *= target_height / actual_height
        deltas.append(
            abs(actual_width - target_width) + abs(actual_height - target_height)
        )
        if deltas[-1] < eps:
            return True
        if len(deltas) > give_up and sorted(deltas[-give_up:]) == deltas[-give_up:]:
            return False
        if set_width * dpi < min_size_px or set_height * dpi < min_size_px:
            return False

def empty_image(width, height):
    """
    Returns an empty canvas/image to draw  on
    :return:
    """
    return Image.new("1", (width, height), 255)  # 255: clear the frame

def place_text(
    img,
    text,
    x_offset=0,
    y_offset=0,
    font_size=40,
    font_name="Forum-Regular",
    fill=0,
):
    """
    Put some text at a location on the image.
    """
    draw = ImageDraw.Draw(img)
    font = ImageFont.truetype("/usr/share/fonts/TTF/DejaVuSans.ttf", font_size)
    draw.text((x_offset, y_offset), text, font=font, fill=fill)

def place_centered_text(
    img,
    text,
    x_offset=0,
    y_offset=0,
    font_size=40,
    font_name="Forum-Regular",
    fill=0,
):
    """
    Put some centered text at a location on the image.
    """
    font = ImageFont.truetype("/usr/share/fonts/TTF/DejaVuSans.ttf", font_size)
    img_width, img_height = img.size
    text_width, _ = font.getsize(text)
    text_height = font_size
    draw_x = (img_width - text_width) // 2 + x_offset
    draw_y = (img_height - text_height) // 2 + y_offset
    place_text(img, text, draw_x, draw_y, font_size, font_name, fill)

def place_text_right(
    img,
    text,
    x_offset=0,
    y_offset=0,
    font_size=40,
    font_name="Forum-Regular",
    fill=0,
):
    font = ImageFont.truetype("/usr/share/fonts/TTF/DejaVuSans.ttf", font_size)
    img_width, img_height = img.size
    text_width, _ = font.getsize(text)
    draw_x = (img_width - text_width) + x_offset
    draw_y = y_offset
    place_text(img, text, draw_x, draw_y, font_size, font_name, fill)

def write_wrapped_lines(
    img,
    text,
    font_size=16,
    y_text=20,
    height=15,
    width=25,
    font_name="Roboto-Light",
):
    lines = textwrap.wrap(text, width)
    num_lines = 0
    for line in lines:
        place_centered_text(img, line, 0, y_text, font_size, font_name)
        y_text += height
        num_lines += 1
    return img