import numpy as np
import pytest
from giant.data.transforms import (
    inv_log_transform,
    local_frame_rotation,
    log_transform,
    Normalizer,
)


def test_log_transform_invertible():
    x = np.array([0.1, 1.0, 10.0, 1000.0], dtype=np.float32)
    np.testing.assert_allclose(inv_log_transform(log_transform(x)), x, rtol=1e-5)


def test_local_frame_rotation_noop_when_aligned():
    N = 8
    pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
    rng = np.random.default_rng(0)
    post_dir = rng.standard_normal((N, 3)).astype(np.float32)
    post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)
    result = local_frame_rotation(pre_dir, post_dir)
    np.testing.assert_allclose(result, post_dir, atol=1e-5)


def test_local_frame_rotation_preserves_angle():
    """Angle between pre_dir and post_dir must equal angle between ẑ and rotated."""
    rng = np.random.default_rng(1)
    N = 200
    pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
    pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
    post_dir = rng.standard_normal((N, 3)).astype(np.float32)
    post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)

    rotated = local_frame_rotation(pre_dir, post_dir)

    cos_before = (pre_dir * post_dir).sum(axis=1)
    cos_after = rotated[:, 2]  # dot with ẑ = z-component (unit vectors)
    np.testing.assert_allclose(cos_after, cos_before, atol=1e-5)


def test_local_frame_rotation_preserves_norm():
    rng = np.random.default_rng(2)
    N = 100
    pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
    pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
    post_dir = rng.standard_normal((N, 3)).astype(np.float32)
    post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)
    result = local_frame_rotation(pre_dir, post_dir)
    np.testing.assert_allclose(np.linalg.norm(result, axis=1), 1.0, atol=1e-5)


def test_normalizer_roundtrip():
    rng = np.random.default_rng(3)
    X = rng.standard_normal((200, 9)).astype(np.float32)
    norm = Normalizer().fit(X)
    np.testing.assert_allclose(norm.inverse_transform(norm.transform(X)), X, atol=1e-5)


def test_normalizer_serialization():
    rng = np.random.default_rng(4)
    X = rng.standard_normal((50, 6)).astype(np.float32)
    norm = Normalizer().fit(X)
    norm2 = Normalizer.from_dict(norm.to_dict())
    np.testing.assert_allclose(norm2.mean, norm.mean)
    np.testing.assert_allclose(norm2.std, norm.std)