from rsf.proj import *

def cubeplot(title, clip="", extra=""):
    return """
    byte gainpanel=all %s |
    grey3 frame1=82 frame2=75 frame3=50 flat=n point1=0.7 point2=0.7
    title="%s" %s labelfat=2 wanttitle=n
    """ % (clip, title, extra)

ns2 = 2
ns3 = 2
"""3D model test"""
# Qdome model
Flow("qdome", None,
     """
     qdome
     n1=200 n2=151 n3=100
     d1=0.01 d2=0.01 d3=0.01
     o1=0 o2=0.5 o3=-0.05 |
     smooth rect1=3 diff1=1 | smooth rect1=4 |
     reverse which=2 |
     noise seed=202107 var=0.00000003 |
     put o2=0 o3=0 d2=1 d3=1
     label1=Time label2="TraceX" label3="TraceY"
     unit1=s unit2= unit3=
     """)
Result("noiseqdome", "qdome", cubeplot("", "clip=0.000760897"))

# random mask
Flow("mask", "qdome",
     """
     window n1=1 | noise type=n rep=y seed=150 |
     mask min=0.2 | dd type=float |
     cut n1=2 f1=74 n2=4 f2=47 | dd type=int
     """)
Flow("masky", "mask", "transp plane=12")
# irregular missing data
Flow("gapqd", "qdome mask", "headercut mask=${SOURCES[1]}")
Result("noisegapqd", "gapqd", cubeplot("", "clip=0.000760897"))
# SNR
Flow("diffgapqd", "qdome gapqd", "add scale=1,-1 ${SOURCES[1]}")
Flow("snrgapqd", "qdome diffgapqd", "snr3 noise=${SOURCES[1]}")

### POCS
Flow("pmask", "mask", "spray axis=1 n=200 d=0.004 | dd type=float")
Flow("pmask3", "pmask", "math output=1-input")

pniter = 50
fniter = 50
fforward = """fft1 | fft3 | fft3 axis=3"""
fbackward = """fft3 axis=3 inv=y | fft3 inv=y | fft1 inv=y"""

fdata = "gapqd"
for iter in range(fniter):
    fold = fdata
    fdata = "fdata%d" % iter
    # 1. Forward 3D Fourier
    # 2. Thresholding
    # 3. Inverse 3D Fourier
    # 4. Multiply by space mask
    # 5. Add data outside of hole
    Flow(fdata, [fold, "pmask3", "gapqd"], fforward + """
         | threshold pclip=%g |
         """ % (5.0 + ((99.0 - 5.0) * (iter**2) / ((pniter - 1)**2))) +
        fbackward + """
         | add mode=p ${SOURCES[1]} | 
         add ${SOURCES[2]}
         """)

Flow("pocsqd", "fdata49", "cp")
Result("noisepocsqd", "pocsqd", cubeplot("", "clip=0.000760897"))
# error
Flow("errpocsqd", "qdome pocsqd", "add ${SOURCES[1]} scale=1,-1 ")
Result("noiseerrpocsqd", "errpocsqd", cubeplot("", "clip=0.000760897"))
# SNR
Flow("snrpocsqd", "qdome errpocsqd", "snr3 noise=${SOURCES[1]}")

# SPF3
# padding x and y axis
Flow("maskpad1", "mask", "window n2=20 | reverse which=2 ")
Flow("maskspf1", "maskpad1 mask", "cat ${SOURCES[1]} axis=2")
Flow("maskpad2", "maskspf1", "window n1=20 | reverse which=1")
Flow("maskspf", "maskpad2 maskspf1", "cat ${SOURCES[1]} axis=1")
Flow("datapad1", "gapqd", "window n3=20 | reverse which=4")
Flow("gapdata1", "datapad1 gapqd", "cat ${SOURCES[1]} axis=3")
Flow("datapad2", "gapdata1", "window n2=20 | reverse which=2")
Flow("gapdata", "datapad2 gapdata1", "cat ${SOURCES[1]} axis=2")

Flow("spfx", "gapdata maskspf",
     """
     fxyspfint3 mask=${SOURCES[1]}
     lambdax=0.001 lambday=0.0008 lambdaf=0.0005
     na1=5 na2=6 ftype=1 verb=y
     """)
Flow("maskspft", "maskspf", "transp")
Flow("spfy", "gapdata maskspft",
     """
     transp plane=23|
     fxyspfint3 mask=${SOURCES[1]}
     lambdax=0.001 lambday=0.0008 lambdaf=0.0005
     na1=5 na2=6 ftype=1 verb=y |
     transp plane=23
     """)
Flow("spfqd", "spfx spfy",
     """
     add ${SOURCES[1]} scale=1,1 | math output="input/2" |
     window f2=20 f3=20 | put d2=1 d3=1 o2=0 o3=0
     """)
Result("noisespfqd", "spfqd", cubeplot("", "clip=0.000760897"))
# error
Flow("errspfqd", "qdome spfqd", "add ${SOURCES[1]} scale=1,-1 ")
Result("noiseerrspfqd", "errspfqd", cubeplot("", "clip=0.000760897"))
# SNR
Flow("snrspfqd", "qdome errspfqd", "snr3 noise=${SOURCES[1]}")

# seislet pocs

# forward and backward
sforward = """
pad n2=256 | seislet adj=y inv=y dip=${SOURCES[3]} eps=0.1 type=b 
"""
sbackward = """
seislet adj=n inv=n dip=${SOURCES[3]} eps=0.1 type=b | window n2=151 
"""
pniter = 50
sniter = 50
datas = []
for islice in range(0, 100):
    dmask = "dipmask%d" % islice
    smask = "smask%d" % islice
    wgap = "wgap%d" % islice
    wdip = "wdip%d" % islice
    Flow(dmask, "mask",
         "window n2=1 f2=%g | spray n=200 axis=1| pad n2=256" % islice)
    Flow(smask, "pmask3", "window n3=1 f3=%g" % islice)
    Flow(wgap, "gapqd", "window n3=1 f3=%g" % islice)
    Flow(wdip, [wgap, dmask],
         "pad n2=256 | dip rect1=20 rect2=20 mask=${SOURCES[1]}")
    sdata = wgap
    for itera in range(sniter):
        sfold = sdata
        sdata = "wgap%s-%s" % (islice, itera)
        # 1. Forward seislet transform
        # 2. Thresholding
        # 3. Inverse seislet transform
        # 4. Multiply by space mask
        # 5. Add data outside of hole
        Flow(sdata, [sfold, smask, wgap, wdip], sforward + """
             | threshold pclip=%g |
            """ % (5.0 + ((99.0 - 5.0) * (itera**2) / ((pniter - 1)**2))) +
            sbackward + """
            | add mode=p ${SOURCES[1]} |
            add ${SOURCES[2]}
            """)
        if itera == (sniter - 1):
            datas.append(sdata)

Flow("stpocsqd", datas, "cat ${SOURCES[1:%d]} axis=3" % len(datas))
Result("noisestpocsqd", "stpocsqd", cubeplot("", "clip=0.000760897"))
Flow("errstpocsqd", "qdome stpocsqd", "add ${SOURCES[1]} scale=1,-1 ")
Result("noiseerrstpocsqd", "errstpocsqd", cubeplot("", "clip=0.000760897"))
# SNR
Flow("snrstpocsqd", "qdome errstpocsqd", "snr3 noise=${SOURCES[1]}")

End()