#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Author        : Lancelot PINCET
# GitHub        : https://github.com/LancelotPincet



# %% Libraries
from smlmlp import block
from arrlp import get_xp, img_correlate
import numpy as np



# %% Function
@block()
def signal_spatial_filter(channels, /, spatial_kernel, signals=None, bkgds=None, noise_corrections=None, *, cuda=False, parallel=False) :
    '''
    This function applyies a spatial filter to enhance signal.
    '''

    # Correct signal length for end of acquisition
    if signals is not None and len(signals[0]) > len(channels[0]):
        signals = [signal[:len(channel)] for channel, signal in zip(channels, signals)]
    
    # Noise corrections
    if noise_corrections is None :
        noise_corrections = [np.float32(1.) for _ in range(len(channels))]

    new_signals = []
    for i in range(len(channels)) :
        bkgd = None if bkgds is None else bkgds[i]
        signal = None if signals is None else signals[i]
        channel = channels[i]
        if bkgd is not None :
            channel = channel - bkgd
        if signal is channel : signal = None
        kernel = spatial_kernel[i]
        new_signal = img_correlate(channel, kernel=kernel, out=signal, cuda=cuda, parallel=parallel, stacks=True)
        noise_corrections[i] *= np.sqrt(np.sum(kernel**2))
        new_signals.append(new_signal)

    return new_signals, noise_corrections