import unittest
from functools import partial

import scipy.sparse as sp
from scipy.spatial.distance import pdist, squareform

import openTSNE
import openTSNE.affinity
import openTSNE.initialization
import numpy as np
from openTSNE.callbacks import VerifyExaggerationError
from sklearn import datasets
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.manifold import SpectralEmbedding
import platform

machine = platform.machine()

TSNE = partial(openTSNE.TSNE, neighbors="exact", negative_gradient_method="bh")

class TestTSNECorrectness(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.tsne = TSNE(early_exaggeration_iter=20, n_iter=100)
        # Set up two modalities, if we want to viually inspect test results
        random_state = np.random.RandomState(0)
        cls.x = np.vstack(
            (random_state.normal(+1, 1, (100, 4)), random_state.normal(-1, 1, (100, 4)))
        )
        cls.x_test = random_state.normal(0, 1, (25, 4))
        cls.iris = datasets.load_iris()

    def test_basic_flow(self):
        """Verify that the basic flow does not crash."""
        embedding = self.tsne.fit(self.x)
        self.assertFalse(np.any(np.isnan(embedding)))

        partial_embedding = embedding.transform(self.x_test, n_iter=20)
        self.assertFalse(np.any(np.isnan(partial_embedding)))

    def test_advanced_flow(self):
        """Verify that the advanced flow does not crash."""
        embedding = self.tsne.prepare_initial(self.x)
        embedding = embedding.optimize(20, exaggeration=12)
        embedding = embedding.optimize(20)  # type: openTSNE.TSNEEmbedding
        self.assertFalse(np.any(np.isnan(embedding)))

        partial_embedding = embedding.prepare_partial(self.x_test)
        partial_embedding = partial_embedding.optimize(20, exaggeration=2)
        partial_embedding = partial_embedding.optimize(20)
        self.assertFalse(np.any(np.isnan(partial_embedding)))

    def test_error_exaggeration_correction(self):
        embedding = self.tsne.prepare_initial(self.x)

        # The callback raises if the KL divergence does not match the true one
        embedding.optimize(
            50,
            exaggeration=5,
            callbacks=[VerifyExaggerationError(embedding)],
            callbacks_every_iters=1,
            inplace=True,
        )

    def test_iris(self):
        x, y = self.iris.data, self.iris.target

        # Evaluate t-SNE optimization using a KNN classifier
        knn = KNeighborsClassifier(n_neighbors=10)
        tsne = TSNE(perplexity=30, initialization="random", random_state=0)

        # Prepare a random initialization
        embedding = tsne.prepare_initial(x)

        # KNN should do poorly on a random initialization
        knn.fit(embedding, y)
        predictions = knn.predict(embedding)
        self.assertLess(accuracy_score(predictions, y), 0.5)

        # Optimize the embedding for a small number of steps so tests run fast
        embedding.optimize(250, inplace=True)

        # Similar points should be grouped together, therefore KNN should do well
        knn.fit(embedding, y)
        predictions = knn.predict(embedding)
        self.assertGreater(accuracy_score(predictions, y), 0.95)

    def test_iris_with_precomputed_distance_matrices(self):
        x, y = self.iris.data, self.iris.target

        distances = squareform(pdist(x))

        # Evaluate t-SNE optimization using a KNN classifier
        knn = KNeighborsClassifier(n_neighbors=10)
        tsne = TSNE(
            perplexity=30, initialization="random", random_state=0, metric="precomputed"
        )

        # Prepare a random initialization
        embedding = tsne.prepare_initial(distances)

        # KNN should do poorly on a random initialization
        knn.fit(embedding, y)
        predictions = knn.predict(embedding)
        self.assertLess(accuracy_score(predictions, y), 0.5)

        # Optimize the embedding for a small number of steps so tests run fast
        embedding.optimize(250, inplace=True)

        # Similar points should be grouped together, therefore KNN should do well
        knn.fit(embedding, y)
        predictions = knn.predict(embedding)
        self.assertGreater(accuracy_score(predictions, y), 0.95)

    def test_iris_bh_transform_equivalency_with_one_by_one(self):
        """Compare one by one embedding vs all at once using BH gradients."""
        x_train, x_test = train_test_split(
            self.iris.data, test_size=0.1, random_state=42
        )

        # Set up the initial embedding
        embedding = openTSNE.TSNE(
            early_exaggeration_iter=0,
            n_iter=50,
            neighbors="exact",
            negative_gradient_method="bh",
        ).fit(x_train)

        params = dict(n_iter=100, perplexity=5)
        # Build up an embedding by adding points one by one
        new_embedding_1 = np.vstack(
            [embedding.transform(np.atleast_2d(point), **params) for point in x_test]
        )
        # Add new points altogether
        new_embedding_2 = embedding.transform(x_test, **params)

        # Verify that the embedding has actually been optimized
        self.assertRaises(
            AssertionError,
            np.testing.assert_almost_equal,
            embedding.prepare_partial(x_test, perplexity=params["perplexity"]),
            new_embedding_1,
        )
        # Check that both methods produced the same embedding
        np.testing.assert_almost_equal(new_embedding_1, new_embedding_2)

    def test_iris_fft_transform_equivalency_with_one_by_one(self):
        """Compare one by one embedding vs all at once using FFT gradients.

        Note that this won't return the exact same embedding both times because
        the grid placed over the embedding will differ when placing points one
        by one vs. when placing them at once. The min/max coords will differ,
        thus changing the overall approximation. They should be quite similar
        though.

        """
        x_train, x_test = train_test_split(
            self.iris.data, test_size=0.1, random_state=42
        )

        # Set up the initial embedding
        embedding = openTSNE.TSNE(
            early_exaggeration_iter=0,
            n_iter=50,
            neighbors="exact",
            negative_gradient_method="fft",
        ).fit(x_train)

        # Changing the gradients using clipping changes how the points move
        # sufficiently so that the interpolation grid is shifted. This test is
        # more reliable when we don't do gradient clipping and reduce the
        # learning rate. We increase the number of iterations so that the points
        # have time to move around
        params = dict(perplexity=5)
        # Build up an embedding by adding points one by one
        new_embedding_1 = np.vstack(
            [embedding.transform(np.atleast_2d(point), **params) for point in x_test]
        )
        # Add new points altogether
        new_embedding_2 = embedding.transform(x_test, **params)

        # Verify that the embedding has actually been optimized
        self.assertRaises(
            AssertionError,
            np.testing.assert_almost_equal,
            embedding.prepare_partial(x_test, perplexity=params["perplexity"]),
            new_embedding_1,
        )

        # Check that both methods produced the same embedding
        np.testing.assert_almost_equal(new_embedding_1, new_embedding_2, decimal=2)

    def test_iris_bh_transform_correctness(self):
        x_train, x_test, y_train, y_test = train_test_split(
            self.iris.data, self.iris.target, test_size=0.33, random_state=42
        )

        # Set up the initial embedding
        embedding = openTSNE.TSNE(
            neighbors="exact",
            negative_gradient_method="bh",
            early_exaggeration_iter=0,
            n_iter=50,
            random_state=0,
        ).fit(x_train)

        # Evaluate t-SNE optimization using a KNN classifier
        knn = KNeighborsClassifier(n_neighbors=10)
        knn.fit(embedding, y_train)

        new_embedding = embedding.transform(x_test, n_iter=100)
        predictions = knn.predict(new_embedding)
        self.assertGreater(accuracy_score(predictions, y_test), 0.95)

    def test_iris_fft_transform_correctness(self):
        x_train, x_test, y_train, y_test = train_test_split(
            self.iris.data, self.iris.target, test_size=0.33, random_state=42
        )

        # Set up the initial embedding
        embedding = openTSNE.TSNE(
            neighbors="exact",
            negative_gradient_method="fft",
            early_exaggeration_iter=0,
            n_iter=50,
            random_state=0,
        ).fit(x_train)

        # Evaluate t-SNE optimization using a KNN classifier
        knn = KNeighborsClassifier(n_neighbors=10)
        knn.fit(embedding, y_train)

        new_embedding = embedding.transform(x_test, n_iter=100)
        predictions = knn.predict(new_embedding)
        self.assertGreater(accuracy_score(predictions, y_test), 0.95)

    def test_bh_transform_with_point_subsets_using_perplexity_nn(self):
        x_train, x_test = train_test_split(
            self.iris.data, test_size=0.33, random_state=42
        )

        # Set up the initial embedding
        init = openTSNE.initialization.pca(x_train)
        affinity = openTSNE.affinity.PerplexityBasedNN(x_train, method="exact")
        embedding = openTSNE.TSNEEmbedding(
            init, affinity, negative_gradient_method="bh", random_state=42
        )
        embedding.optimize(n_iter=50, inplace=True)

        # The test set contains 50 samples, so let's verify on half of those
        transform_params = dict(n_iter=0, early_exaggeration_iter=0)
        new_embedding_1 = embedding.transform(x_test, **transform_params)[:25]
        new_embedding_2 = embedding.transform(x_test[:25], **transform_params)

        np.testing.assert_equal(new_embedding_1, new_embedding_2)

    def test_fft_transform_with_point_subsets_using_perplexity_nn(self):
        x_train, x_test = train_test_split(
            self.iris.data, test_size=0.33, random_state=42
        )

        # Set up the initial embedding
        init = openTSNE.initialization.pca(x_train)
        affinity = openTSNE.affinity.PerplexityBasedNN(x_train, method="exact")
        embedding = openTSNE.TSNEEmbedding(
            init, affinity, negative_gradient_method="fft", random_state=42
        )
        embedding.optimize(n_iter=100, inplace=True)

        # The test set contains 50 samples, so let's verify on half of those
        transform_params = dict(n_iter=0, early_exaggeration_iter=0)
        new_embedding_1 = embedding.transform(x_test, **transform_params)[:25]
        new_embedding_2 = embedding.transform(x_test[:25], **transform_params)

        np.testing.assert_equal(new_embedding_1, new_embedding_2)


class TestTSNECorrectnessUsingNonStandardDof(TestTSNECorrectness):
    @classmethod
    def setUpClass(cls):
        cls.tsne = TSNE(early_exaggeration_iter=20, n_iter=100, dof=0.8)
        # Set up two modalities, if we want to viually inspect test results
        random_state = np.random.RandomState(0)
        cls.x = np.vstack(
            (random_state.normal(+1, 1, (100, 4)), random_state.normal(-1, 1, (100, 4)))
        )
        cls.x_test = random_state.normal(0, 1, (25, 4))
        cls.iris = datasets.load_iris()


class TestTSNECorrectnessUsingPrecomputedDistanceMatrix(unittest.TestCase):
    # remove this once #1004706 fixed by upstream
    @unittest.skipIf(('i686' in machine) or machine == 'i386'
                       or machine == 'ppc64le',
                     reason=f"Skipping test on {machine} architecture")
    def test_iris(self):
        x = datasets.load_iris().data
        print(platform.machine())
        x += np.random.normal(0, 1e-3, x.shape)  # iris contains duplicate rows

        # We run this for only a few iterations since this will check for
        # correctness. If we let it run for longer, this test fails. The reason
        # it fails is that slight differences (e.g. 16th decimal) in the
        # distance matrices produce slightly different P matrices (detectable at
        # precision 18 decimals), which compounds during optimization, resulting
        # in slightly different embeddings (visually indistinguishable). If the
        # computation was, however, wrong, we would see a difference after only
        # a few iterations. Early exaggeration appears to have a much stronger
        # effect on this compounting, so we disable it here.
        # See also: https://github.com/pavlin-policar/openTSNE/issues/247
        distances = squareform(pdist(x))
        params = dict(
            early_exaggeration_iter=0,
            n_iter=500,
            initialization="random",
            random_state=0,
        )
        embedding1 = TSNE(metric="precomputed", **params).fit(distances)
        embedding2 = TSNE(metric="euclidean", **params).fit(x)

        self.assertTrue(
            sp.linalg.norm(embedding1.affinities.P - embedding2.affinities.P) < 1e-16
        )
        np.testing.assert_almost_equal(embedding1, embedding2)


class TestSpectralInitializationCorrectness(unittest.TestCase):
    def test_spectral_agreement_with_sklearn(self):
        # Generate some random data and stretch it, to give it some structure
        np.random.seed(42)
        x = np.random.randn(100, 20)
        x[:,0] *= 5

        # Perform spectral embedding via sklearn and via openTSNE
        P = openTSNE.affinity.PerplexityBasedNN(x).P
        embedding1 = openTSNE.initialization.spectral(P, tol=0, add_jitter=False)
        embedding2 = SpectralEmbedding(affinity='precomputed').fit_transform(P)

        np.testing.assert_almost_equal(
            np.abs(np.corrcoef(embedding1[:,0], embedding2[:,0])[0,1]), 1
        )
        np.testing.assert_almost_equal(
            np.abs(np.corrcoef(embedding1[:,1], embedding2[:,1])[0,1]), 1
        )


class TestEarlyExaggerationCollapse(unittest.TestCase):
    """In some cases, the BH implementation was producing a collapsed embedding
    for all data points. For more information, see #233, #234."""
    def test_early_exaggeration_does_not_collapse(self):
        n_samples = [100, 150, 200]
        n_dims = [5, 10, 20]

        np.random.seed(42)
        for n in n_samples:
            for d in n_dims:
                x = np.random.randn(n, d)
                embedding = openTSNE.TSNE(random_state=42).fit(x)
                self.assertGreater(np.max(np.abs(embedding)), 1e-8)


class TestDataMatricesWithDuplicatedRows(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        from sklearn.preprocessing import KBinsDiscretizer

        # Load up contrived example where we have a large number of duplicated
        # rows. This is similar to the Titanic data set, which is problematic.
        np.random.seed(0)
        iris = datasets.load_iris()
        x, y = iris.data, iris.target

        discretizer = KBinsDiscretizer(n_bins=2, strategy="uniform")
        x = discretizer.fit_transform(x).toarray()

        idx = np.random.choice(x.shape[0], size=1000, replace=True)
        cls.x, cls.y = x[idx], y[idx]

    def test_works_without_error(self):
        openTSNE.TSNE(
            early_exaggeration=100, negative_gradient_method="bh", random_state=0
        ).fit(self.x)