import numpy as np
import pandas as pd
import pytest
from numpy.testing import assert_array_almost_equal, assert_array_equal
from sklearn.utils.estimator_checks import parametrize_with_checks

from ktch.datasets import load_landmark_mosquito_wings
from ktch.landmark import (
    GeneralizedProcrustesAnalysis,
    centroid_size,
    combine_landmarks_and_curves,
    define_curve_sliders,
)

data_landmark_mosquito_wings = load_landmark_mosquito_wings(as_frame=True)
data_landmark_mosquito_wings.coords

X = data_landmark_mosquito_wings.coords.to_numpy().reshape(-1, 18 * 2)


def test_gpa_shape():
    gpa = GeneralizedProcrustesAnalysis()
    gpa.fit_transform(X)
    X_transformed = gpa.fit_transform(X)

    assert X.shape == X_transformed.shape


@pytest.mark.parametrize("n_dim", [2, 3])
def test_centroid_size(n_dim):
    x = np.random.uniform(0, 100, (10, n_dim))
    cs_r = np.sqrt(np.sum((x - x.mean(axis=0)) ** 2))
    cs_t = centroid_size(x)

    assert_array_almost_equal(cs_r, cs_t)


###########################################################
#
#   define_curve_sliders
#
###########################################################


def test_define_curve_sliders_open():
    """Test define_curve_sliders with open curve."""
    result = define_curve_sliders([0, 1, 2, 3, 4])
    expected = np.array([[0, 1, 2], [1, 2, 3], [2, 3, 4]])
    assert_array_equal(result, expected)


def test_define_curve_sliders_closed_curve_pattern():
    """Test define_curve_sliders for closed curve using repeated anchor index."""
    result = define_curve_sliders([0, 1, 2, 3, 0])
    expected = np.array([[0, 1, 2], [1, 2, 3], [2, 3, 0]])
    assert_array_equal(result, expected)


def test_define_curve_sliders_minimum():
    """Test define_curve_sliders with minimum 3 points."""
    result = define_curve_sliders([5, 6, 7])
    expected = np.array([[5, 6, 7]])
    assert_array_equal(result, expected)


def test_define_curve_sliders_too_few_points():
    """Test define_curve_sliders raises error with less than 3 points."""
    with pytest.raises(ValueError, match="At least 3 points"):
        define_curve_sliders([0, 1])


def test_define_curve_sliders_combine_multiple():
    """Test combining multiple curves."""
    curve1 = define_curve_sliders([0, 1, 2, 3])
    curve2 = define_curve_sliders([3, 4, 5, 6])
    combined = np.vstack([curve1, curve2])
    expected = np.array([[0, 1, 2], [1, 2, 3], [3, 4, 5], [4, 5, 6]])
    assert_array_equal(combined, expected)


###########################################################
#
#   combine_landmarks_and_curves
#
###########################################################


def test_combine_landmarks_and_curves_basic():
    """Test basic combine_landmarks_and_curves functionality."""
    landmarks = np.array(
        [
            [[0, 0], [1, 0], [1, 1]],
            [[0, 0], [1, 0], [1, 1]],
        ]
    )
    curves = [
        [np.array([[0.2, 0.2], [0.4, 0.4], [0.6, 0.6], [0.8, 0.8]])],
        [np.array([[0.2, 0.2], [0.4, 0.4], [0.6, 0.6], [0.8, 0.8]])],
    ]

    combined, slider_matrix, info = combine_landmarks_and_curves(
        landmarks, curves, curve_landmarks=[(0, 2)]
    )

    assert combined.shape == (2, 7, 2)  # 3 landmarks + 4 curve points
    assert slider_matrix.shape == (4, 3)  # all 4 curve points slide
    assert info["n_landmarks"] == 3
    assert info["n_curve_points"] == 4


def test_combine_landmarks_and_curves_slider_indices():
    """Test that slider indices are correct with anchor flags."""
    landmarks = np.array(
        [
            [[0, 0], [5, 0]],
        ]
    )
    curves = [
        [np.array([[1, 1], [2, 2], [3, 3], [4, 4]])],
    ]

    combined, slider_matrix, info = combine_landmarks_and_curves(
        landmarks, curves, anchor_first=True, anchor_last=True
    )

    # With anchor_first=True, anchor_last=True:
    # Points 0,1 are landmarks (indices 0,1)
    # Curve points are at indices 2,3,4,5
    # First (index 2) and last (index 5) are anchors
    # Sliders are indices 3,4

    assert slider_matrix.shape[0] == 2  # 2 sliders
    assert np.all(slider_matrix[:, 1] == [3, 4])  # slider indices


def test_combine_landmarks_and_curves_with_curve_landmarks():
    """Test that curve_landmarks produces correct slider topology."""
    landmarks = np.array(
        [
            [[0, 0], [5, 0], [5, 5]],
        ]
    )
    curves = [
        [np.array([[1, 1], [2, 2], [3, 3]])],
    ]

    combined, slider_matrix, info = combine_landmarks_and_curves(
        landmarks, curves, curve_landmarks=[(0, 2)]
    )

    # 3 landmarks (0,1,2) + 3 curve points (3,4,5)
    # All curve points slide with LM0 and LM2 as anchors
    assert slider_matrix.shape == (3, 3)
    # First semilandmark: before=LM0, slider=3, after=4
    assert np.all(slider_matrix[0] == [0, 3, 4])
    # Middle: before=3, slider=4, after=5
    assert np.all(slider_matrix[1] == [3, 4, 5])
    # Last semilandmark: before=4, slider=5, after=LM2
    assert np.all(slider_matrix[2] == [4, 5, 2])


###########################################################
#
#   GPA with semilandmarks
#
###########################################################


@pytest.fixture
def semilandmark_test_data():
    """Create test data for GPA with semilandmarks."""
    np.random.seed(42)
    base = np.array(
        [
            [0.0, 0.0],
            [1.0, 0.5],
            [2.0, 1.0],
            [3.0, 0.5],
            [4.0, 0.0],
        ]
    )
    X = np.array(
        [
            base + np.random.randn(5, 2) * 0.1,
            base + np.random.randn(5, 2) * 0.1,
            base + np.random.randn(5, 2) * 0.1,
        ]
    )
    X_flat = X.reshape(3, -1)
    curves = define_curve_sliders([0, 1, 2, 3, 4])
    return X_flat, curves


def test_gpa_with_semilandmarks_procrustes(semilandmark_test_data):
    """Test GPA with semilandmarks using Procrustes distance criterion."""
    X_flat, curves = semilandmark_test_data

    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=curves,
        sliding_criterion="procrustes_distance",
    )
    X_aligned = gpa.fit_transform(X_flat)

    assert X_aligned.shape == X_flat.shape
    assert gpa.mu_ is not None
    assert gpa.mu_.shape == (5, 2)


def test_gpa_with_semilandmarks_bending_energy(semilandmark_test_data):
    """Test GPA with semilandmarks using bending energy criterion."""
    X_flat, curves = semilandmark_test_data

    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=curves,
        sliding_criterion="bending_energy",
    )
    X_aligned = gpa.fit_transform(X_flat)

    assert X_aligned.shape == X_flat.shape
    assert gpa.mu_ is not None


def test_gpa_backward_compatibility(semilandmark_test_data):
    """Test that GPA without curves produces same results as before."""
    X_flat, _ = semilandmark_test_data

    # Without curves parameter
    gpa1 = GeneralizedProcrustesAnalysis(n_dim=2)
    X1 = gpa1.fit_transform(X_flat)

    # With curves=None explicitly
    gpa2 = GeneralizedProcrustesAnalysis(n_dim=2, curves=None)
    X2 = gpa2.fit_transform(X_flat)

    assert_array_almost_equal(X1, X2)


###########################################################
#
#   parallel n_jobs
#
###########################################################


@pytest.mark.parametrize("n_jobs", [None, 1, 3])
def test_gpa_semilandmark_parallel_bending(n_jobs, semilandmark_test_data):
    """Test that parallel semilandmark GPA (bending energy) produces identical results."""
    X_flat, curves = semilandmark_test_data

    gpa_baseline = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=curves,
        sliding_criterion="bending_energy",
        n_jobs=None,
    )
    X_baseline = gpa_baseline.fit_transform(X_flat)

    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=curves,
        sliding_criterion="bending_energy",
        n_jobs=n_jobs,
    )
    X_result = gpa.fit_transform(X_flat)

    assert_array_almost_equal(X_result, X_baseline)
    assert_array_almost_equal(gpa.mu_, gpa_baseline.mu_)


@pytest.mark.parametrize("n_jobs", [None, 1, 3])
def test_gpa_semilandmark_parallel_procrustes(n_jobs, semilandmark_test_data):
    """Test that parallel semilandmark GPA (Procrustes distance) produces identical results."""
    X_flat, curves = semilandmark_test_data

    gpa_baseline = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=curves,
        sliding_criterion="procrustes_distance",
        n_jobs=None,
    )
    X_baseline = gpa_baseline.fit_transform(X_flat)

    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=curves,
        sliding_criterion="procrustes_distance",
        n_jobs=n_jobs,
    )
    X_result = gpa.fit_transform(X_flat)

    assert_array_almost_equal(X_result, X_baseline)
    assert_array_almost_equal(gpa.mu_, gpa_baseline.mu_)


@pytest.mark.parametrize("n_jobs", [None, 1, 3])
def test_gpa_shape_parallel(n_jobs):
    """Test that parallel shape alignment produces identical results."""
    gpa_baseline = GeneralizedProcrustesAnalysis(n_dim=2, n_jobs=None)
    X_baseline = gpa_baseline.fit_transform(X)

    gpa = GeneralizedProcrustesAnalysis(n_dim=2, n_jobs=n_jobs)
    X_result = gpa.fit_transform(X)

    assert_array_almost_equal(X_result, X_baseline)
    assert_array_almost_equal(gpa.mu_, gpa_baseline.mu_)


###########################################################
#
#   size-and-shape mode
#
###########################################################


def test_gpa_size_and_shape():
    """Test GPA in size-and-shape mode (scaling=False)."""
    gpa = GeneralizedProcrustesAnalysis(n_dim=2, scaling=False)
    X_transformed = gpa.fit_transform(X)

    assert X.shape == X_transformed.shape
    assert gpa.mu_ is not None


@pytest.mark.parametrize("n_jobs", [None, 1, 3])
def test_gpa_size_and_shape_parallel(n_jobs):
    """Test that parallel size-and-shape alignment produces identical results."""
    gpa_baseline = GeneralizedProcrustesAnalysis(n_dim=2, scaling=False, n_jobs=None)
    X_baseline = gpa_baseline.fit_transform(X)

    gpa = GeneralizedProcrustesAnalysis(n_dim=2, scaling=False, n_jobs=n_jobs)
    X_result = gpa.fit_transform(X)

    assert_array_almost_equal(X_result, X_baseline)
    assert_array_almost_equal(gpa.mu_, gpa_baseline.mu_)


###########################################################
#
#   sklearn API compliance
#
###########################################################


def test_gpa_unfitted_transform_raises():
    """Test that transform on an unfitted GPA raises NotFittedError."""
    from sklearn.exceptions import NotFittedError

    gpa = GeneralizedProcrustesAnalysis(n_dim=2)
    with pytest.raises(NotFittedError):
        gpa.transform(X)


def test_gpa_fit_then_transform():
    """Test that fit followed by transform produces same results as fit_transform."""
    gpa1 = GeneralizedProcrustesAnalysis(n_dim=2)
    X_ft = gpa1.fit_transform(X)

    gpa2 = GeneralizedProcrustesAnalysis(n_dim=2)
    gpa2.fit(X)
    X_t = gpa2.transform(X)

    assert_array_almost_equal(X_ft, X_t)
    assert_array_almost_equal(gpa1.mu_, gpa2.mu_)


def test_gpa_transform_new_data():
    """Test that transform aligns unseen data to the learned mean."""
    np.random.seed(42)
    n_landmarks, n_dim = 5, 2
    base = np.random.randn(n_landmarks, n_dim)

    # Training data: rotated/scaled copies of base
    X_train = []
    for _ in range(5):
        angle = np.random.uniform(0, 2 * np.pi)
        R = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]])
        scale = np.random.uniform(0.5, 2.0)
        offset = np.random.randn(n_dim)
        X_train.append((scale * (base @ R.T) + offset).ravel())
    X_train = np.array(X_train)

    gpa = GeneralizedProcrustesAnalysis(n_dim=2)
    gpa.fit(X_train)

    # New data: another rotated/scaled copy
    angle = np.random.uniform(0, 2 * np.pi)
    R = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]])
    X_new = (1.5 * (base @ R.T) + np.array([3.0, -2.0])).ravel().reshape(1, -1)

    X_aligned = gpa.transform(X_new)

    assert X_aligned.shape == X_new.shape
    # Aligned shape should be close to the mean shape
    aligned_2d = X_aligned.reshape(n_landmarks, n_dim)
    # Procrustes aligns to mu_, so disparity should be small
    _, _, disp = __import__("scipy").spatial.procrustes(gpa.mu_, aligned_2d)
    assert disp < 1e-10


def test_gpa_fit_accepts_y():
    """Test that fit and fit_transform accept y parameter (ignored)."""
    gpa = GeneralizedProcrustesAnalysis(n_dim=2)
    y_dummy = np.ones(len(X))

    X1 = gpa.fit(X, y=y_dummy).transform(X)
    X2 = gpa.fit_transform(X, y=y_dummy)

    assert X1.shape == X.shape
    assert X2.shape == X.shape


def test_gpa_validate_data_rejects_nan():
    """Test that NaN in input raises ValueError via validate_data."""
    X_bad = X.copy().astype(float)
    X_bad[0, 0] = np.nan
    gpa = GeneralizedProcrustesAnalysis(n_dim=2)
    with pytest.raises(ValueError, match="Input X contains NaN"):
        gpa.fit(X_bad)


def test_gpa_validate_data_rejects_inf():
    """Test that inf in input raises ValueError via validate_data."""
    X_bad = X.copy().astype(float)
    X_bad[0, 0] = np.inf
    gpa = GeneralizedProcrustesAnalysis(n_dim=2)
    with pytest.raises(ValueError, match="Input X contains infinity"):
        gpa.fit(X_bad)


def test_gpa_invalid_n_dim_raises_on_fit():
    """Test that invalid n_dim raises ValueError on fit, not on construction."""
    gpa = GeneralizedProcrustesAnalysis(n_dim=4)
    # Construction should succeed
    assert gpa.n_dim == 4
    # Fit should fail
    with pytest.raises(ValueError, match="n_dim must be 2 or 3"):
        gpa.fit(X)


def test_gpa_clone():
    """Test that sklearn clone works without side effects."""
    from sklearn.base import clone

    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        tol=1e-5,
        scaling=False,
        debug=True,
    )
    gpa_cloned = clone(gpa)

    assert gpa_cloned.n_dim == 2
    assert gpa_cloned.tol == 1e-5
    assert gpa_cloned.scaling is False
    assert gpa_cloned.debug is True
    assert not hasattr(gpa_cloned, "mu_")


def test_gpa_debug_warning_not_in_init():
    """Test that debug warning fires on fit, not on construction."""
    import warnings as _warnings

    # Construction should NOT warn
    with _warnings.catch_warnings(record=True) as record:
        _warnings.simplefilter("always")
        GeneralizedProcrustesAnalysis(n_dim=2, debug=True)
    future_warnings = [w for w in record if issubclass(w.category, FutureWarning)]
    assert len(future_warnings) == 0

    # Fit SHOULD warn
    gpa = GeneralizedProcrustesAnalysis(n_dim=2, debug=True)
    with pytest.warns(FutureWarning, match="debug"):
        gpa.fit(X)


def test_gpa_n_features_in_set_after_fit():
    """Test that n_features_in_ is set after fitting."""
    gpa = GeneralizedProcrustesAnalysis(n_dim=2)
    assert not hasattr(gpa, "n_features_in_")

    gpa.fit(X)
    assert gpa.n_features_in_ == X.shape[1]


###########################################################
#
#   parameter validation
#
###########################################################


def test_gpa_invalid_sliding_criterion():
    """Test that invalid sliding_criterion raises error."""
    X = np.random.randn(3, 10)
    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=np.array([[0, 1, 2]]),
        sliding_criterion="invalid",
    )
    with pytest.raises(ValueError, match="sliding_criterion must be"):
        gpa.fit(X)


def test_gpa_invalid_curves_shape():
    """Test that invalid curves shape raises error."""
    X = np.random.randn(3, 10)
    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=np.array([0, 1, 2]),  # 1D instead of 2D
    )
    with pytest.raises(ValueError, match="curves must be a 2D array"):
        gpa.fit(X)


def test_gpa_curves_index_out_of_range():
    """Test that out-of-range curve indices raise error."""
    X = np.random.randn(3, 10)  # 5 landmarks * 2 dims
    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=np.array([[0, 1, 100]]),  # index 100 is out of range
    )
    with pytest.raises(ValueError, match="curves indices must be in range"):
        gpa.fit(X)


def test_gpa_surfaces_not_implemented():
    """Test that surfaces parameter raises NotImplementedError."""
    X = np.random.randn(3, 10)
    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        surfaces=np.array([0, 1, 2]),
    )
    with pytest.raises(NotImplementedError, match="Surface semilandmarks"):
        gpa.fit(X)


def test_gpa_centroid_size_zero_raises():
    """Test that specimens with centroid size 0 raise ValueError."""
    X_degenerate = np.array(
        [
            [1.0, 2.0, 1.0, 2.0],  # 2 landmarks at the same point
            [3.0, 4.0, 5.0, 6.0],
        ]
    )
    gpa = GeneralizedProcrustesAnalysis(n_dim=2)
    with pytest.raises(ValueError, match="centroid size 0"):
        gpa.fit(X_degenerate)


###########################################################
#
#   closed-form bending energy sliding
#
###########################################################


def test_bending_energy_sliding_reduces_energy():
    """Test that sliding reduces bending energy."""
    from ktch.landmark._kernels import tps_bending_energy

    np.random.seed(42)
    base = np.array(
        [
            [0.0, 0.0],
            [1.0, 0.5],
            [2.0, 1.0],
            [3.0, 0.5],
            [4.0, 0.0],
        ]
    )
    # Perturb the semilandmarks
    specimen = base.copy()
    specimen[1] += np.array([0.0, 0.2])
    specimen[2] += np.array([0.0, -0.3])
    specimen[3] += np.array([0.0, 0.15])

    curves = define_curve_sliders([0, 1, 2, 3, 4])
    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=curves,
        sliding_criterion="bending_energy",
    )

    # Compute bending energy before sliding
    be_before = tps_bending_energy(base, specimen)

    # Slide
    slid = gpa._slide_bending_energy(specimen, base, curves)

    # Compute bending energy after sliding
    be_after = tps_bending_energy(base, slid)

    assert be_after < be_before


def test_bending_energy_sliding_only_moves_sliders():
    """Test that sliding only moves slider points, not anchors."""
    np.random.seed(42)
    base = np.array(
        [
            [0.0, 0.0],
            [1.0, 0.5],
            [2.0, 1.0],
            [3.0, 0.5],
            [4.0, 0.0],
        ]
    )
    specimen = base + np.random.randn(5, 2) * 0.1

    curves = define_curve_sliders([0, 1, 2, 3, 4])
    # Sliders are indices 1, 2, 3; anchors are 0, 4
    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=curves,
        sliding_criterion="bending_energy",
    )

    slid = gpa._slide_bending_energy(specimen, base, curves)

    # Anchor points (0 and 4) should not move
    assert_array_almost_equal(slid[0], specimen[0])
    assert_array_almost_equal(slid[4], specimen[4])

    # At least some slider points should have moved
    assert not np.allclose(slid[1:4], specimen[1:4])


###########################################################
#
#   re-projection onto curves
#
###########################################################


def test_reproject_onto_curves_basic():
    """Test that re-projection places points on the line segment."""
    base = np.array(
        [
            [0.0, 0.0],
            [1.0, 0.0],
            [2.0, 0.0],
            [3.0, 0.0],
            [4.0, 0.0],
        ]
    )
    # Slide point 2 off-curve (perpendicular displacement)
    slid = base.copy()
    slid[2] = [2.0, 0.5]

    curves = define_curve_sliders([0, 1, 2, 3, 4])
    gpa = GeneralizedProcrustesAnalysis(n_dim=2, curves=curves)
    result = gpa._reproject_onto_curves(slid, base, curves)

    # Point 2 should be projected back onto the segment [1,0]->[3,0] at (2,0)
    assert_array_almost_equal(result[2], [2.0, 0.0])
    # Anchors should be unchanged
    assert_array_almost_equal(result[0], base[0])
    assert_array_almost_equal(result[4], base[4])


def test_reproject_onto_curves_preserves_anchors():
    """Test that re-projection does not move anchor points."""
    base = np.array(
        [
            [0.0, 0.0],
            [1.0, 1.0],
            [2.0, 0.0],
        ]
    )
    slid = base.copy()
    slid[1] = [1.5, 1.5]  # move slider off-curve

    curves = define_curve_sliders([0, 1, 2])
    gpa = GeneralizedProcrustesAnalysis(n_dim=2, curves=curves)
    result = gpa._reproject_onto_curves(slid, base, curves)

    assert_array_almost_equal(result[0], base[0])
    assert_array_almost_equal(result[2], base[2])


def test_reproject_onto_curves_3d():
    """Test re-projection in 3D space."""
    base = np.array(
        [
            [0.0, 0.0, 0.0],
            [1.0, 1.0, 1.0],
            [2.0, 2.0, 2.0],
        ]
    )
    slid = base.copy()
    slid[1] = [1.0, 1.5, 0.5]  # move off the line

    curves = define_curve_sliders([0, 1, 2])
    gpa = GeneralizedProcrustesAnalysis(n_dim=3, curves=curves)
    result = gpa._reproject_onto_curves(slid, base, curves)

    # The segment goes from (0,0,0) to (2,2,2).
    # Projection of (1,1.5,0.5) onto this line: t = dot(P-A,AB)/dot(AB,AB)
    # AB = (2,2,2), P-A = (1,1.5,0.5), dot = 2+3+1 = 6, |AB|^2 = 12
    # t = 0.5, projected = (1,1,1)
    assert_array_almost_equal(result[1], [1.0, 1.0, 1.0])


def test_reproject_clamps_to_segment():
    """Test that re-projection clamps to segment endpoints."""
    base = np.array(
        [
            [0.0, 0.0],
            [1.0, 0.0],
            [2.0, 0.0],
        ]
    )
    slid = base.copy()
    slid[1] = [-1.0, 0.0]  # beyond point A

    curves = define_curve_sliders([0, 1, 2])
    gpa = GeneralizedProcrustesAnalysis(n_dim=2, curves=curves)
    result = gpa._reproject_onto_curves(slid, base, curves)

    # Should clamp to segment start (point A = [0,0])
    assert_array_almost_equal(result[1], [0.0, 0.0])


def test_reproject_two_segments_selects_nearest():
    """Test that re-projection picks the nearest of the two segments."""
    # L-shaped curve: B=(0,0), S_orig=(1,0), A=(1,1)
    # Segment B->S_orig is horizontal, segment S_orig->A is vertical
    base = np.array(
        [
            [0.0, 0.0],  # B (anchor)
            [1.0, 0.0],  # S_orig (slider)
            [1.0, 1.0],  # A (anchor)
        ]
    )
    curves = define_curve_sliders([0, 1, 2])
    gpa = GeneralizedProcrustesAnalysis(n_dim=2, curves=curves)

    # Point near the horizontal segment B->S_orig
    slid_h = base.copy()
    slid_h[1] = [0.5, 0.1]
    result_h = gpa._reproject_onto_curves(slid_h, base, curves)
    assert_array_almost_equal(result_h[1], [0.5, 0.0])

    # Point near the vertical segment S_orig->A
    slid_v = base.copy()
    slid_v[1] = [1.1, 0.5]
    result_v = gpa._reproject_onto_curves(slid_v, base, curves)
    assert_array_almost_equal(result_v[1], [1.0, 0.5])


def test_reproject_two_segments_3d():
    """Test 2-segment re-projection in 3D with non-collinear curve."""
    # V-shaped curve in 3D: B=(0,0,0), S_orig=(1,0,0), A=(1,1,0)
    base = np.array(
        [
            [0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0],
            [1.0, 1.0, 0.0],
        ]
    )
    curves = define_curve_sliders([0, 1, 2])
    gpa = GeneralizedProcrustesAnalysis(n_dim=3, curves=curves)

    # Displace slider above the plane, closer to segment S_orig->A
    slid = base.copy()
    slid[1] = [1.05, 0.5, 0.3]
    result = gpa._reproject_onto_curves(slid, base, curves)

    # Projection onto S_orig->A: S_orig=(1,0,0), A=(1,1,0), SA=(0,1,0)
    # P-S=(0.05,0.5,0.3), t = 0.5, proj = (1,0.5,0)
    # Projection onto B->S_orig: B=(0,0,0), S_orig=(1,0,0), BS=(1,0,0)
    # P-B=(1.05,0.5,0.3), t = 1.05 clamped to 1.0, proj = (1,0,0)
    # dist to S_orig->A proj: (0.05)^2 + 0 + 0.09 = 0.0925
    # dist to B->S_orig proj: (0.05)^2 + 0.25 + 0.09 = 0.3425
    # Nearest is S_orig->A projection
    assert_array_almost_equal(result[1], [1.0, 0.5, 0.0])


def test_sliding_with_reprojection_reduces_energy():
    """Test that sliding with re-projection still reduces bending energy."""
    from ktch.landmark._kernels import tps_bending_energy

    np.random.seed(42)
    base = np.array(
        [
            [0.0, 0.0],
            [1.0, 0.5],
            [2.0, 1.0],
            [3.0, 0.5],
            [4.0, 0.0],
        ]
    )
    specimen = base.copy()
    specimen[1] += np.array([0.0, 0.2])
    specimen[2] += np.array([0.0, -0.3])
    specimen[3] += np.array([0.0, 0.15])

    curves = define_curve_sliders([0, 1, 2, 3, 4])
    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=curves,
        sliding_criterion="bending_energy",
    )

    be_before = tps_bending_energy(base, specimen)

    # Slide + reproject (single specimen)
    x_slid = gpa._slide_bending_energy(specimen, base, curves)
    x_slid = gpa._reproject_onto_curves(x_slid, specimen, curves)

    be_after = tps_bending_energy(base, x_slid)
    assert be_after < be_before


###########################################################
#
#   3D curve semilandmarks
#
###########################################################


@pytest.fixture
def semilandmark_3d_test_data():
    """Create 3D test data for GPA with curve semilandmarks."""
    np.random.seed(123)
    base = np.array(
        [
            [0.0, 0.0, 0.0],
            [1.0, 0.5, 0.2],
            [2.0, 1.0, 0.4],
            [3.0, 0.5, 0.2],
            [4.0, 0.0, 0.0],
        ]
    )
    X = np.array(
        [
            base + np.random.randn(5, 3) * 0.05,
            base + np.random.randn(5, 3) * 0.05,
            base + np.random.randn(5, 3) * 0.05,
        ]
    )
    X_flat = X.reshape(3, -1)
    curves = define_curve_sliders([0, 1, 2, 3, 4])
    return X_flat, curves


def test_gpa_with_3d_curve_semilandmarks_bending_energy(semilandmark_3d_test_data):
    """Test GPA with 3D curve semilandmarks using bending energy criterion."""
    X_flat, curves = semilandmark_3d_test_data

    gpa = GeneralizedProcrustesAnalysis(
        n_dim=3,
        curves=curves,
        sliding_criterion="bending_energy",
    )
    X_aligned = gpa.fit_transform(X_flat)

    assert X_aligned.shape == X_flat.shape
    assert gpa.mu_ is not None
    assert gpa.mu_.shape == (5, 3)


def test_gpa_with_3d_curve_semilandmarks_procrustes(semilandmark_3d_test_data):
    """Test GPA with 3D curve semilandmarks using Procrustes distance criterion."""
    X_flat, curves = semilandmark_3d_test_data

    gpa = GeneralizedProcrustesAnalysis(
        n_dim=3,
        curves=curves,
        sliding_criterion="procrustes_distance",
    )
    X_aligned = gpa.fit_transform(X_flat)

    assert X_aligned.shape == X_flat.shape
    assert gpa.mu_ is not None
    assert gpa.mu_.shape == (5, 3)


def test_3d_tangent_computation():
    """Test that tangent vectors in 3D are correctly computed."""
    X = np.array(
        [
            [0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0],
            [2.0, 0.0, 0.0],
            [3.0, 0.0, 0.0],
            [4.0, 0.0, 0.0],
        ]
    )
    curves = define_curve_sliders([0, 1, 2, 3, 4])
    gpa = GeneralizedProcrustesAnalysis(n_dim=3, curves=curves)

    tangents = gpa._compute_tangents(X, curves)

    # All tangents should be unit vectors along x-axis
    for t in tangents:
        assert_array_almost_equal(t, [1.0, 0.0, 0.0])

    # Non-axis-aligned case
    X2 = np.array(
        [
            [0.0, 0.0, 0.0],
            [1.0, 1.0, 1.0],
            [2.0, 2.0, 2.0],
        ]
    )
    curves2 = define_curve_sliders([0, 1, 2])
    tangents2 = gpa._compute_tangents(X2, curves2)
    expected = np.array([1.0, 1.0, 1.0]) / np.sqrt(3)
    assert_array_almost_equal(tangents2[0], expected)


def test_bending_energy_sliding_moves_along_tangent():
    """Test that sliding displacement is parallel to the tangent direction."""
    base = np.array(
        [
            [0.0, 0.0],
            [1.0, 0.0],
            [2.0, 0.0],
            [3.0, 0.0],
            [4.0, 0.0],
        ]
    )
    specimen = base.copy()
    specimen[2] += np.array([0.0, 0.5])  # Move middle point off the line

    curves = define_curve_sliders([0, 1, 2, 3, 4])
    gpa = GeneralizedProcrustesAnalysis(
        n_dim=2,
        curves=curves,
        sliding_criterion="bending_energy",
    )

    # Compute tangents from specimen (same as _slide_bending_energy does)
    tangents = gpa._compute_tangents(specimen, curves)

    slid = gpa._slide_bending_energy(specimen, base, curves)

    slider_idx = curves[:, 1]
    for j, s_idx in enumerate(slider_idx):
        displacement = slid[s_idx] - specimen[s_idx]
        if np.linalg.norm(displacement) < 1e-14:
            continue
        # Displacement should be parallel to tangent:
        # cross product (2D) should be zero
        cross = displacement[0] * tangents[j][1] - displacement[1] * tangents[j][0]
        assert abs(cross) < 1e-10


###########################################################
#
#   sklearn estimator checks (parametrize_with_checks)
#
###########################################################

# GPA requires n_features % n_dim == 0 and non-degenerate specimens
# (centroid_size > 0). Many sklearn checks use test data that violates
# these domain constraints, so we xfail those checks.
_N_DIM_REASON = "sklearn test data has n_features not divisible by n_dim=2"
_SINGLE_LM_REASON = (
    "sklearn test data maps to 1 landmark (2 features / n_dim=2), "
    "giving centroid_size=0"
)

_EXPECTED_FAILURES = {
    # sklearn test data has n_features not divisible by n_dim=2
    "check_fit_score_takes_y": _N_DIM_REASON,
    "check_dont_overwrite_parameters": _N_DIM_REASON,
    "check_estimators_dtypes": _N_DIM_REASON,
    "check_pipeline_consistency": _N_DIM_REASON,
    "check_estimators_nan_inf": _N_DIM_REASON,
    "check_estimators_pickle": _N_DIM_REASON,
    "check_f_contiguous_array_estimator": _N_DIM_REASON,
    "check_transformer_data_not_an_array": _N_DIM_REASON,
    "check_transformer_general": _N_DIM_REASON,
    "check_transformer_preserve_dtypes": _N_DIM_REASON,
    "check_methods_sample_order_invariance": _N_DIM_REASON,
    "check_methods_subset_invariance": _N_DIM_REASON,
    "check_dict_unchanged": _N_DIM_REASON,
    "check_fit2d_predict1d": _N_DIM_REASON,
    "check_fit2d_1feature": _N_DIM_REASON,
    # sklearn test data maps to 1 landmark per specimen (centroid_size=0)
    "check_estimators_overwrite_params": _SINGLE_LM_REASON,
    "check_estimators_fit_returns_self": _SINGLE_LM_REASON,
    "check_readonly_memmap_input": _SINGLE_LM_REASON,
    "check_fit_idempotent": _SINGLE_LM_REASON,
    "check_fit_check_is_fitted": _SINGLE_LM_REASON,
    "check_n_features_in": _SINGLE_LM_REASON,
}


def _expected_failed_checks(estimator):
    return _EXPECTED_FAILURES


@parametrize_with_checks(
    [GeneralizedProcrustesAnalysis(n_dim=2)],
    expected_failed_checks=_expected_failed_checks,
)
def test_sklearn_estimator_checks(estimator, check):
    check(estimator)


###########################################################
#
#   set_output(transform="pandas")
#
###########################################################


def test_gpa_set_output_pandas():
    """Test that set_output(transform='pandas') returns a DataFrame."""
    gpa = GeneralizedProcrustesAnalysis(n_dim=2)
    gpa.set_output(transform="pandas")
    X_aligned = gpa.fit_transform(X)

    assert isinstance(X_aligned, pd.DataFrame)
    assert X_aligned.shape == X.shape


def test_gpa_set_output_pandas_preserves_columns():
    """Test that DataFrame output preserves input column names."""
    cols = [f"lm{i}_{ax}" for i in range(18) for ax in ("x", "y")]
    X_df = pd.DataFrame(X, columns=cols)

    gpa = GeneralizedProcrustesAnalysis(n_dim=2)
    gpa.set_output(transform="pandas")
    X_aligned = gpa.fit_transform(X_df)

    assert isinstance(X_aligned, pd.DataFrame)
    assert list(X_aligned.columns) == cols


def test_gpa_set_output_pandas_transform():
    """Test set_output with separate fit and transform calls."""
    gpa = GeneralizedProcrustesAnalysis(n_dim=2)
    gpa.set_output(transform="pandas")
    gpa.fit(X)
    X_aligned = gpa.transform(X)

    assert isinstance(X_aligned, pd.DataFrame)
    assert X_aligned.shape == X.shape


def test_gpa_transform_rejects_wrong_n_features():
    """Test that transform rejects input with wrong number of features."""
    gpa = GeneralizedProcrustesAnalysis(n_dim=2)
    gpa.fit(X)

    X_wrong = np.random.randn(5, X.shape[1] + 2)
    with pytest.raises(ValueError):
        gpa.transform(X_wrong)
