from rsf.proj import *
from Tflow import Tflow
import os

# Size of the model
n1=200
n2=256

# Generate and display the model
Flow('sigmoid',None,
        '''
        sigmoid n1=%d n2=%d d2=.008 |
        smooth rect1=3 diff1=1 adj=1| smooth rect1=3 |
        put label2=Distance
        ''' % (n1,n2) )
Result('sigmoid', 'grey title=Input')


# Compare computation cost
tim = [[],[]]
samples = range(4,22)
for sample in samples:
    nn2         = int(2**sample)
    sigmoid     = 'sigmoid%d' % sample
    pad         = 'pad%d' % sample
    sigmoidpad  = 'sigmoidpad%d' % sample
    dippad      = 'dippad%d' % sample
    seed        = 'seed%d' % sample
    rt          = 'rt%d' % sample
    rtseis      = 'rtseis%d' % sample
    dip         = 'dip%d' % sample
    pwdseis     = 'pwdseis%d' % sample
    siz         = 'siz%d' %sample

    # Generate models with different trace numbers
    Flow(sigmoid,None,
         '''
         sigmoid n1=%d n2=%d d2=.008 |
         smooth rect1=3 diff1=1 adj=1| smooth rect1=3 |
         put label2=Distance
         ''' % (n1,nn2) )

    # Prepare for RT-seislet transform
    Flow(pad,sigmoid,'math output=1 | pad beg1=50 end1=50')
    Flow(sigmoidpad,sigmoid,'pad beg1=50 end1=50 | bandpass fhi=60')
    Flow(dippad,[sigmoidpad,pad],
         '''
         dip order=2 p0=0 verb=y niter=10 rect1=3 rect2=3
         mask=${SOURCES[1]} 
         ''')
    Flow(seed,dippad,'window n2=1 | math output=x1')
    Flow(dippad,[sigmoidpad,pad],
         '''
         dip order=2 p0=0 verb=y niter=10 rect1=3 rect2=3
         mask=${SOURCES[1]}
         ''')
    Flow(rt,[dippad,seed],
         '''
         pwpaint order=2 seed=${SOURCES[1]} i0=%d eps=0.1 
         ''' % (nn2/2.0))

    # Number of traces
    Flow(siz,sigmoid,
         'math output=1 | stack axis=2 norm=n | window n1=1 f1=1')

    # Time of RT-seislet transform
    time = 'tim%d-rt' % sample
    Tflow(time,[sigmoidpad, rt, siz],
          'rtseislet rt=${SOURCES[1]} eps=0.1 adj=y inv=y unit=y type=b')
    tim[0].append(time)

    # Time of PWD-seislet transform
    time = 'tim%d-dip' % sample
    Flow(dip,sigmoidpad,'dip rect1=10 rect2=10 p0=0 pmin=-100')
    Tflow(time,[sigmoidpad, dip, siz],
          'seislet dip=${SOURCES[1]} eps=0.1 adj=y inv=y unit=y type=b')
    tim[1].append(time)

# Combine time files
Flow('tim1', tim[0], 'cat axis=1 ${SOURCES[1:%d]}' % len(samples))
Flow('tim2', tim[1], 'cat axis=1 ${SOURCES[1:%d]}' % len(samples))
Flow('tim','tim1 tim2','cat axis=2 ${SOURCES[1]}')

Result('tim',
       '''
       graph wanttitle=n dash=0,1 screenratio=1. labelsz=9 plotfat=9
       label2="CPU time" unit2=s label1="number of traces" unit1=
       ''')

End()