#!/usr/bin/env python3

"""
Unit tests for functions in detection.py
cmdoret, 20200904
"""
import numpy as np
import scipy.sparse as sp
import pytest
import pareidolia.detection as pad

# Run test on square and rectangle matrices with even and odd dimensions,
MAP_PARAMS = ("shape", ((2, 2), (3, 3), (3, 5), (5, 3), (2, 5), (5, 2)))


def gen_mats(shape, n_mats=4, density=0.8):
    """
    Helper function generating a list of random sparse matrices with identical
    sparsity structure.
    """
    print(shape)
    # Get total number of pixels in the matrix
    n_coords = shape[0] * shape[1]
    # Pick a number of (1D) pixel indices to switch on
    picked_coords = np.random.choice(n_coords, round(density * n_coords), replace=False)
    # Convert 1D coords to 2D
    picked_rows = picked_coords // shape[1]
    picked_cols = picked_coords % shape[1]
    # Generate empty sparse matrices and fill the right pixels
    mats = [None] * n_mats
    for i in range(n_mats):
        mat = sp.csr_matrix(shape)
        mat[picked_rows, picked_cols] = np.random.random(size=len(picked_coords))
        mats[i] = mat

    return mats


@pytest.mark.parametrize(*MAP_PARAMS)
def test_median_bg(shape):
    """Test computation of median background from multiple sparse matrices."""
    # Generate multiple matrices of identical shape and sparsity
    mats = gen_mats(shape)
    # Compute median background
    obs_bg = pad.median_bg(mats)
    # Use numpy's utilities to compute the median background in dense format
    dense_mats = [m.toarray() for m in mats]
    exp_bg = np.apply_along_axis(np.median, 2, np.dstack(dense_mats))
    # The median bg should be the same as numpy's result on dense matrices
    assert np.all(obs_bg == exp_bg)
    bad_mats = mats.copy()
    # A sparse matrix with different format should yield ValueError
    with pytest.raises(ValueError):
        bad_mats[0] = mats[0].tocoo()
        pad.median_bg(bad_mats)
    with pytest.raises(ValueError):  # Same for sparsity structure
        bad_mats[0] = mats[0]
        bad_mats[0].data[0] = 0
        bad_mats[0].eliminate_zeros()
        pad.median_bg(bad_mats)
    with pytest.raises(ValueError):  # Same for dense matrix
        bad_mats[0] = mats[0].toarray()
        pad.median_bg(bad_mats)


@pytest.mark.parametrize(*MAP_PARAMS)
def test_reps_bg_diff(shape):
    """Test computation of distribution of differences to median background"""
    mats = gen_mats(shape)
    obs_diff = pad.reps_bg_diff(mats)
    # Test the number of points in the distribution
    assert obs_diff.shape[0] == len(mats[0].data) * len(mats)


@pytest.mark.parametrize(*MAP_PARAMS)
def test_get_sse_mat(shape):
    """Test generation of SSE matrix with sparse operations"""
    # Compare with dense equivalent
    mats = gen_mats(shape)
    bg_dense = pad.median_bg(mats).toarray()
    se_dense = np.array([(m.toarray() - bg_dense) ** 2 for m in mats])
    sse_dense = np.sum(se_dense, axis=0)
    sse_sparse = pad.get_sse_mat(mats)
    assert np.all(sse_dense == sse_sparse)


@pytest.mark.parametrize(*MAP_PARAMS)
def test_get_win_density_shape(shape):
    """Test output shape when computing local pixel density"""
    mats = gen_mats(shape)
    for mat in mats:
        density = pad.get_win_density(mat)
        # Check if shape is correct
        assert density.shape == mat.shape
        # Density cannot be smaller than in input mat
        assert density.nnz >= mat.nnz
        # Check if lower triangle is empty when running
        # in symmetric upper mode
        density = pad.get_win_density(mat, sym_upper=True)
        assert density.sum() == sp.triu(density).sum()
        assert density.shape == mat.shape


@pytest.mark.parametrize("density", (0.1, 0.5, 0.8))
def test_get_win_density_values(density):
    """Fuzzy test the results of local pixel density"""
    # Perform local density estimation on a few large matrices
    mats = gen_mats((1000, 1000), n_mats=5, density=density)
    obs_density = np.mean([pad.get_win_density(m).mean() for m in mats])
    # Check if average density match expectations
    assert abs(obs_density - density) < 0.01
