

import numpy as np
import scipy.ndimage as ndi
import scipy.stats as stats


def nanmean(x):
    n = np.sum(np.isfinite(x))
    return np.nanmean(x) if n>0 else np.nan


def do_test(mean_s, mean_b, var_rand, var_sys, size, q, dtype='f4'):
    """
    mean_s:    mean of sample
    mean_b:    mean of background
    var_rand:  random variance of sample
    var_sys:   systematic variance of sample
    size:      size of sample
    q:         threshold
    """
    mean_s = np.array(mean_s)
    size = np.array(size)

    z_values = np.full(mean_s.shape, np.nan)
    m = size > 0

    z_values[m] = mean_s[m] - mean_b
    z_values[m] /= np.sqrt(var_rand / size[m] + var_sys)

    return z_values.astype(dtype), z_values > stats.norm.ppf(q)


def welch_satterthwaite_equation(s1, s2, n1, n2):
    """ Estimate degree of freedom. """
    a = (s1**2 / n1 + s2**2 / n2)**2
    b = s1**4 / (n1**2 * (n1-1)) + s2**2 / (n2**2 * (n2-1))
    return a/b


def local_mean(img, size):
    """
    neighborhood size: 1,5,9,13,25
    """
    if size == 1:
        return img.copy(), np.ones_like(img)

    elif size == 5:
        footprint = np.array([
            [0, 1, 0],
            [1, 1, 1],
            [0, 1, 0]
        ])
    elif size == 13:
        footprint = np.array([
            [0, 0, 1, 0, 0],
            [0, 1, 1, 1, 0],
            [1, 1, 1, 1, 1],
            [0, 1, 1, 1, 0],
            [0, 0, 1, 0, 0],
        ])
    elif size == 37:
        footprint = np.array([
            [0,0,1,1,1,0,0],
            [0,1,1,1,1,1,0],
            [1,1,1,1,1,1,1],
            [1,1,1,1,1,1,1],
            [1,1,1,1,1,1,1],
            [0,1,1,1,1,1,0],
            [0,0,1,1,1,0,0],
        ])
    elif size in [9,25,49,81,121]:
        size = int(np.sqrt(size))
        footprint = np.ones((size,size))
    else:
        raise ValueError('size of neighborhood %d not supported' % size)

    # only compute if any values in across-track direction
    domain = np.any(img, axis=1)

    # keep dtype of image
    footprint = footprint.astype(img.dtype)

    mean = np.full(img.shape, np.nan, dtype=img.dtype)
    mean[domain,:] = ndi.generic_filter(
            img[domain,:], nanmean, footprint=footprint, mode='constant', cval=np.nan
    )

    n = np.full(img.shape, np.nan, dtype=int)
    n[domain,:] = ndi.generic_filter(
            np.isfinite(img).astype(int)[domain,:], np.count_nonzero,
            footprint=footprint, mode='constant', cval=0
    )

    return mean, n
