# ----------------------------------------------------------------------------
# Copyright (c) 2013--, scikit-bio development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE.txt, distributed with this software.
# ----------------------------------------------------------------------------

"""Tests for MMvec implementation."""

import os
import io
import sys
import unittest

import numpy as np
import numpy.testing as npt
import pandas as pd
import pandas.testing as pdt
from scipy.spatial.distance import pdist
from scipy.stats import spearmanr
from scipy.sparse import coo_array

from skbio.util import get_data_path
from skbio.stats.composition import clr_inv as softmax
from skbio.stats.ordination import mmvec, MMvec, MMvecResult
from skbio.stats.ordination._mmvec import random_multimodal, _MMvecModel


class TestRandomMultimodal(unittest.TestCase):
    """Tests for random_multimodal simulation helper."""

    def test_numerical_precision(self):
        """Generated outputs that match expected values."""
        res = random_multimodal(
            n_features_x=4, n_features_y=6, n_samples=8, n_components=2, seed=7
        )
        x_counts, y_counts, design, coefs, x_main, x_bias, y_main, y_bias = res

        x_counts_exp = pd.DataFrame([
            [2, 3, 2, 3],
            [1, 4, 2, 3],
            [3, 3, 1, 3],
            [1, 4, 2, 3],
            [2, 4, 2, 2],
            [2, 3, 3, 2],
            [3, 3, 1, 3],
            [1, 2, 3, 4],
        ], index=[
            f"sample_{i}" for i in range(8)
        ], columns=[
            f"x_feature_{j}" for j in range(4)
        ], dtype=float)
        pdt.assert_frame_equal(x_counts, x_counts_exp)

        y_counts_exp = pd.DataFrame([
            [14, 10, 12, 19, 20, 25],
            [14, 13, 10, 11, 21, 31],
            [ 7, 21, 14, 13, 25, 20],
            [11, 15, 20,  9, 15, 30],
            [11, 20, 12,  7, 24, 26],
            [ 9, 13, 12,  9, 23, 34],
            [14, 22,  9, 19, 19, 17],
            [11,  7, 15, 14, 14, 39],
        ], index=[
            f"sample_{i}" for i in range(8)
        ], columns=[
            f"y_feature_{j}" for j in range(6)
        ], dtype=float)
        pdt.assert_frame_equal(y_counts, y_counts_exp)

        design_exp = np.vstack([np.ones(8), np.array([
            -1.0, -0.714, -0.429, -0.143, 0.143, 0.429, 0.714, 1.0
        ])]).T
        npt.assert_array_equal(design.round(3), design_exp)

        coefs_exp = np.array([
            [ 0.002,  0.597, -0.548, -1.781],
            [-0.909, -1.983,  0.120,  2.680],
        ])
        npt.assert_array_equal(coefs.round(3), coefs_exp)

        x_main_exp = np.array([
            [-0.979, -0.809],
            [ 1.061, -0.808],
            [-0.033,  0.884],
            [-0.584, -0.112]
        ])
        npt.assert_array_equal(x_main.round(3), x_main_exp)

        x_bias_exp = np.array([ 0.762, -1.199, 0.075, 0.577]).reshape(-1, 1)
        npt.assert_array_equal(x_bias.round(3), x_bias_exp)

        y_main_exp = np.array([
            [ 0.11 ,  0.064, -1.225,  0.076,  1.359],
            [-1.547,  0.859,  0.119, -0.641,  2.   ],
        ])
        npt.assert_array_equal(y_main.round(3), y_main_exp)

        y_bias_exp = np.array([-0.189,  0.683, -0.067,  0.667,  1.439]).reshape(1, -1)
        npt.assert_array_equal(y_bias.round(3), y_bias_exp)

        self.assertEqual(x_counts.shape, (8, 4))
        self.assertEqual(y_counts.shape, (8, 6))
        self.assertEqual(design.shape, (8, 2))
        self.assertEqual(coefs.shape, (2, 4))
        self.assertEqual(x_main.shape, (4, 2))
        self.assertEqual(x_bias.shape, (4, 1))
        self.assertEqual(y_main.shape, (2, 5))
        self.assertEqual(y_bias.shape, (1, 5))

        self.assertEqual(x_counts.index[0], "sample_0")
        self.assertEqual(x_counts.columns[0], "x_feature_0")
        self.assertEqual(y_counts.columns[0], "y_feature_0")

    def test_output_shapes(self):
        """Generated outputs should have expected shapes and IDs."""
        nx = 123  # number of X features
        ny = 456  # number of Y features
        ns = 78   # number of samples
        nd = 9    # number of dimensions

        res = random_multimodal(
            n_features_x=nx, n_features_y=ny, n_samples=ns, n_components=nd, seed=42
        )
        x_counts, y_counts, design, coefs, x_main, x_bias, y_main, y_bias = res

        self.assertEqual(x_counts.shape, (ns, nx))
        self.assertEqual(y_counts.shape, (ns, ny))
        self.assertEqual(design.shape, (ns, 2))
        self.assertEqual(coefs.shape, (2, nx))
        self.assertEqual(x_main.shape, (nx, nd))
        self.assertEqual(x_bias.shape, (nx, 1))
        self.assertEqual(y_main.shape, (nd, ny - 1))
        self.assertEqual(y_bias.shape, (1, ny - 1))

        self.assertListEqual(list(x_counts.index), [f"sample_{i}" for i in range(ns)])
        self.assertListEqual(
            list(x_counts.columns), [f"x_feature_{j}" for j in range(nx)]
        )
        self.assertListEqual(
            list(y_counts.columns), [f"y_feature_{k}" for k in range(ny)]
        )

    def test_reproducible_with_seed(self):
        """Same seed should provide reproducible output."""
        res1 = random_multimodal(seed=42)
        res2 = random_multimodal(seed=42)

        pdt.assert_frame_equal(res1[0], res2[0])
        pdt.assert_frame_equal(res1[1], res2[1])
        for i in range(2, len(res1)):
            npt.assert_allclose(res1[i], res2[i])


class TestMMvecRecovery(unittest.TestCase):
    """Test that MMvec recovers true embeddings from synthetic data."""

    # @unittest.skip("Skipping a test that requires long runtime.")
    def test_recovers_embeddings(self):
        """Verify model recovers true embeddings from synthetic data."""
        # Simulate a dataset of 150 samples
        X, Y, _, _, x_embed, x_bias, y_embed, y_bias = random_multimodal(
            8, 8, 150, 2, x_noise_std=2, x_total=1000, y_total=10000, seed=1
        )
        # Leave out last 10 samples as a test set
        n = 10
        X_train, X_test = X.iloc[:-n], X.iloc[-n:]
        Y_train, Y_test = Y.iloc[:-n], Y.iloc[-n:]

        result = mmvec(
            X_train,
            Y_train,
            dimensions=2,
            optimizer="adam",
            max_iter=500,  # may need to increase for harder cases
            batch_size=50,
            learning_rate=1e-3,
            beta_1=0.8,
            beta_2=0.9,
            seed=0,
            output_format="numpy",
        )

        # Check result type
        self.assertIsInstance(result, MMvecResult)

        # Verify X embeddings are recovered
        # Compare pairwise distances using Spearman correlation
        x_r, x_p = spearmanr(
            pdist(result.x_embeddings[:, :-1]),  # exclude bias
            pdist(x_embed),
        )
        self.assertGreater(x_r, 0.5, f"X embedding correlation too low: {x_r}")
        self.assertLess(x_p, 0.05, f"X embedding p-value too high: {x_p}")

        # Verify Y embeddings are recovered
        # Compare pairwise distances between Y features
        y_r, y_p = spearmanr(
            pdist(result.y_embeddings[1:, :-1]),  # exclude ref & bias
            pdist(y_embed.T),
        )
        self.assertGreater(y_r, 0.5, f"Y embedding correlation too low: {y_r}")
        self.assertLess(y_p, 0.05, f"Y embedding p-value too high: {y_p}")

        # Verify conditional probabilities match
        d1 = x_embed.shape[0]

        # Compute expected probabilities from true parameters
        x_aug = np.hstack((np.ones((d1, 1)), x_bias, x_embed))
        y_aug = np.vstack((y_bias, np.ones((1, y_embed.shape[1])), y_embed))
        exp = softmax(np.hstack((np.zeros((d1, 1)), x_aug @ y_aug)), validate=False)
        res = softmax(result.ranks, validate=False)

        s_r, s_p = spearmanr(np.ravel(res), np.ravel(exp))
        self.assertGreater(s_r, 0.5, f"Probability correlation too low: {s_r}")
        self.assertLess(s_p, 0.05, f"Probability p-value too high: {s_p}")

        # Compute Q^2 score on test data. The value should be positive for a reasonably
        # trained model (better than predicting the mean).
        q2 = result.score(X_test, Y_test)
        # Exact value may vary by platform and optimizer implementation, so we check
        # that it's above a reasonable threshold rather than exact.
        # self.assertEqual(round(q2, 7), 0.2664482)
        self.assertGreater(q2, -1.0, f"Q-squared score too low: {q2}")


class TestMMvecGradients(unittest.TestCase):
    """Test that analytical gradients match numerical gradients."""

    def test_full_model_gradients(self):
        """Verify all model gradients against numerical differentiation."""
        n_features_x = 5
        n_features_y = 8
        n_components = 2
        n_samples = 20

        # Create model
        rng = np.random.default_rng(42)
        model = _MMvecModel(
            n_features_x=n_features_x,
            n_features_y=n_features_y,
            n_components=n_components,
            x_prior_mean=0.0,
            x_prior_scale=1.0,
            y_prior_mean=0.0,
            y_prior_scale=1.0,
            rng=rng,
        )

        # Generate small test data
        X = np.random.randint(0, 100, size=(n_samples, n_features_x)).astype(
            np.float64
        )
        Y = np.random.randint(0, 100, size=(n_samples, n_features_y)).astype(
            np.float64
        )
        X_coo = coo_array(X)
        data = X_coo.data
        weights = data / data.sum()

        # Compute analytical gradients with fixed seed for reproducibility
        seed = 123
        size = 10
        norm = n_features_x / size
        _, grads = model.loss_and_grad(
            X_coo, Y, size, norm, weights, np.random.default_rng(seed)
        )
        grads = dict(zip(["x_main", "x_bias", "y_main", "y_bias"], grads))

        # Verify each parameter numerically
        # Use same seed for each call to get consistent batch
        eps = 1e-5
        for param_name in ["x_main", "x_bias", "y_main", "y_bias"]:
            param = getattr(model, param_name)
            numerical_grad = np.zeros_like(param)

            for idx in np.ndindex(param.shape):
                original = param[idx]
                param[idx] = original + eps
                loss_plus, _ = model.loss_and_grad(
                    X_coo, Y, size, norm, weights, np.random.default_rng(seed)
                )
                param[idx] = original - eps
                loss_minus, _ = model.loss_and_grad(
                    X_coo, Y, size, norm, weights, np.random.default_rng(seed)
                )
                param[idx] = original  # restore

                numerical_grad[idx] = (loss_plus - loss_minus) / (2 * eps)

            # Check relative error
            rel_error = np.abs(grads[param_name] - numerical_grad) / (
                np.abs(numerical_grad) + 1e-8
            )
            max_error = rel_error.max()
            self.assertLess(
                max_error,
                1e-3,
                f"{param_name}: max relative error {max_error}",
            )


class TestMMvecBasic(unittest.TestCase):
    """Basic functionality tests for MMvec."""

    def test_output_shapes(self):
        """Check that output shapes are correct."""
        ndim = 3
        X, Y, *_ = random_multimodal(10, 15, 50, ndim, seed=42)
        res = mmvec(X, Y, dimensions=ndim, max_iter=10, seed=42)
        n_features_x, n_features_y = X.shape[1], Y.shape[1]
        self.assertEqual(res.x_embeddings.shape, (n_features_x, ndim + 1))
        self.assertEqual(res.y_embeddings.shape, (n_features_y, ndim + 1))
        self.assertEqual(res.ranks.shape, (n_features_x, n_features_y))

    def test_ranks_row_centered(self):
        """Check that ranks are row-centered."""
        X, Y, *_ = random_multimodal(8, 10, 50, seed=42)
        obs = mmvec(X, Y, dimensions=2, max_iter=50, seed=42)
        row_means = obs.ranks.to_numpy().mean(axis=1)
        npt.assert_allclose(row_means, 0, atol=1e-10)

    def test_reproducibility(self):
        """Check that results are reproducible with same seed."""
        X, Y, *_ = random_multimodal(8, 10, 50, seed=42)
        obs1 = mmvec(X, Y, dimensions=2, max_iter=50, seed=42)
        obs2 = mmvec(X, Y, dimensions=2, max_iter=50, seed=42)
        pdt.assert_frame_equal(obs1.ranks, obs2.ranks, rtol=1e-10)

    def test_batch_norm_modes(self):
        """Test that both batch normalization modes work."""
        X, Y, *_ = random_multimodal(8, 10, 50, seed=42)
        params = dict(dimensions=2, optimizer="adam", max_iter=50, seed=42)
        obs1 = mmvec(X, Y, batch_norm="legacy", **params)
        obs2 = mmvec(X, Y, batch_norm="unbiased", **params)

        # Results should differ due to different normalization
        with self.assertRaises(AssertionError):
            pdt.assert_frame_equal(obs1.ranks, obs2.ranks, rtol=1e-10)


class TestMMvecEstimator(unittest.TestCase):
    """Test MMvec estimator class."""

    def setUp(self):
        """Create test data."""
        self.X, self.Y, *_ = random_multimodal(8, 10, 50, seed=42)

    def test_predict(self):
        """Test predict method returns valid Y distributions."""
        model = MMvec(n_components=2, max_iter=50, seed=42).fit(self.X, self.Y)

        # Predict on new samples
        X_test = self.X.iloc[:5]
        Y_pred = model.predict(X_test)

        # Predictions should have correct shape
        self.assertEqual(Y_pred.shape, (5, 10))

        # Predictions should sum to 1 per row
        npt.assert_allclose(Y_pred.sum(axis=1), 1.0, rtol=1e-6)

        # Predictions should be positive
        self.assertTrue((Y_pred.to_numpy() >= 0).all())

        # Column names should match Y features
        self.assertListEqual(list(Y_pred.columns), list(self.Y.columns))

    def test_predict_match(self):
        """Check if results match the old ranks path."""
        model = MMvec(n_components=2, max_iter=50, seed=42).fit(self.X, self.Y)
        X_test = self.X.iloc[:5].to_numpy()
        obs = model._predict(X_test)
        X_props = X_test / X_test.sum(axis=1, keepdims=True)
        probs = softmax(model.ranks_.to_numpy(), validate=False)
        exp = X_props @ probs
        npt.assert_allclose(obs, exp)

    def test_score(self):
        """Test score method returns valid Q-squared value."""
        # Split data into train and test sets
        X_train, X_test = self.X.iloc[:40], self.X.iloc[40:]
        Y_train, Y_test = self.Y.iloc[:40], self.Y.iloc[40:]
        model = MMvec(n_components=2, max_iter=100, seed=42).fit(X_train, Y_train)
        q2 = model.score(X_test, Y_test)
        self.assertIsInstance(q2, float)

        # Q^2 should be <= 1.0 (perfect prediction)
        self.assertLessEqual(q2, 1.0)

    def test_predict_zero_sample_raises(self):
        """Predict with zero-count sample should raise ValueError."""
        model = MMvec(
            n_components=2, max_iter=10, seed=42
        ).fit(self.X, self.Y)

        # Create X with a zero row
        X_ = self.X.iloc[:3].copy()
        X_.iloc[0, :] = 0
        with self.assertRaises(ValueError) as ctx:
            model.predict(X_)
        self.assertIn("all-zero counts", str(ctx.exception))

    def test_mmvec_result_str(self):
        """Test string representation of MMvecResult."""
        result = mmvec(self.X, self.Y, dimensions=2, max_iter=10, seed=42)
        obs = str(result)
        self.assertIn("MMvecResult", obs)
        self.assertIn("Features", obs)
        self.assertIn("Targets", obs)
        self.assertIn("Components", obs)
        self.assertIn("Iterations", obs)

    def test_predict_numpy_array(self):
        """Test predict with numpy array input (not DataFrame)."""
        model = MMvec(n_components=2, max_iter=50, seed=42).fit(self.X, self.Y)
        predictions = model.predict(self.X.to_numpy()[:5])
        self.assertEqual(predictions.shape[0], 5)
        npt.assert_allclose(predictions.sum(axis=1), 1.0, rtol=1e-6)

    def test_score_numpy_array(self):
        """Test score with numpy array inputs (not DataFrames)."""
        X_train, X_test = self.X.iloc[:40], self.X.iloc[40:]
        Y_train, Y_test = self.Y.iloc[:40], self.Y.iloc[40:]
        model = MMvec(n_components=2, max_iter=100, seed=42).fit(X_train, Y_train)
        q2 = model.score(X_test.to_numpy(), Y_test.to_numpy())
        self.assertIsInstance(q2, float)
        self.assertLessEqual(q2, 1.0)

    def test_score_zero_raises(self):
        """Score with zero-count Y sample should raise ValueError."""
        model = MMvec(n_components=2, max_iter=100, seed=42).fit(self.X, self.Y)
        Y_ = self.Y.iloc[:5].copy()
        Y_.iloc[0, :] = 0
        with self.assertRaises(ValueError) as ctx:
            model.score(self.X.iloc[:5], Y_)
        self.assertIn("all-zero counts", str(ctx.exception))

    def test_check_is_fitted(self):
        """Should raise if fit has not been called."""
        model = MMvec(n_components=2, max_iter=10, seed=42)
        with self.assertRaises(ValueError) as ctx:
            model._check_is_fitted()
        self.assertIn("not fitted", str(ctx.exception))

    def test_str(self):
        """String representation of estimator should require fitted state."""
        model = MMvec(n_components=2, max_iter=10, seed=42)
        model.fit(self.X, self.Y)
        obs = str(model)
        self.assertIn("MMvec", obs)
        self.assertIn("Features", obs)
        self.assertIn("Targets", obs)
        self.assertIn("Components", obs)
        self.assertIn("Iterations", obs)

    def test_get_params(self):
        """Estimator parameters should be returned as a dict."""
        model = MMvec(
            n_components=4,
            optimizer="adam",
            max_iter=123,
            learning_rate=0.02,
            batch_size=11,
            x_prior_mean=1.5,
            x_prior_scale=0.8,
            y_prior_mean=-2.5,
            y_prior_scale=1.2,
            beta_1=0.8,
            beta_2=0.85,
            clipnorm=3.0,
            batch_norm="legacy",
            seed=7,
            verbose=True,
            output_format="numpy",
        )

        obs = model.get_params()
        exp = {
            "n_components": 4,
            "optimizer": "adam",
            "max_iter": 123,
            "learning_rate": 0.02,
            "batch_size": 11,
            "x_prior_mean": 1.5,
            "x_prior_scale": 0.8,
            "y_prior_mean": -2.5,
            "y_prior_scale": 1.2,
            "beta_1": 0.8,
            "beta_2": 0.85,
            "clipnorm": 3.0,
            "batch_norm": "legacy",
            "seed": 7,
            "verbose": True,
            "output_format": "numpy",
        }
        self.assertDictEqual(obs, exp)

    def test_set_params(self):
        """set_params should update known parameters and reject unknown ones."""
        model = MMvec(n_components=2, max_iter=10, seed=42)

        result = model.set_params(n_components=5, optimizer="adam", max_iter=20)
        self.assertIs(result, model)
        self.assertEqual(model.n_components, 5)
        self.assertEqual(model.optimizer, "adam")
        self.assertEqual(model.max_iter, 20)

        with self.assertRaises(ValueError) as ctx:
            model.set_params(not_a_real_param=1)
        self.assertIn("Invalid parameter", str(ctx.exception))


class TestMMvecValidation(unittest.TestCase):
    """Test input validation for MMvec."""

    def setUp(self):
        """Create valid test data."""
        self.X, self.Y, *_ = random_multimodal(8, 10, 50, seed=42)

    def test_mismatched_sample_counts(self):
        """Mismatched sample counts should raise."""
        X = self.X.iloc[:40]  # 40 samples
        Y = self.Y.iloc[:50]  # 50 samples
        with self.assertRaises(ValueError) as ctx:
            mmvec(X, Y, max_iter=1)
        self.assertIn("same number of samples", str(ctx.exception))

    def test_all_zero_columns(self):
        """Columns (features) with all zero counts should raise."""
        X = self.X.copy()
        for i in (1, 2, 3):
            X.iloc[:, i] = 0
        with self.assertRaises(ValueError) as ctx:
            mmvec(X, self.Y, max_iter=1)
        self.assertIn("X contains all-zero columns", str(ctx.exception))

        Y = self.Y.copy()
        for i in (4, 5, 6):
            Y.iloc[:, i] = 0
        with self.assertRaises(ValueError) as ctx:
            mmvec(self.X, Y, max_iter=1)
        self.assertIn("y contains all-zero columns", str(ctx.exception))

    def test_all_zero_rows(self):
        """Rows (samples) with all zero counts should raise."""
        X = self.X.copy()
        for i in (1, 2, 3):
            X.iloc[i, :] = 0
        with self.assertRaises(ValueError) as ctx:
            mmvec(X, self.Y, max_iter=1)
        self.assertIn("X contains all-zero rows", str(ctx.exception))

        Y = self.Y.copy()
        for i in (4, 5, 6):
            Y.iloc[i, :] = 0
        with self.assertRaises(ValueError) as ctx:
            mmvec(self.X, Y, max_iter=1)
        self.assertIn("y contains all-zero rows", str(ctx.exception))

    def test_prior_scales(self):
        """x_prior_scale and y_prior_scale must be positive."""
        msg = "x_prior_scale must be positive"
        with self.assertRaises(ValueError) as ctx:
            mmvec(self.X, self.Y, x_prior_scale=0, max_iter=1)
        self.assertIn(msg, str(ctx.exception))
        with self.assertRaises(ValueError) as ctx:
            mmvec(self.X, self.Y, x_prior_scale=-0.5, max_iter=1)
        self.assertIn(msg, str(ctx.exception))

        msg = "y_prior_scale must be positive"
        with self.assertRaises(ValueError) as ctx:
            mmvec(self.X, self.Y, y_prior_scale=0, max_iter=1)
        self.assertIn(msg, str(ctx.exception))
        with self.assertRaises(ValueError) as ctx:
            mmvec(self.X, self.Y, y_prior_scale=-0.5, max_iter=1)
        self.assertIn(msg, str(ctx.exception))


class TestMMvecInputTypes(unittest.TestCase):
    """Test different input types for MMvec."""

    def test_sparse_pandas_input(self):
        """Test with sparse pandas DataFrame."""
        X_dense, Y_dense, *_ = random_multimodal(8, 10, 50, seed=42)

        # Convert to sparse pandas DataFrame
        X_sparse = X_dense.astype(pd.SparseDtype("float64", 0))
        Y_sparse = Y_dense.astype(pd.SparseDtype("float64", 0))

        result = mmvec(X_sparse, Y_sparse, dimensions=2, max_iter=10, seed=42)

        # Should produce valid results
        self.assertEqual(result.ranks.shape, (8, 10))
        npt.assert_allclose(
            result.ranks.to_numpy().mean(axis=1), 0, atol=1e-10
        )


class TestMMvecParameterBehavior(unittest.TestCase):
    """Test parameter behavior for MMvec."""

    def setUp(self):
        """Create test data."""
        self.X, self.Y, *_ = random_multimodal(8, 10, 50, seed=42)

    def test_nonzero_prior_mean_affects_results(self):
        """Non-zero prior means should affect results."""
        params = dict(dimensions=2, max_iter=50, seed=42)
        result_default = mmvec(
            self.X, self.Y, **params, x_prior_mean=0.0, y_prior_mean=0.0
        )
        result_nonzero = mmvec(
            self.X, self.Y, **params, x_prior_mean=5.0, y_prior_mean=5.0
        )

        # Results should differ
        self.assertFalse(
            np.allclose(
                result_default.x_embeddings.to_numpy(),
                result_nonzero.x_embeddings.to_numpy(),
            )
        )

    def test_verbose_output(self):
        """Verbose mode should produce output."""
        # Capture stdout
        captured = io.StringIO()
        sys.stdout = captured
        try:
            mmvec(
                self.X,
                self.Y,
                dimensions=2,
                optimizer="adam",
                max_iter=5,
                verbose=True,
                seed=42,
            )
        finally:
            sys.stdout = sys.__stdout__

        output = captured.getvalue()
        self.assertIn("Epoch", output)
        self.assertIn("Loss", output)


class TestMMvecResultVerification(unittest.TestCase):
    """Test output structure and properties."""

    def setUp(self):
        """Create test data."""
        self.X, self.Y, *_ = random_multimodal(8, 10, 50, seed=42)

    def test_index_columns_preserved(self):
        """Index and column names should be preserved from input."""
        result = mmvec(self.X, self.Y, dimensions=2, max_iter=10, seed=42)
        # X IDs preserved
        pdt.assert_index_equal(result.x_embeddings.index, self.X.columns)
        pdt.assert_index_equal(result.ranks.index, self.X.columns)
        # Y IDs preserved
        pdt.assert_index_equal(result.y_embeddings.index, self.Y.columns)
        pdt.assert_index_equal(result.ranks.columns, self.Y.columns)

    def test_convergence_structure(self):
        """Convergence DataFrame should have correct structure."""
        result = mmvec(self.X, self.Y, dimensions=2, max_iter=10, seed=42)
        self.assertTupleEqual(result.convergence.shape, (11,))

    def test_loss_decreases(self):
        """Loss should generally decrease during training."""
        result = mmvec(
            self.X, self.Y, dimensions=2, optimizer="adam", max_iter=100, seed=42
        )
        losses = result.convergence.to_numpy()
        # Compare first 10% average to last 10% average
        n = len(losses)
        early_loss = losses[: n // 10].mean()
        late_loss = losses[-n // 10 :].mean()
        self.assertLess(late_loss, early_loss)


class TestMMvecLBFGS(unittest.TestCase):
    """Test L-BFGS optimizer for MMvec."""

    def setUp(self):
        """Create test data."""
        res = random_multimodal(8, 10, 50, 2, seed=42)
        self.X, self.Y = res[0], res[1]
        self.x_embed_true, self.y_embed_true = res[4], res[6]

    def test_lbfgs_produces_valid_results(self):
        """L-BFGS optimizer should produce valid results."""
        result = mmvec(self.X, self.Y, 2, optimizer="lbfgs", max_iter=100, seed=42)

        # Check result type
        self.assertIsInstance(result, MMvecResult)

        # Check shapes
        self.assertEqual(result.x_embeddings.shape, (8, 3))
        self.assertEqual(result.y_embeddings.shape, (10, 3))
        self.assertEqual(result.ranks.shape, (8, 10))

        # Ranks should be row-centered
        row_means = result.ranks.to_numpy().mean(axis=1)
        npt.assert_allclose(row_means, 0, atol=1e-10)

    def test_lbfgs_reproducibility(self):
        """L-BFGS should be deterministic with same seed."""
        args = dict(dimensions=2, optimizer="lbfgs", max_iter=50, seed=123)
        res1 = mmvec(self.X, self.Y, **args)
        res2 = mmvec(self.X, self.Y, **args)
        pdt.assert_frame_equal(res1.ranks, res2.ranks, rtol=1e-10)

    def test_lbfgs_recovers_structure(self):
        """L-BFGS should recover embedding structure from synthetic data."""
        # Use more data for better recovery
        res = random_multimodal(
            8, 8, 150, 2, x_noise_std=2, x_total=1000, y_total=10000, seed=1
        )
        X, Y = res[0], res[1]
        x_embed_true = res[4]
        result = mmvec(X, Y, dimensions=2, optimizer="lbfgs", max_iter=500, seed=42)
        # Check correlation of pairwise distances
        x_r, _ = spearmanr(
            pdist(result.x_embeddings.to_numpy()[:, :-1]), pdist(x_embed_true)
        )
        self.assertGreater(x_r, 0.3, f"X embedding correlation too low: {x_r}")

    def test_lbfgs_score_on_test_data(self):
        """L-BFGS model should produce reasonable Q-squared score on test data."""
        # Split data
        X_train = self.X.iloc[:40]
        X_test = self.X.iloc[40:]
        Y_train = self.Y.iloc[:40]
        Y_test = self.Y.iloc[40:]

        model = MMvec(
            n_components=2, optimizer="lbfgs", max_iter=100, seed=42
        ).fit(X_train, Y_train)

        # Compute Q^2 score on test data
        q2 = model.score(X_test, Y_test)

        # Q^2 should be a valid float
        self.assertIsInstance(q2, float)

    def test_lbfgs_verbose_output(self):
        """L-BFGS verbose mode should produce output."""
        captured = io.StringIO()
        sys.stdout = captured
        try:
            mmvec(
                self.X,
                self.Y,
                dimensions=2,
                optimizer="lbfgs",
                max_iter=50,
                verbose=True,
                seed=42,
            )
        finally:
            sys.stdout = sys.__stdout__

        output = captured.getvalue()
        self.assertIn("Iteration", output)
        self.assertIn("Loss", output)

    def test_seed_generator(self):
        """seed as np.random.Generator should work."""
        rng = np.random.default_rng(42)
        result = mmvec(
            self.X,
            self.Y,
            dimensions=2,
            optimizer="lbfgs",
            max_iter=50,
            seed=rng,
        )
        self.assertIsInstance(result, MMvecResult)
        self.assertEqual(result.ranks.shape, (8, 10))

    def test_numpy_array_inputs(self):
        """mmvec should accept numpy arrays (not just DataFrames)."""
        result = mmvec(
            self.X.to_numpy(),
            self.Y.to_numpy(),
            dimensions=2,
            optimizer="lbfgs",
            max_iter=50,
            seed=42,
        )
        self.assertIsInstance(result, MMvecResult)
        self.assertEqual(result.ranks.shape, (8, 10))

    def test_invalid_optimizer_raises(self):
        """Invalid optimizer should raise ValueError."""
        with self.assertRaises(ValueError) as ctx:
            mmvec(
                self.X,
                self.Y,
                optimizer="invalid",
                max_iter=1,
            )

        self.assertIn("Optimizer must be", str(ctx.exception))


class TestMMvecCaseStudies(unittest.TestCase):
    """Test MMvec on real-world case study datasets.

    These tests verify that the MMvec implementation can reproduce the biological
    findings from the original MMvec publication and example notebooks.

    """

    def test_biocrust_wetting(self):
        """Co-occurrence of Cyanobacteria and known metabolites in wetup biocrust.

        This test reproduces the finding from the soils example notebook: the model
        should learn that Cyanobacteria (specifically rplo 1) co-occur with a known
        set of 13 metabolites associated with Microcoleus vaginatus, a desert crust
        cyanobacterium.

        """
        # Load the biocrust wetting dataset
        subdir = os.path.join("data", "soils")
        microbes = pd.read_table(
            get_data_path("microbes.tsv", subdir), index_col=0
        ).T
        metabolites = pd.read_table(
            get_data_path("metabolites.tsv", subdir), index_col=0
        ).T

        # Fit the model with low latent dimension for faster testing
        result = MMvec(
            n_components=1,
            optimizer="lbfgs",
            max_iter=500,
            x_prior_scale=1.0,
            y_prior_scale=1.0,
            seed=42,
        ).fit(microbes, metabolites)

        # Define expected Cyanobacteria-associated metabolites
        expected_metabolites = [
            "(3-methyladenine)",
            "7-methyladenine",
            "4-guanidinobutanoate",
            "uracil",
            "xanthine",
            "hypoxanthine",
            "(N6-acetyl-lysine)",
            "cytosine",
            "N-acetylornithine",
            "succinate",
            "adenosine",
            "guanine",
            "adenine",
        ]

        # Define Cyanobacteria taxon
        cyanobacteria_id = "rplo 1 (Cyanobacteria)"

        # Count how many of the expected metabolites have positive rank for
        # Cyanobacteria.
        ranks_cyanobacteria = result.ranks_.loc[cyanobacteria_id]
        positive_count = (ranks_cyanobacteria.loc[expected_metabolites] > 0).sum()

        # The original notebook expects all 13 metabolites to have positive ranks for
        # Cyanobacteria. We use a more lenient threshold (11) to account for
        # optimization variability.
        # self.assertEqual(positive_count, 13)
        self.assertGreaterEqual(positive_count, 11)

        # Check numerical values against expected outputs for reproducibility.
        # NOTE: Due to optimization variability and floating-point accumulation order
        # changes from computational optimizations, these values may not match exactly,
        # but should be close.
        # obs = result.score(microbes, metabolites)
        # exp = 0.217611  # (Linux)
        # self.assertAlmostEqual(obs, exp, places=4)

    @unittest.skip("Skipping a test that requires long runtime.")
    def test_cystic_fibrosis(self):
        """Co-occurrence of Pseudomonas and rhamnolipids in cystic fibrosis sputum.
        
        This test reproduces the finding from the example notebook: the model should
        learn that Pseudomonas microbes co-occur with rhamnolipids and other
        Pseudomonas-associated metabolites.

        Note: The actual data files are not kept in the codebase due to their sizes. To
        run this test, place the local data files in the expected directory. It is
        anticipated that the test will cost several to dozens of seconds.

        """
        # Load the cystic fibrosis dataset
        subdir = os.path.join("data", "cf")
        microbes = pd.read_table(
            get_data_path("microbes.tsv", subdir), index_col=0
        ).T
        metabolites = pd.read_table(
            get_data_path("metabolites.tsv", subdir), index_col=0
        ).T
        microbe_meta = pd.read_table(
            get_data_path("microbe_meta.tsv", subdir), index_col=0
        )
        metabolite_meta = pd.read_table(
            get_data_path("metabolite_meta.tsv", subdir), index_col=0
        )

        # Find Pseudomonas microbes in the metadata (n=39)
        pseudomonas = microbe_meta[
            microbe_meta["Taxon"].str.contains("Pseudomonas")
        ].index.tolist()

        # Find expert-annotated metabolites in the metadata (n=20)
        expert_metabolites = metabolite_meta[
            metabolite_meta["expert_annotation"].notnull()
        ].index.tolist()

        # Fit the model with parameters similar to the original
        result = mmvec(
            microbes,
            metabolites,
            dimensions=3,
            optimizer="lbfgs",
            max_iter=500,
            x_prior_scale=1.0,
            y_prior_scale=1.0,
            seed=42,
        )

        # For the first Pseudomonas taxon in the ranks table, count expert-annotated
        # metabolites with positive ranks.
        ranks_pseudomonas = result.ranks[result.ranks.index.isin(pseudomonas)]
        first_pseudomonas = ranks_pseudomonas.iloc[0]

        # The original analysis expects 19 out of 20 metabolites with positive ranks.
        # A slightly lower threshold (15) is used here for reproducibility.
        positive_count = (first_pseudomonas.loc[expert_metabolites] > 0).sum()
        self.assertEqual(positive_count, 19)
        self.assertGreaterEqual(positive_count, 15)


if __name__ == "__main__":
    unittest.main()
