import numpy as np
from skimage import transform


class ScaleTranslationTransform(transform.SimilarityTransform):
    """
    A custom transformation model that only considers uniform scaling and translation.

    This class is a simplified version of `skimage.transform.SimilarityTransform`
    that ignores rotation and shear. It is used to estimate the transformation
    from a set of source and destination points.
    """

    def estimate(self, src, dst):
        # src and dst are (N, 2) arrays
        if src.shape[0] < 2:
            return False

        # 1. Compute centroids
        src_mean = src.mean(axis=0)
        dst_mean = dst.mean(axis=0)

        # 2. Center the points
        src_centered = src - src_mean
        dst_centered = dst - dst_mean

        # 3. Compute scale s
        # s = sum(src_centered * dst_centered) / sum(src_centered^2)
        # We sum over both dimensions and all points
        num = np.sum(src_centered * dst_centered)
        den = np.sum(src_centered**2)

        if den == 0:
            return False

        s = num / den

        # 4. Compute translation
        # t = dst_mean - s * src_mean
        t = dst_mean - s * src_mean

        # 5. Set parameters
        # Update the matrix
        self.params = np.array([[s, 0, t[0]], [0, s, t[1]], [0, 0, 1]])

        return True


def test():
    # Create synthetic data
    # src: square
    src = np.array([[0, 0], [10, 0], [10, 10], [0, 10]])

    # Case 1: Pure translation
    dst_trans = src + [5, 5]
    model = ScaleTranslationTransform()
    model.estimate(src, dst_trans)
    print("Case 1 (Translation):")
    print(f"Scale: {model.scale} (Expected 1.0)")
    print(f"Translation: {model.translation} (Expected [5. 5.])")
    print(f"Rotation: {model.rotation} (Expected 0)")

    # Case 2: Pure scale
    dst_scale = src * 2.0
    model = ScaleTranslationTransform()
    model.estimate(src, dst_scale)
    print("\nCase 2 (Scale):")
    print(f"Scale: {model.scale} (Expected 2.0)")
    print(f"Translation: {model.translation} (Expected [0. 0.])")

    # Case 3: Scale + Translation
    dst_st = src * 0.5 + [10, 20]
    model = ScaleTranslationTransform()
    model.estimate(src, dst_st)
    print("\nCase 3 (Scale + Translation):")
    print(f"Scale: {model.scale} (Expected 0.5)")
    print(f"Translation: {model.translation} (Expected [10. 20.])")

    # Case 4: Rotation (should fail to capture rotation, but find best scale/trans)
    # Rotate 90 degrees: (x,y) -> (-y, x)
    # src: (0,0), (10,0), (10,10), (0,10)
    # dst: (0,0), (0,10), (-10,10), (-10,0)
    # Centroids: src(5, 5), dst(-5, 5)
    # src_c: (-5,-5), (5,-5), (5,5), (-5,5)
    # dst_c: (5,-5), (5,5), (-5,5), (-5,-5)
    # num = (-25 + 25) + (25 - 25) + (-25 + 25) + (25 - 25) = 0?
    # Let's check.
    # (-5)*5 + (-5)*(-5) = -25 + 25 = 0
    # 5*5 + (-5)*5 = 25 - 25 = 0
    # ...
    # So scale should be 0? That makes sense as projection of orthogonal vectors.

    dst_rot = np.array([[0, 0], [0, 10], [-10, 10], [-10, 0]])
    model = ScaleTranslationTransform()
    model.estimate(src, dst_rot)
    print("\nCase 4 (Rotation 90deg):")
    print(f"Scale: {model.scale}")
    print(f"Translation: {model.translation}")

    # Case 5: Rotation small (10 deg)
    # Should find scale ~1 and translation approx correct
    theta = np.radians(10)
    c, s = np.cos(theta), np.sin(theta)
    rot_mat = np.array([[c, -s], [s, c]])
    dst_rot_small = src @ rot_mat.T

    model = ScaleTranslationTransform()
    model.estimate(src, dst_rot_small)
    print("\nCase 5 (Rotation 10deg):")
    print(f"Scale: {model.scale}")
    print(f"Translation: {model.translation}")

    # Compare with SimilarityTransform
    sim_model = transform.SimilarityTransform()
    sim_model.estimate(src, dst_rot_small)
    print(f"Similarity Scale: {sim_model.scale}")
    print(f"Similarity Rotation: {sim_model.rotation}")


if __name__ == "__main__":
    test()