
import cantera           as ct
import numpy             as np
import random            as rnd
import matplotlib.pyplot as plt
from scipy import integrate

#==========================================================================

class pmsr :

    #--------------------------------------------------------------------------

    def __init__(s) :

        s.np     = 1000          # number of particles
        s.dt     = 0.1E-3        # step size for inlet/outlet/pairing
        s.tau_r  = 10.E-3        # rxr residence time
        s.tau_m  = 1.E-3         # mixing timescale (~micromixing, ~diffusive)
        s.tau_p  = 1.E-3         # pairing timescale (~macromixing, ~convective)
        s.ntau_r = 5.0           # number of run times to simulate for
        s.mech   = 'h2o2.yaml'   # chemical mechanism
        s.mixf   = 0.1111        # fuel/(fuel+air) stream ratio (mass basis)
        s.X0     = 'O2:0.5'
        s.X1     = 'H2:1.0'
        s.T0     = 300.0
        s.T1     = 300.0
        s.P      = 101325.0

        #----------- 

        s.trun   = s.ntau_r*s.tau_r   # run time
        s.nsteps = int(s.trun / s.dt) 

        #----------- set the gas and streams

        s.gas = ct.Solution(s.mech)
        s.gas.TPX = s.T0,s.P,s.X0
        s.Y0 = s.gas.Y
        s.h0 = s.gas.enthalpy_mass
        s.gas.TPX = s.T1,s.P,s.X1
        s.Y1 = s.gas.Y
        s.h1 = s.gas.enthalpy_mass

        #----------- set the initial state

        s.Yinit = s.Y1*s.mixf + s.Y0*(1.0-s.mixf)
        s.hinit = s.h1*s.mixf + s.h0*(1.0-s.mixf)
        s.gas.HPY = s.hinit,s.P,s.Yinit
        s.gas.equilibrate("HP")
        s.Yinit = s.gas.Y
        s.Tinit = s.gas.T

        #----------- Initialize the particle arrays

        s.Y   = np.tile(s.Yinit,(s.np,1))
        s.h   = np.ones(s.np) * s.hinit
        s.T   = np.ones(s.np) * s.Tinit

        #----------- set pair indexing arrays

        s.npp = s.np/2                # number of particle pairs

        s.ip1 = np.arange(s.npp)      # indices of particles in pairs (p in (p,q))
        s.ip2 = np.arange(s.npp,s.np) # indices of particles in pairs (q in (p,q))
        s.ipp = np.arange(s.npp)      # all pair indices

        #----------- statistics arrays

        s.Tavg = np.zeros(s.nsteps+1)
        s.Tavg[0] = s.Tinit
        
    #--------------------------------------------------------------------------

    def driver(s) :

        s.ee = np.exp(-2.0*s.dt/s.tau_m)                 # convenience factor used below           

        ode = integrate.ode(s.rhsf)
        #ode.set_integrator('vode', method='bdf', nsteps=10000, order=1, min_step=s.dt/2.0)
        ode.set_integrator('lsoda', method='bdf', nsteps=10000, rtol=1.0E-3, atol=1.0E-6, max_order_s=2, max_order_ns=2)
    
        for istep in range(s.nsteps) :
            
            print "step ", istep, " of ", s.nsteps, " Tavg = ", s.Tavg[istep], s.T[s.ip1[0]], s.T[s.ip2[0]]

            s.io_and_generate_pairs()

            #---------- mix the particles (analytically)

            h = s.h.copy()
            Y = s.Y.copy()

            s.h[s.ip1] = 0.5*s.ee*( h[s.ip1] - h[s.ip2] ) + 0.5*( h[s.ip1] + h[s.ip2] ) 
            s.h[s.ip2] = 0.5*s.ee*( h[s.ip2] - h[s.ip1] ) + 0.5*( h[s.ip2] + h[s.ip1] )

            s.Y[s.ip1,:] = 0.5*s.ee*( Y[s.ip1,:] - Y[s.ip2,:] ) + 0.5*( Y[s.ip1,:] + Y[s.ip2,:] )
            s.Y[s.ip2,:] = 0.5*s.ee*( Y[s.ip2,:] - Y[s.ip1,:] ) + 0.5*( Y[s.ip2,:] + Y[s.ip1,:] )
                
            #---------- react the particles

            times = np.array([0.0,s.dt])                 # start and end time (local)

            for s.ip in range(s.np) :

                ode.set_initial_value(s.Y[s.ip,:], 0.0)
                ode.integrate(s.dt)
                s.Y[s.ip,:] = ode.y
                
                s.gas.HPY = s.h[s.ip],s.P,s.Y[s.ip,:]
                s.T[s.ip] = s.gas.T

            #----------

            s.Tavg[istep+1] = np.mean(s.T)

        print s.Tavg



    #--------------------------------------------------------------------------

    def rhsf(s,t,y_local) :

        s.gas.HPY = s.h[s.ip], s.P, y_local
        return s.gas.net_production_rates * s.gas.molecular_weights / s.gas.density


    #--------------------------------------------------------------------------

    def io_and_generate_pairs(s) :

        #---------- To avoid integer rounding of # particles for io and pairing and F/A split:
        #---------- Sample the number of particles stochastically for a fractional:
        #---------- e.g. if we want 3.5 particles per step, then 50% of the time we'll get 3
        #---------- and 50% of the time we'll get 4

        fractional_part, npp_io = np.modf(0.5*s.np*s.dt/s.tau_r)
        if np.random.rand() <= fractional_part  : 
            npp_io = npp_io + 1                                      # number of pairs that enter/leave per step

        fractional_part,npp_pair = np.modf(0.5*s.np*s.dt/s.tau_p)
        if np.random.rand() <= fractional_part  : 
            npp_pair = npp_pair + 1                                  # number of pairs to rearrange per step

        np_io    = 2*npp_io                                          # number of particles that enter/leave per step
        np_pair  = 2*npp_pair                                        # number of particles to rearrange per step

        fractional_part, npf = np.modf(np_io*s.mixf)
        if np.random.rand() <= fractional_part : 
            npf = npf+1                                              # number of fuel stream particles on io
        npa = np_io - npf                                            # number of air stream particles on io

        #---------- 
        
        rnd.shuffle(s.ipp)                                           # shuffle the pair list to grab random pairs

        #---------- 

        ipp_io     = s.ipp[:npp_io]                                  # grab io pair indices
        ipp_pair   = s.ipp[npp_io:npp_io+npp_pair]                   # grab pair indices to re-pair
        ipp_io_pair = np.hstack(( ipp_io,ipp_pair ))                 # list of io and re-pair pairs

        ip_io      = np.hstack(( s.ip1[ipp_io], s.ip2[ipp_io] ))     # list of particles subject to io
        ip_pair    = np.hstack(( s.ip1[ipp_pair], s.ip2[ipp_pair] )) # list of particles to re-pair
        ip_io_pair = np.hstack(( ip_io,ip_pair ))                    # list of io and re-pair particles

        #---------- reset the io particles:

        s.Y[ip_io[:npa],:]        = s.Y0
        s.Y[ip_io[npa:npa+npf],:] = s.Y1

        s.h[ip_io[:npa]]          = s.h0
        s.h[ip_io[npa:npa+npf]]   = s.h1

        #---------- shuffle the io and pairing particles

        rnd.shuffle(ip_io_pair)

        #---------- reset the pairs
       
        s.ip1[ipp_io_pair] = ip_io_pair[: npp_io+npp_pair]
        s.ip2[ipp_io_pair] = ip_io_pair[npp_io+npp_pair :]



#==========================================================================


pm = pmsr()
pm.driver()









