"""Tests for E-Agglo (agglomerative clustering algorithm)."""

__maintainer__ = []

import numpy as np
import pandas as pd

from aeon.segmentation import EAggloSegmenter


def test_fit_default_params_univariate():
    """Test univariate data and default parameters.

    These numbers are generated from the original implementation
    in R, with the following code:

    set.seed(1234)
    X <- c(rnorm(15, mean = -6), 0, rnorm(16, mean = 6))
    X <- as.matrix(X[c(1,2,17,18)])
    ret = e.agglo(X)
    """
    X = pd.Series([-7.207066, -5.722571, 5.889715, 5.488990])

    cluster_expected = [0, 0, 1, 1]
    fit_expected = [104.77424, 134.51387, 186.92586, -31.15431]

    model = EAggloSegmenter()
    fitted_model = model.fit(X)

    cluster_actual = fitted_model.cluster_
    fit_actual = fitted_model.gof_

    assert np.allclose(cluster_actual, cluster_expected)
    assert np.allclose(fit_actual, fit_expected)


def test_fit_other_params_univariate():
    """Test univariate data with alternative starting clusters."""
    X = np.array([-7.207066, -5.722571, 5.889715, 5.488990])

    cluster_expected = [0, 0, 1, 1]
    fit_expected = [1182.754, 1772.526, -295.421]

    model = EAggloSegmenter(member=np.array([0, 0, 1, 2]), alpha=2)
    fitted_model = model.fit(X)
    cluster_actual = fitted_model.cluster_
    fit_actual = fitted_model.gof_

    assert np.allclose(cluster_actual, cluster_expected)
    assert np.allclose(fit_actual, fit_expected)


def test_fit_default_params_multivariate():
    """Test multivariate data with default parameters.

    These numbers are generated from the original implementation
    in R, with the following code:

    set.seed(1234)
    X <- c(rnorm(15, mean = -6), 0, rnorm(16, mean = 6))
    X <- as.matrix(cbind(X[c(1,2,17,18,19)],X[c(3,4,5,6,7)]))
    ret = e.agglo(X)
    """
    X = np.array(
        [
            [-7.207, -4.916],
            [-5.723, -8.346],
            [5.890, -5.571],
            [5.489, -5.494],
            [5.089, -6.575],
        ]
    )

    cluster_expected = [0, 0, 1, 1, 1]
    fit_expected = [118.58235, 132.67919, 156.54531, 208.61512, -33.52743]

    model = EAggloSegmenter()
    fitted_model = model.fit(X, axis=0)

    cluster_actual = fitted_model.cluster_
    fit_actual = fitted_model.gof_

    assert np.allclose(cluster_actual, cluster_expected)
    assert np.allclose(fit_actual, fit_expected)


def test_len_penalty():
    """Test multivariate data with penalty function as string input."""
    X = pd.DataFrame(
        [
            [-7.207, -4.916],
            [-5.723, -8.346],
            [5.890, -5.571],
            [5.489, -5.494],
            [5.089, -6.575],
        ]
    )

    cluster_expected = [0, 0, 1, 1, 1]
    fit_expected = [112.58235, 127.67919, 152.54531, 205.61512, -35.52743]

    model = EAggloSegmenter(penalty="len_penalty")
    fitted_model = model._fit(X)

    cluster_actual = fitted_model.cluster_
    fit_actual = fitted_model.gof_

    assert np.allclose(cluster_actual, cluster_expected)
    assert np.allclose(fit_actual, fit_expected)


def test_custom_penalty():
    """Test multivariate data with functional input as penalty."""
    X = pd.DataFrame([-7.207066, -5.722571, 5.889715, 5.488990])

    cluster_expected = [0, 0, 1, 1]
    fit_expected = [105.77424, 135.84720, 188.92586, -29.15431]

    model = EAggloSegmenter(penalty=lambda x: np.mean(np.diff(np.sort(x))))
    fitted_model = model._fit(X)

    cluster_actual = fitted_model.cluster_
    fit_actual = fitted_model.gof_

    assert np.allclose(cluster_actual, cluster_expected)
    assert np.allclose(fit_actual, fit_expected)
