import cupy as cp

from beamformer.das import TFM
from utils.Config import ImagingConfig
from utils.ScanData import ScanData

tfm_cache = [None]


def pwi_process(s: ScanData, icfg: ImagingConfig, pwi):
    return (s
            .filter_max_persent(icfg.bscan_max/1000, bid=True)
            .sum(cond=s.d == 3)
            .dct(icfg.dct_start, icfg.dct_end)
            .call(lambda m: m.astype(cp.int16))
            .call(pwi)
            .call(cp.asarray, order='C')
            .argrelextrema()
            .conv_guass(b=icfg.beta * 0.01)
            .clip(icfg.focus_start, icfg.focus_end)
            .filter_max_persent(icfg.focus_max/100, mmax=icfg.focus_mmax if icfg.uafm else None)
            .time_gain_compensation_linear_max(icfg.tgcl, mmax=icfg.focus_mmax if icfg.uafm else None)
            .cpu()
            .get()
            )


def tfm_process(s: ScanData, icfg: ImagingConfig, disable_cache: bool, tfm: TFM):
    # print(icfg.changed_field, icfg.changed_field in ['dct_start', 'dct_end', 'bscan_max'], disable_cache)
    if icfg.changed_field in ['dct_start', 'dct_end', 'bscan_max'] or tfm_cache[0] is None or disable_cache:
        tfm_cache[0] = (s.dct(icfg.dct_start, icfg.dct_end)
                        .filter_max_persent(icfg.bscan_max/1000, bid=True)
                        .call(lambda m: m.astype(cp.int16))
                        .call(tfm))
    return (tfm_cache[0]
            .call(cp.asarray, order='C')
            .argrelextrema(axis=1)
            .conv_guass(b=icfg.beta * 0.01, axis=1)
            .clip(icfg.focus_start, icfg.focus_end)
            .filter_max_persent(icfg.focus_max, mmax=icfg.focus_mmax if icfg.uafm else None)
            .time_gain_compensation_linear_max(icfg.tgcl, mmax=icfg.focus_mmax if icfg.uafm else None)
            .cpu()
            .get()
            )