import unittest
from functools import partial

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

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):
    def test_iris(self):
        x = datasets.load_iris().data
        x += np.random.normal(0, 1e-3, x.shape)  # iris contains duplicate rows

        distances = squareform(pdist(x))
        params = dict(initialization="random", random_state=0)
        embedding1 = TSNE(metric="precomputed", **params).fit(distances)
        embedding2 = TSNE(metric="euclidean", **params).fit(x)

        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)
