"""Unit tests for Exponential Barycenter mean."""

import logging

import geomstats.backend as gs
import tests.conftest
from geomstats.geometry.euclidean import Euclidean
from geomstats.geometry.special_euclidean import SpecialEuclidean
from geomstats.geometry.special_orthogonal import SpecialOrthogonal
from geomstats.learning.exponential_barycenter import ExponentialBarycenter
from geomstats.learning.frechet_mean import FrechetMean


class TestExponentialBarycenter(tests.conftest.TestCase):
    def setup_method(self):
        logger = logging.getLogger()
        logger.disabled = True
        self.se_mat = SpecialEuclidean(n=3)
        self.so_vec = SpecialOrthogonal(n=3, point_type="vector")
        self.so = SpecialOrthogonal(n=3)
        self.n_samples = 4

    @tests.conftest.np_and_autograd_only
    def test_estimate_and_belongs_se(self):
        point = self.se_mat.random_point(self.n_samples)
        estimator = ExponentialBarycenter(self.se_mat)
        estimator.fit(point)
        barexp = estimator.estimate_
        result = self.se_mat.belongs(barexp)
        expected = True
        self.assertAllClose(result, expected)

        point = self.so_vec.random_uniform(self.n_samples)
        estimator = ExponentialBarycenter(self.so_vec)
        estimator.fit(point)
        barexp = estimator.estimate_
        result = self.so_vec.belongs(barexp)
        expected = True
        self.assertAllClose(result, expected)

    def test_estimate_one_sample_se(self):
        point = self.se_mat.random_point()
        estimator = ExponentialBarycenter(self.se_mat)
        estimator.fit(gs.expand_dims(point, axis=0))
        result = estimator.estimate_
        expected = point
        self.assertAllClose(result, expected)

        point = self.so_vec.random_uniform(1)
        estimator = ExponentialBarycenter(self.so_vec)
        estimator.fit(gs.expand_dims(point, axis=0))
        result = estimator.estimate_
        expected = point
        self.assertAllClose(result, expected)

    @tests.conftest.np_and_autograd_only
    def test_estimate_and_reach_max_iter_se(self):
        point = self.se_mat.random_point(1)
        estimator = ExponentialBarycenter(self.se_mat).set(max_iter=2)
        points = gs.array([point, point])
        estimator.fit(points)
        result = estimator.estimate_
        expected = point
        self.assertAllClose(result, expected)

        point = self.so_vec.random_uniform(1)
        estimator = ExponentialBarycenter(self.so_vec).set(max_iter=2)
        points = gs.array([point, point])
        estimator.fit(points)
        result = estimator.estimate_
        expected = point
        self.assertAllClose(result, expected)

    @tests.conftest.np_and_autograd_only
    def test_estimate_so_matrix(self):
        points = self.so.random_uniform(2)

        mean_vec = ExponentialBarycenter(space=self.so)
        mean_vec.fit(points)

        logs = self.so.log(points, mean_vec.estimate_)
        result = gs.sum(logs, axis=0)
        expected = gs.zeros_like(points[0])
        self.assertAllClose(result, expected)

    @tests.conftest.np_and_autograd_only
    def test_estimate_and_belongs_so(self):
        point = self.so.random_uniform(self.n_samples)
        estimator = ExponentialBarycenter(self.so)
        estimator.fit(point)
        barexp = estimator.estimate_
        result = self.so.belongs(barexp)
        expected = True
        self.assertAllClose(result, expected)

        point = self.so_vec.random_uniform(self.n_samples)
        estimator = ExponentialBarycenter(self.so_vec)
        estimator.fit(point)
        barexp = estimator.estimate_
        result = self.so_vec.belongs(barexp)
        expected = True
        self.assertAllClose(result, expected)

    @tests.conftest.np_and_autograd_only
    def test_estimate_one_sample_so(self):
        point = self.so.random_uniform(1)
        estimator = ExponentialBarycenter(self.so)
        estimator.fit(gs.expand_dims(point, axis=0))
        result = estimator.estimate_
        expected = point
        self.assertAllClose(result, expected)

        point = self.so_vec.random_uniform(1)
        estimator = ExponentialBarycenter(self.so_vec)
        estimator.fit(gs.expand_dims(point, axis=0))
        result = estimator.estimate_
        expected = point
        self.assertAllClose(result, expected)

    @tests.conftest.np_and_autograd_only
    def test_estimate_and_reach_max_iter_so(self):
        point = self.so.random_uniform(self.n_samples)
        estimator = ExponentialBarycenter(self.so).set(max_iter=2)
        estimator.fit(point)
        barexp = estimator.estimate_
        result = self.so.belongs(barexp)
        expected = True
        self.assertAllClose(result, expected)

        point = self.so_vec.random_uniform(self.n_samples)
        estimator = ExponentialBarycenter(self.so_vec).set(max_iter=2)
        estimator.fit(point)
        barexp = estimator.estimate_
        result = self.so_vec.belongs(barexp)
        expected = True
        self.assertAllClose(result, expected)

    @tests.conftest.np_and_autograd_only
    def test_coincides_with_frechet_so(self):
        gs.random.seed(0)
        point = self.so.random_uniform(self.n_samples)

        estimator = ExponentialBarycenter(self.so).set(max_iter=40, epsilon=1e-10)
        estimator.fit(point)
        result = estimator.estimate_

        frechet_estimator = FrechetMean(self.so, method="adaptive").set(
            max_iter=40,
            epsilon=1e-10,
            init_step_size=1.0,
        )
        frechet_estimator.fit(point)
        expected = frechet_estimator.estimate_

        self.assertAllClose(result, expected, atol=1e-6)

    @tests.conftest.np_and_autograd_only
    def test_estimate_weights(self):
        point = self.so.random_uniform(self.n_samples)
        estimator = ExponentialBarycenter(self.so).set(verbose=True)
        weights = gs.arange(self.n_samples)
        estimator.fit(point, weights=weights)
        barexp = estimator.estimate_
        result = self.so.belongs(barexp)
        expected = True
        self.assertAllClose(result, expected)

        point = self.so_vec.random_uniform(self.n_samples)
        estimator = ExponentialBarycenter(self.so_vec)
        estimator.fit(point, weights=weights)
        barexp = estimator.estimate_
        result = self.so_vec.belongs(barexp)
        expected = True
        self.assertAllClose(result, expected)

    def test_linear_mean(self):
        euclidean = Euclidean(3)
        point = euclidean.random_point(self.n_samples)

        estimator = ExponentialBarycenter(euclidean)

        estimator.fit(point)
        result = estimator.estimate_

        expected = gs.mean(point, axis=0)

        self.assertAllClose(result, expected)
