# . . Set up generic project files
from rsf.proj import *
from rsf.prog import RSFROOT

import math,os,fdmod,encode,wemig,stiffness

# . . Set parameter for the modelling. 
par = {
    # Model Dimensions
    'nx':288, 'ox':0, 'dx':0.005,  'lx':'x', 'ux':'km',
    'ny':288, 'oy':0, 'dy':0.005,  'ly':'y', 'uy':'km',
    'nz':288, 'oz':0, 'dz':0.005,  'lz':'z', 'uz':'km',

    # Time parameters
    'nt':1000,'ot':0, 'dt':0.0005,  'lt':'t', 'ut':'s',
    'kt':100,'frq':40,

    # Shot locations
    'nsx':1, 'dsx':0.005, 'osx':0.120,
    'nsy':1, 'dsy':0.005, 'osy':0.120,

    # EWEFD2D parameters
    'snap':'n','jsnap':1000,'height':6,
    'nb':48,'nbell':5,'jdata':10,'dabc':'y',
    'free':'n','ssou':'y',
}

# . . Initialize parameters in fdm module   
fdmod.param(par)
# . . Thomsen parameters
par['vp']=2.0
par['vs']=1.1547
par['ro']=2000

# --------------------------------------------------------------------
# . . 3D Elastic source
fdmod.wavelet('wav_',par['frq'],par)
Flow('souz','wav_','math output=input*1')
Flow('soux','wav_','math output=input*1')
Flow('souy','wav_','math output=input*1')
Flow('wave',['souz','soux','souy'],
    '''
    cat axis=2 space=n ${SOURCES[1:3]} |
    transp plane=12 |
    transp plane=23 |
    transp plane=12
    ''')

# --------------------------------------------------------------------
# . . Coordinates Section
# . . Locate source position
xsou=par['osx'];
ysou=par['osy']
zsou=par['oz'];

# . . 3D Sources
fdmod.point3d('ss',xsou,ysou,zsou,par) # . . 2D Point Sources
fdmod.horizontal3d('rr',0,par)    # . . 2D receivers

# --------------------------------------------------------------------
# . .  3D model section
# . . Create zero array size of 2D model
Flow('zero',None,
     '''
     spike nsp=1 mag=0.0
     n1=%(nz)d o1=%(oz)g d1=%(dz)g
     n2=%(nx)d o2=%(ox)g d2=%(dx)g 
     n3=%(nx)d o3=%(ox)g d3=%(dx)g |
     put label1=%(lz)s label2=%(lx)s label3=%(ly)s 
     unit1=%(uz)s unit2=%(ux)s unit3=%(uy)s
     ''' % par)

# . . Create 3D Vp, Vs and rho models
Flow('vp','zero','math output="%(vp)g"' %par)
Flow('vs','zero','math output="%(vs)g"' %par)
Flow('ro','zero','math output="%(ro)g"' %par)
stiffness.iso3d('cc','vp','vs','ro',par)

# ------------------------------------------------------------
# . . Elastic Modeling Section
par['begmod']=152 #32
par['endmod']=par['nx']
par['dmod']=24

mindex=range(par['begmod'],par['endmod'],par['dmod'])

mpirun = WhereIs('mpirun')
par['ntest']=10
rindex=range(0,par['ntest'],1)

#for msou in mindex:
#    mtag="-%03d" % msou
#    par['mm']=msou
#    par['mtag']=mtag
#    Flow('rr'+mtag,'rr','window n2=%(mm)d' % par)
#    Flow('cc'+mtag,'cc','window n1=%(mm)d n2=%(mm)d n3=%(mm)d ' %par )
#    Flow('ro'+mtag,'ro','window n1=%(mm)d n2=%(mm)d n3=%(mm)d ' %par )
#    # . . GPU Timing
#    # . . Set up timing file
#    for rsou in rindex:
#        rtag="-%02d" % rsou
#        alltag=mtag+rtag
#
#       Flow('timegpu3d-SIN'+alltag,None,
#           '''
#            spike n1=3 k1=1 mag=0 d1=%(dmod)g o1=%(begmod)g | 
#            dd form=ascii out=${TARGETS[0]}@
#            ''' %par)
#
#       # . . Run GPU modelling
#       Flow(['Dgpu3d-SIN'+alltag,'Wgpu-SIN'+alltag],
#           ['wave','cc'+mtag,'ro'+mtag,'ss','rr'+mtag,'timegpu3d-SIN'+alltag],
#            '''
#            /usr/bin/time -o timegpu3d-SIN%s.rsf@ ''' %alltag
#           +'--format="\t%e \t%U \t%S" ' +
#            '''
#            %(SDIR)s/sfewefd3d_gpu_p2p
#            jdata=%(jdata)d dabc=%(dabc)s free=%(free)s opot=n
#            verb=y snap=%(snap)s jsnap=%(jsnap)d ssou=%(ssou)s
#            nb=%(nb)d nbell=%(nbell)d ngpu=1
#            in1=${SOURCES[0]} ccc=${SOURCES[1]}
#            den=${SOURCES[2]} sou=${SOURCES[3]}
#            rec=${SOURCES[4]} wfl=${TARGETS[1]}
#            ''' % par)
#        # . . Back to native
#       Flow('gputiming3d-SIN'+alltag,['timegpu3d-SIN'+alltag,'Dgpu3d-SIN'+alltag],
#           'dd form=native')
#    Flow('gputiming3d-SIN'+mtag,map(lambda x: 'gputiming3d-SIN%s' %mtag+'-%02d.rsf' % x,rindex),
#            'cat ${SOURCES[1:%d]} axis=2 norm=y | stack axis=2 norm=y' %par['ntest'])
#    ####################################################
#    # . . GPU Timing
#    # . . Set up timing file
#    for rsou in rindex:
#        rtag="-%02d" % rsou
#        alltag=mtag+rtag
#       Flow('timegpu3d-P2P'+alltag,None,
#             '''
#             spike n1=3 k1=1 mag=0 d1=%(dmod)g o1=%(begmod)g | 
#             dd form=ascii out=${TARGETS[0]}@
#             ''' %par)
#        # . . Run GPU modelling
#       Flow(['Dgpu3d-P2P'+alltag,'Wgpu3d-P2P'+alltag],
#            ['wave','cc'+mtag,'ro'+mtag,'ss','rr'+mtag,'timegpu3d-P2P'+alltag],
#            '''
#            /usr/bin/time -o timegpu3d-P2P%s.rsf@ ''' %alltag
#           +'--format="\t%e \t%U \t%S" ' +
#            '''
#            %(SDIR)s/sfewefd3d_gpu_p2p
#            jdata=%(jdata)d dabc=%(dabc)s free=%(free)s opot=n
#            verb=y snap=%(snap)s jsnap=%(jsnap)d ssou=%(ssou)s
#            nb=%(nb)d nbell=%(nbell)d ngpu=2
#            in1=${SOURCES[0]} ccc=${SOURCES[1]}
#            den=${SOURCES[2]} sou=${SOURCES[3]}
#            rec=${SOURCES[4]} wfl=${TARGETS[1]}
#            ''' % par)
#       Flow('gputiming3d-P2P'+alltag,['timegpu3d-P2P'+alltag,'Dgpu3d-P2P'+alltag],'dd form=native')
#
#    # . . Back to native
#    Flow('gputiming3d-P2P'+mtag,map(lambda x:'gputiming3d-P2P%s' %mtag +'-%02d.rsf'% x,rindex),
#            'cat ${SOURCES[1:%d]} axis=2 norm=y | stack axis=2 norm=y' %par['ntest'])
#
#####################################################
#    # . . GPU Timing
#    # . . Set up timing file
#    for rsou in rindex:
#        rtag="-%02d" % rsou
#        alltag=mtag+rtag
#       Flow('timegpu3d-MPI'+alltag,None,
#             '''
#             spike n1=3 k1=1 mag=0 d1=%(dmod)g o1=%(begmod)g | 
#             dd form=ascii out=${TARGETS[0]}@
#             ''' %par)
#        # . . Run GPU modelling
#       Flow(['Dgpu3d-MPI'+alltag,'Wgpu3d-MPI'+alltag],
#            ['wave','cc'+mtag,'ro'+mtag,'ss','rr'+mtag,'timegpu3d-MPI'+alltag],
#            '''
#            /usr/bin/time -o timegpu3d-MPI%s.rsf@ ''' %alltag
#                        +'--format="\t%e \t%U \t%S" /usr/bin/mpirun -np 2' +
#            '''
#            %(SDIR)s/sfewefd3d_gpu_mpi
#               jdata=%(jdata)d dabc=%(dabc)s free=%(free)s opot=n
#            verb=y snap=%(snap)s jsnap=%(jsnap)d ssou=%(ssou)s
#            nb=%(nb)d nbell=%(nbell)d ngpu=2
#            in1=${SOURCES[0]} ccc=${SOURCES[1]}
#            den=${SOURCES[2]} sou=${SOURCES[3]}
#            rec=${SOURCES[4]} wfl=${TARGETS[1]}
#            ''' % par)
#       Flow('gputiming3d-MPI'+alltag,['timegpu3d-MPI'+alltag,'Dgpu3d-MPI'+alltag],'dd form=native')
#    # . . Back to native
#    Flow('gputiming3d-MPI'+mtag,map(lambda x: 'gputiming3d-MPI%s'%mtag +'-%02d.rsf'% x,rindex),
#            'cat ${SOURCES[1:%d]} axis=2 norm=y | stack axis=2 norm=y' %par['ntest'])
#
#####################################################
#    # . . CPU Timing
#    # . . Set up timing file
#
#    for rsou in rindex:
#       rtag="-%02d" % rsou
#       alltag=mtag+rtag
#               Flow('timecpu3d'+alltag,None,
#             '''
#             spike n1=3 k1=1 mag=0 d1=%(dmod)g o1=%(begmod)g | 
#             dd form=ascii out=${TARGETS[0]}@ 
#             ''' %par)
#
#        # . . Perform CPU MODELLING
#       Flow(['Dcpu3d'+alltag,'Wcpu'+alltag],
#            ['wave','cc'+mtag,'ro'+mtag,'ss','rr'+mtag,'timecpu3d'+alltag],
#            '''
#            /usr/bin/time -o timecpu3d%s.rsf@ ''' %alltag
#                +'--format="\t%e \t%U \t%S" ' +
#            '''
#            %(SDIR)s/sfewefd3d_omp
#            jdata=%(jdata)d dabc=%(dabc)s free=%(free)s opot=n
#            verb=y snap=y jsnap=%(jsnap)d ssou=%(ssou)s
#            nb=%(nb)d nbell=%(nbell)d
#            in1=${SOURCES[0]} ccc=${SOURCES[1]}
#            den=${SOURCES[2]} sou=${SOURCES[3]}
#            rec=${SOURCES[4]} wfl=${TARGETS[1]}
#            ''' % par)
#        # . . Back to native
#       Flow('cputiming3d'+alltag,['timecpu3d'+alltag,'Dcpu3d'+alltag],'dd form=native')
#    Flow('cputiming3d'+mtag,map(lambda x: 'cputiming3d%s' %mtag +'-%02d.rsf' % x,rindex),
#            'cat ${SOURCES[1:%d]} axis=2 norm=y | stack axis=2 norm=y' %par['ntest'])
#
## . . Concatenate all together
#par['nmod']=len(mindex)
#
#Flow('gputiming3d-SIN', map(lambda x: 'gputiming3d-SIN-'+'%03d.rsf' % x,mindex), 'cat #${SOURCES[1:%d]} axis=2 | window n1=1 f1=0 | put d1=24 o1=90' % par['nmod'])
#Flow('gputiming3d-P2P', map(lambda x: 'gputiming3d-P2P-'+'%03d.rsf' % x,mindex), 'cat #${SOURCES[1:%d]} axis=2 | window n1=1 f1=0 | put d1=24 o1=90' % par['nmod'])
#Flow('gputiming3d-MPI', map(lambda x: 'gputiming3d-MPI-'+'%03d.rsf' % x,mindex), 'cat #${SOURCES[1:%d]} axis=2 | window n1=1 f1=0 | put d1=24 o1=90' % par['nmod'])
#Flow('cputiming3d', map(lambda x: 'cputiming3d-'+'%03d.rsf' % x,mindex), 'cat ${SOURCES[1:%d]} #axis=2 | window n1=1 f1=0 | put d1=24 o1=90' % par['nmod'])
#
## . . Compute graphs for runtime and speedup metrics
#Flow('Runtime3d','gputiming3d-SIN gputiming3d-P2P gputiming3d-MPI cputiming3d','cat ${SOURCES[1:4]} #axis=2')
#Flow('Speedup3d-1','gputiming3d-SIN cputiming3d','math other=${SOURCES[1]} output="other/input"')
#Flow('Speedup3d-2','gputiming3d-P2P cputiming3d','math other=${SOURCES[1]} output="other/input"')
#Flow('Speedup3d-3','gputiming3d-MPI cputiming3d','math other=${SOURCES[1]} output="other/input"')
#Flow('Speedup3d','Speedup3d-1 Speedup3d-2 Speedup3d-3','cat ${SOURCES[1:3]} axis=2 space=n')
#
#Plot('Runtimes3d1','Runtime3d','graph wantaxis=n title=" " plotfat=5  ')
#Plot('Runtimes3d2','Runtime3d','graph label2=Runtimes unit2=s label1="Model Dimension" title=" " #unit1="N\^3\_" plotfat=5 symbolsz=10 symbol="o"')
#Result('Runtimes3d','Runtimes3d1 Runtimes3d2','Overlay')
#
#Plot('Speedup3d1','Speedup3d ','graph wantaxis=n title="" plotfat=5')
#Plot('Speedup3d2','Speedup3d','graph label2=Speedup unit2="" label1="Model Dimension" #unit1="N\^3\_" title="" plotfat=10 symbol="o" ')
##Result('Speedup3d','Speedup3d1 Speedup3d2','Overlay')
#

###########
# . . OUR TEST CPU results
Flow('CPUx cpux.rsf@  cpux.rsf',None,
    '''
    echo ; echo n1=8 o1=0 d1=1 esize=0 in=${TARGETS[1]} > ${TARGETS[2]} &&
    echo 128 152 176 200 224 248 272 296 > ${TARGETS[1]} &&
    dd < ${TARGETS[2]} form=native 
    ''' %par)
Flow('CPUy cpuy.rsf@ cpuy.rsf',None,
    '''
    echo ; echo n1=8 o1=0 d1=1 esize=0 in=${TARGETS[1]} > ${TARGETS[2]} &&
    echo 159.6 236 360.2 521.6 787.5 1098 1490 1926 > ${TARGETS[1]} &&
    dd < ${TARGETS[2]} form=native 
    ''' %par)
Flow('CPU','CPUx CPUy','cmplx ${SOURCES[1]}')

###########
# . . OUR TEST SINGLE results
Flow('SINx sinx.rsf@ sinx.rsf',None,
    '''
    echo ; echo n1=8 o1=0 d1=1 esize=0 in=${TARGETS[1]} > ${TARGETS[2]} && 
    echo 128 152 176 200 224 224 224 224 > ${TARGETS[1]} && 
    dd < ${TARGETS[2]} form=native 
    ''' %par)
Flow('SINy siny.rsf@ siny.rsf',None,
    '''
    echo ; echo n1=8 o1=0 d1=1 esize=0 in=${TARGETS[1]} > ${TARGETS[2]} && 
    echo 14.22 21.18 29.29 41.1 53.01 53.01 53.01 53.01 > ${TARGETS[1]} && 
    dd < ${TARGETS[2]} form=native 
    ''' %par)
Flow('SIN','SINx SINy','cmplx ${SOURCES[1]}')

###########
# . . OUR TEST MPI results
Flow('MPIx mpix.rsf@ mpix.rsf',None,
    '''
    echo ; echo n1=8 o1=0 d1=1 esize=0 in=${TARGETS[1]} > ${TARGETS[2]} && 
    echo 128 152 176 200 224 248 272 296 > ${TARGETS[1]} && 
    dd < ${TARGETS[2]} form=native 
    ''' %par)
Flow('MPIy mpiy.rsf@ mpiy.rsf',None,
    '''
    echo ; echo n1=8 o1=0 d1=1 esize=0 in=${TARGETS[1]} > ${TARGETS[2]} && 
    echo 12.79  17.93   24.21   33.06   42.28   49.72   63.26   79.16 > ${TARGETS[1]} && 
    dd < ${TARGETS[2]} form=native 
    ''' %par)
Flow('MPI','MPIx MPIy','cmplx ${SOURCES[1]}')

############
# . . OUR TEST CPU results
Flow('P2Px p2px.rsf@ p2px.rsf',None,
    '''
    echo ; echo n1=8 o1=0 d1=1 esize=0 in=${TARGETS[1]} > ${TARGETS[2]} && 
    echo 128 152 176 200 224 248 272 296 > ${TARGETS[1]} && 
    dd < ${TARGETS[2]} form=native 
    ''' %par)
Flow('P2Py p2py.rsf@  p2py.rsf',None,
    '''
    echo ; echo n1=8 o1=0 d1=1 esize=0 in=${TARGETS[1]} > ${TARGETS[2]} && 
    echo  11.26 15.92   21.28   28.9    36.65   45.27   58.26   68.95 > ${TARGETS[1]} && 
    dd < ${TARGETS[2]} form=native 
    ''' %par)
Flow('P2P','P2Px P2Py','cmplx ${SOURCES[1]}')

###############
# . . RUNTIME
Flow('ALLRUNTIME','SIN MPI P2P CPU','cat ${SOURCES[1:4]} axis=2')

Plot('ALLRUNTIME1','ALLRUNTIME','graph wantaxis=n title="" plotfat=5')
Plot('ALLRUNTIME2','ALLRUNTIME',
     '''
     graph label2=Runtime unit2=s 
     label1="Model Dimension" unit1="N\^3\_" 
     title="" plotfat=12 symbol="o" 
     ''')
Result('ALLRUNTIME','ALLRUNTIME1 ALLRUNTIME2','Overlay')

###############
# . . RUNTIME
Flow('SPEEDUP1y','CPUy SINy','math GPU=${SOURCES[1]} output="input/GPU-14.86" | window n1=5 | pad n1=8 | math output="input+14.86" ')
Flow('SPEEDUP1','SINx SPEEDUP1y','cmplx ${SOURCES[1]} ')
Flow('SPEEDUP2y','CPUy MPIy','math GPU=${SOURCES[1]} output="input/GPU" ')
Flow('SPEEDUP2','CPUx SPEEDUP2y','cmplx ${SOURCES[1]} ')
Flow('SPEEDUP3y','CPUy P2Py','math GPU=${SOURCES[1]} output="input/GPU" ')
Flow('SPEEDUP3','CPUx SPEEDUP3y','cmplx ${SOURCES[1]} ')

Flow('SPEEDUP','SPEEDUP1 SPEEDUP2 SPEEDUP3','cat axis=2 ${SOURCES[1:3]} ')

Plot('SPEED1','SPEEDUP','graph wantaxis=n title="" plotfat=5')
Plot('SPEED2','SPEEDUP',
     '''
     graph label2=Speedup unit2="" 
     label1="Model Dimension" unit1="N\^3\_" 
     title="" plotfat=12 symbol="o" 
     ''')
Result('SPEEDUP','SPEED1 SPEED2','Overlay')

# . . MPI vs P2P speed up metrics
Flow('P2PvMPIy','MPIy P2Py','math P2P=${SOURCES[1]} output="input/P2P" ')
Flow('P2PvMPI','MPIx P2PvMPIy','cmplx ${SOURCES[1]} ')
Plot('P2PvMPI1','P2PvMPI','graph wantaxis=n title="" plotfat=5')
Plot('P2PvMPI2','P2PvMPI',
     '''
     graph label2=Speedup unit2="" 
     label1="Model Dimension" unit1="N\^3\_" 
     title="" plotfat=12 symbol="o" 
     ''')
Result('P2PvMPI','P2PvMPI1 P2PvMPI2','Overlay')

End()