# -*- coding: utf-8 -*-
# Copyright 2019 The pyXem developers
#
# This file is part of pyXem.
#
# pyXem is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# pyXem is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with pyXem.  If not, see <http://www.gnu.org/licenses/>.

import numpy as np
import pytest

from pyxem.utils.segment_utils import (
    norm_cross_corr,
    separate_watershed,
    get_gaussian2d,
)


@pytest.mark.parametrize(
    "img, template, corr_expt",
    [
        (
            np.array([[0.1, 50.0, 0.0], [0.0, 0.0, 44], [0.3, 37, 11.1]]),
            np.array([[0.2, 0.0, 1.4], [10.9, 22.0, 15.0], [0.28, 36, 0.36]]),
            0.28918685246563913,
        ),
        (
            np.array([[1, 9, 5], [42, 42, 44], [6, 37, 11]]),
            np.array([[0, 1, 0], [1, 0, 0], [40, 90, 35]]),
            0.075079495399324,
        ),
        (
            np.array([[1, 0, 1], [0, 1, 1], [1, 0, 0]]),
            np.array([[0, 1, 0], [1, 0, 0], [0, 1, 1]]),
            -1.0,
        ),
        (
            np.array([[1, 9, 5], [42, 42, 44], [6, 37, 11]]),
            np.array([[1, 9, 5], [42, 42, 44], [6, 37, 11]]),
            1.0,
        ),
    ],
)
def test_norm_cross_corr(img, template, corr_expt):
    c = norm_cross_corr(img, template)
    np.testing.assert_allclose(c, corr_expt)


@pytest.fixture
def vdf_image():
    stest = np.zeros((7, 6))
    stest[1:3, :2] = 1
    stest[2:5, 1:3] = 2
    stest[1:3, 4:] = 3
    stest[4, 4:] = 4
    stest[6:, :] = 5
    return stest


@pytest.mark.parametrize(
    "min_distance, min_size, max_size, "
    "max_number_of_grains, exclude_border, threshold, "
    "sep_shape_expt",
    [
        (1, 1, np.inf, np.inf, 0, True, (4, 6, 7)),
        (1, 1, np.inf, np.inf, 1, False, (1, 6, 7)),
        (1, 5, np.inf, np.inf, 0, True, (2, 6, 7)),
        (1, 1, np.inf, 4, 0, True, (3, 6, 7)),
        (1, 1, 3, np.inf, 0, True, (1, 6, 7)),
        (3, 1, np.inf, np.inf, 0, True, (1, 6, 7)),
    ],
)
def test_separate_watershed(
    vdf_image,
    min_distance,
    min_size,
    max_size,
    max_number_of_grains,
    exclude_border,
    threshold,
    sep_shape_expt,
):
    sep = separate_watershed(
        vdf_image,
        min_distance=min_distance,
        min_size=min_size,
        max_size=max_size,
        max_number_of_grains=max_number_of_grains,
        exclude_border=exclude_border,
        threshold=threshold,
    )
    np.testing.assert_equal(sep.shape, sep_shape_expt)


@pytest.mark.parametrize(
    "a, xo, yo, x, y, sigma, gauss_shape_expt",
    [
        (10, 5, 5, np.arange(10), np.arange(10), 0.2, (10,)),
        (
            10,
            5,
            5,
            np.tile(np.arange(10), (5, 1)),
            np.tile(np.arange(10), (5, 1)),
            0.2,
            (5, 10),
        ),
    ],
)
def test_get_gaussian2d(a, xo, yo, x, y, sigma, gauss_shape_expt):
    gauss = get_gaussian2d(a, xo, yo, x, y, sigma)
    assert isinstance(gauss, np.ndarray)
    assert gauss.dtype == float
    np.testing.assert_equal(gauss.shape, gauss_shape_expt)
