from rsf.proj import *
import fdmod,stiffness

# ------------------------------------------------------------
#
# . . PARAMETERS 
# 
# ------------------------------------------------------------
par = {

# . . Time parameters
'nt':14000, 'dt':0.0005, 'ot':0, 'lt':'Time', 'ut':'s',
'kt':130,'frq':15,

# . . Space parameters
'nx':3792, 'dx':6.25,  'ox':4000., 'lx':'Distance', 'ux':'m',
'ny':1,    'dy':6.25,  'oy':0.,   'ly':'y', 'uy':'m',
'nz':1896, 'dz':6.25,  'oz':0.,   'lz':'Depth', 'uz':'m',

# . . Source Parameters
'nsx':1,'dsx':0.004,'osx':0.512,
'nsy':1,'dsy':0.004,'osy':0.512,

# . . Modelling parameters
'nb':40,'verb':'y','jdata':1,'free':'n','ssou':'y',
'opot':'n','snap':'y','jsnap':100,'dabc':'y','nbell':11,

# . . Plotting parameters
'scalebar':1,'ratio':1,'height':0.8,
}

par['ncut'] = par['nz']/2;
par['begx'] = par['nx']/2;
par['begy'] = par['ny']/2;
par['xmin'] = par['ox']
par['zmin'] = par['oz']
par['tmin']=0
par['tmax']=par['tmin']+(par['nt']-1)*par['dt']
par['zmax']=par['oz']+(par['nz']-1)*par['dz']
par['xmax']=par['ox']+(par['nx']-1)*par['dx']
par['endx']=par['nx'] - par['begx']
par['labelattr']=' labelsz=8 '

##########################

# . . make source wavelet
fdmod.wavelet('wavtmp',par['frq'],par)

Flow('souz','wavtmp','math output=input*1')
Flow('soux','wavtmp','math output=input*1')

Flow('wave-2d',['souz','soux'],
     '''
     cat axis=2 space=n ${SOURCES[1:2]} |
     transp plane=12 |
     transp plane=23 |
     transp plane=12
     ''')
fdmod.ewavelet('wave-2d','',par)

# . . set up source and receiver
xsou=par['ox']+(par['nx']/2*par['dx']);
zsou=0.03

fdmod.point('ss-2d',xsou,zsou,par)
fdmod.horizontal('rr-2d',zsou,par)

#########################
# BP VELOCITY MODEL
par['dxbeg']=12.5

# . . Velocity model
velsegy = 'vel_z6.25m_x12.5m_exact.segy'
Fetch(velsegy+'.gz',dir='2004_BP_Vel_Benchmark',
      server='ftp://software.seg.org',top='pub/datasets/2D')
zcat = WhereIs('gzcat') or WhereIs('zcat')
Flow(velsegy,velsegy+'.gz',zcat + ' $SOURCE',stdin=0)

# . . P-wave velocity model
Flow('VP-BP',velsegy,
    '''
    segyread read=data | 
    put d1=%(dz)g d2=%(dxbeg)g | transp plane=12 | 
    spline d1=%(dx)g n1=10790 o1=0 | transp | 
    put label1=Depth label2=Distance unit1=m unit2=m | 
    math output="input" 
    ''' %par)

# . . S-wave velocity model
Flow('VS-BP','VP-BP','math output="input/sqrt(3.0)" ')

# . . Density model
densegy = 'density_z6.25m_x12.5m.segy'
Fetch(densegy+'.gz',dir='2004_BP_Vel_Benchmark',
      server='ftp://software.seg.org',top='pub/datasets/2D')
zcat = WhereIs('gzcat') or WhereIs('zcat')
Flow(densegy,densegy+'.gz',zcat + ' $SOURCE',stdin=0)

Flow('RO-BP',densegy,
    '''
    segyread read=data | 
    put d1=%(dz)g d2=%(dxbeg)g | transp plane=12 | 
    spline d1=%(dx)g n1=10790 o1=0| transp | 
    put label1=Depth label2=Distance unit1=m unit2=m | 
    math output="1000*input" 
    ''' %par)

# . . Cut out small sections
Flow('VPcut','VP-BP',
        'window n1=%(nz)d n2=%(nx)d min1=%(oz)g min2=%(ox)g| smooth rect1=2 rect2=2' %par)
Flow('VScut','VS-BP',
        'window n1=%(nz)d n2=%(nx)d min1=%(oz)g min2=%(ox)g| smooth rect1=2 rect2=2' %par)
Flow('ROcut','RO-BP',
        'window n1=%(nz)d n2=%(nx)d min1=%(oz)g min2=%(ox)g| smooth rect1=5 rect2=5' %par)

# . . Create 2D input stiffness tensor
stiffness.iso2d('BPcut','VPcut','VScut','ROcut',par)

# . . GPU Data Modeling
Flow(['BPdat_gpu','BPwfld_gpu'], ['wave-2d','BPcut','ROcut','ss-2d','rr-2d'],
    '''
    ewefd2d_gpu 
    in1=${SOURCES[0]} ccc=${SOURCES[1]} den=${SOURCES[2]} 
    sou=${SOURCES[3]} rec=${SOURCES[4]} wfl=${TARGETS[1]} 
    nt=%(nt)d jdata=%(jdata)d verb=%(verb)s free=%(free)s 
    ssou=%(ssou)s opot=%(opot)s snap=%(snap)s jsnap=%(jsnap)d 
    dabc=%(dabc)s nb=%(nb)d nbell=%(nbell)d
    ''' %par)

# . . CPU Data Modeling
Flow(['BPdat_cpu','BPwfld_cpu'], ['wave-2d','BPcut','ROcut','ss-2d','rr-2d'],
    '''
    ewefd2d_omp
    ccc=${SOURCES[1]} den=${SOURCES[2]}
    sou=${SOURCES[3]} rec=${SOURCES[4]} wfl=${TARGETS[1]}
    nt=%(nt)d jdata=%(jdata)d verb=%(verb)s free=%(free)s 
    ssou=%(ssou)s opot=%(opot)s snap=%(snap)s jsnap=%(jsnap)d 
    dabc=%(dabc)s nb=%(nb)d nbell=%(nbell)d
    ''' %par)

#####################################
#
# . . Figures section

# . . Wavefield/velocity overlay
Flow('wavevel','BPwfld_gpu VPcut',
    '''
    window n3=1 n4=1 f4=80 min1=%(oz)g min2=%(ox)g n1=%(nz)d n2=%(nx)d|
    scale axis=123 |
    math other=${SOURCES[1]} output="input-(other-3562)/10000"
    ''' %par )
Result('wavevel',fdmod.cgrey('screenht=8 screenratio=0.5 pclip=97 wantscalebar=n',par))

Flow('DataDiff','BPdat_cpu BPdat_gpu','math x=${SOURCES[1]} output="input-x" ')
Flow('Difftest','BPdat_cpu BPdat_gpu DataDiff',
    '''
    cat ${SOURCES[1]} ${SOURCES[2]} axis=4 space=n | 
    window n2=1 j3=8 |
    transp plane=12 |
    scale axis=12 |
    byte gainpanel=a pclip=99.5 | 
    put label1=%(lt)s unit1=%(ut)s label2=%(lx)s unit2=%(ux)s
    ''' %par)

Plot('CPUtest', 'Difftest','window n3=1 f3=0 |'+ fdmod.dgrey('title="(b)" labelsz=12 titlesz=14',par))
Plot('GPUtest', 'Difftest','window n3=1 f3=1 |'+ fdmod.dgrey('title="(a)" labelsz=12 titlesz=14',par))
Plot('Difftest','Difftest','window n3=1 f3=2 |'+ fdmod.dgrey('title="(c)" labelsz=12 titlesz=14',par))

# . . Section comparison
Result('BPdata','GPUtest CPUtest Difftest','SideBySideAniso' )


End()