import cupy as cp

from flandre.beamformer.das import gen_pwi
from flandre.beamformer.dist import direct_dist
from flandre.utils.Config import DeviceConfig
from flandre.utils.Msg import ImageArgMsg
from flandre.utils.RfMat import RfMat
from flandre.utils.RfSequence import RfSequence

if __name__ == '__main__':
    arg = ImageArgMsg(
        sender='',
        t_end=2900,
        t_start=0,
        v2=1524,
        dct_center=1086,
        dct_bandwidth=915,
        f_rows=6002,
        beta=40,
        tgc=0,
        g8=80
    )

    dc = DeviceConfig(v2=1540, rows=5999)
    pwi, _, la = gen_pwi(direct_dist(dc, p=cp), dc)
    seq = RfSequence('/run/media/lambda/b86dccdc-f134-464b-a310-6575ee9ae85c/us/baby789,S=(256 6002),M=PWI,U=120/')
    for i, frame in enumerate(seq.frames):
        data = RfMat.from_rf_frame(frame, device='gpu')
        data = data.dct_center(arg.dct_center, arg.dct_bandwidth)
        data = data.call(lambda m: m.astype(cp.int16))
        data = data.call(pwi)
        data = data.call(cp.asarray, order='C')
        data = data.argrelextrema()
        data = data.conv_guass(b=arg.beta * 0.01)
        data = data.crop(arg.t_start, arg.t_end)
        data = data.time_gain_compensation_global((1 - arg.g8 * (1.0 / 128)) ** 2)
        data = data.rotate90()
        data.png(f'/home/lambda/source/scarlet/flandre/@DS/test/{i}.png', pre=300)