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

# . . Import modules
import math,os
import fdmod,encode,wemig,stiffness

# . . Set parameter for the modelling. 
par = {
    # Model Dimensions
    'nx':4882, 'ox':0, 'dx':0.005,  'lx':'x', 'ux':'km',
    'ny':1,   'oy':0, 'dy':0.005,  'ly':'y', 'uy':'km',
    'nz':4812, '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.48,
    'nsy':1, 'dsy':0.005, 'osy':0,

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

# . . Initialize parameters in fdm module   
fdmod.param(par)
par['nframe']=20
par['iframe']=8

# . . Thomsen parameters
par['vp']=2.0
par['vs']=1.1547
par['ro']=2000

# . . Set up a wavelet
fdmod.wavelet('wav_',par['frq'],par)

# --------------------------------------------------------------------
# . . Source Section
#
# . . Elastic source
Flow('souz','wav_','math output=input*1')
Flow('soux','wav_','math output=input*1')
Flow('souy','wav_','math output=input*1')

# . . Make 2D elastic source
Flow('wave',['souz','soux'],
    '''
    cat axis=2 space=n ${SOURCES[1:2]} |
    transp plane=12 |
    transp plane=23 |
    transp plane=12
    ''')

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

# . . 2D Sources
fdmod.point('ss',xsou,zsou,par) # . . 2D Point Sources
fdmod.horizontal('rr',0,par)    # . . 2D receivers

# --------------------------------------------------------------------
# . .  2D 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 |
     put label1=%(lz)s label2=%(lx)s unit1=%(uz)s unit2=%(ux)s
     ''' % par)

# . . Create 2D 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.iso2d('cc','vp','vs','ro',par)

# ------------------------------------------------------------
# . . Elastic Modeling Section
#
# . . 2D Modelling cases CPU
par['begmod']=432
par['endmod']=3792
par['dmod']=240

mindex=range(par['begmod'],par['endmod'],par['dmod'])
par['ntest']=1
rindex=range(0,par['ntest'],1)

par['SDIR']=os.path.join(RSFROOT,'bin')

for msou in mindex:
    mtag="-%04d" % 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 ' %par )
    Flow('ro'+mtag,'ro','window n1=%(mm)d n2=%(mm)d ' %par )

    # . . GPU Timing
    # . . Output timing file
    for rsou in rindex:
        rtag="-%02d" % rsou
        alltag=mtag+rtag

        Flow('timegpu'+alltag,None,
                        '''
                        spike n1=3 k1=1 mag=0 d1=%(dmod)g o1=%(begmod)g | 
                        dd form=ascii out=${TARGETS[0]}@ 
                        ''' %par)

        # . . Do GPU modelling
        Flow(['Dgpu'+alltag,'timegpu%s.rsf@' % alltag],
                        ['wave','cc'+mtag,'ro'+mtag,'ss','rr'+mtag,'timegpu'+alltag],
                '''
                /usr/bin/time -o ${TARGETS[1]} --format="\t%%e \t%%U \t%%S" 
                %(SDIR)s/sfewefd2d_gpu
                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
                in1=${SOURCES[0]} ccc=${SOURCES[1]}
                den=${SOURCES[2]} sou=${SOURCES[3]}
                rec=${SOURCES[4]} 
                ''' % par)
        # . . Back to native format
        Flow('gputiming'+alltag,['timegpu'+alltag,'Dgpu'+alltag],'dd form=native')
        Flow('gputiming'+mtag,map(lambda x: 'gputiming%s' %mtag +'-%02d.rsf' % x,rindex),
             'cat ${SOURCES[1:%d]} axis=2 norm=y | window n1=1|put d1=1 o1=1 ' %par['ntest'])

    # . . CPU Timing
    for rsou in rindex:
        rtag="-%02d" % rsou
        alltag=mtag+rtag
        Flow('timecpu'+alltag,None,
                        '''
                        spike n1=3 k1=1 mag=0 d1=%(dmod)g o1=%(begmod)g |
                        dd form=ascii out=${TARGETS[0]}@
                        ''' %par)

        # . . Run modelling
        Flow(['Dcpu'+alltag,'timecpu%s.rsf@' % alltag],
                        ['wave','cc'+mtag,'ro'+mtag,'ss','rr'+mtag,'timecpu'+alltag],
                '''
                /usr/bin/time -o ${TARGETS[1]} --format="\t%%e \t%%U \t%%S" 
                %(SDIR)s/sfewefd2d_omp
                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
                in1=${SOURCES[0]} ccc=${SOURCES[1]}
                den=${SOURCES[2]} sou=${SOURCES[3]}
                rec=${SOURCES[4]} 
                ''' % par)
        # . . Back to native format
        Flow('cputiming'+alltag,['timecpu'+alltag,'Dcpu'+alltag],'dd form=native')
        Flow('cputiming'+mtag,map(lambda x: 'cputiming%s' %mtag +'-%02d.rsf' % x,rindex),
             'cat ${SOURCES[1:%d]} axis=2 norm=y | window n1=1 |put d1=1 o1=1 ' %par['ntest'])

par['nmod']=len(mindex)

Flow('gputiming', map(lambda x: 'gputiming-'+'%04d.rsf' % x,mindex), 'cat ${SOURCES[1:%d]} axis=2 | put d2=240 o2=528 ' % par['nmod'])
Flow('cputiming', map(lambda x: 'cputiming-'+'%04d.rsf' % x,mindex), 'cat ${SOURCES[1:%d]} axis=2 | put d2=240 o2=528 ' % par['nmod'])

# . . Fake the modelling runtime results
Flow('Runtime-ALL','gputiming cputiming','cat ${SOURCES[1]} axis=3')
Flow('gputiming-AVE','gputiming','stack axis=1')
Flow('cputiming-AVE','cputiming','stack axis=1')
Flow('Runtime-AVE','gputiming-AVE cputiming-AVE','cat ${SOURCES[1]} axis=2')

Flow('Speedup','gputiming-AVE cputiming-AVE','math other=${SOURCES[1]} output="other/input"')

Plot('Runtimes1','Runtime-AVE','graph wantaxis=n title=" " plotfat=5  ')
Plot('Runtimes2','Runtime-AVE','graph label2=Runtimes unit2=s label1="Model Dimension" title=" " unit1="N\^2\_" plotfat=5 symbolsz=8 symbol="o"')
Result('Runtimes','Runtimes1 Runtimes2','Overlay')

Plot('Speedup1','Speedup','graph wantaxis=n title="" plotfat=5')
Plot('Speedup2','Speedup','graph label2=Speedup unit2="" label1="Model Dimension" unit1="N\^2\_" title="" plotfat=10 symbol="o" ')
Result('Speedup','Speedup1 Speedup2','Overlay')


End()