MOON/compute.py

from collections.abc import Iterable

Number = float | int


def _is_square_matrix(matrix: Iterable[Iterable[Number]]) -> bool:
    """Return True if the supplied iterable is a square matrix.

    A proper matrix must contain at least one row and every row must be
    the same length, and the number of rows must equal that length.
    """
    rows = list(matrix)
    if not rows:
        return False
    row_len = len(rows[0])
    return all(len(row) == row_len for row in rows) and len(rows) == row_len


def _iter_to_list(matrix: Iterable[Iterable[Number]]) -> list[list[Number]]:
    """Convert the supplied iterable of iterables into a list of lists."""
    return [list(row) for row in matrix]


def transpose(matrix: Iterable[Iterable[Number]]) -> list[list[Number]]:
    """Return the transpose of a 2D matrix."""
    rows = _iter_to_list(matrix)
    if not rows:
        return []

    row_len = len(rows[0])
    if not all(len(r) == row_len for r in rows):
        raise ValueError("All rows must have the same length.")

    return [list(col) for col in zip(*rows)]


def rotate(
    matrix: Iterable[Iterable[Number]], clockwise: bool = True
) -> list[list[Number]]:
    """Rotate a 2D matrix by 90 degrees.

    Parameters
    ----------
    matrix : Iterable[Iterable[Number]]
        Input matrix.
    clockwise : bool, default=True
        Rotate 90 degrees clockwise if True, counter-clockwise if False.

    Returns
    -------
    list[list[Number]]
        The rotated matrix.
    """
    rows = _iter_to_list(matrix)
    if not rows:
        return []

    row_len = len(rows[0])
    if not all(len(r) == row_len for r in rows):
        raise ValueError("All rows must have the same length.")

    if clockwise:
        # 90° im Uhrzeigersinn: Zeilen umkehren, dann transponieren
        return [list(col) for col in zip(*rows[::-1])]
    else:
        # 90° gegen den Uhrzeigersinn: transponieren, dann Zeilen umkehren
        return [list(col) for col in zip(*rows)][::-1]


def matmul(a: Iterable[Iterable[Number]], b: Iterable[Iterable[Number]]) -> list[list[Number]]:
    """Multiply two matrices."""
    a_rows = _iter_to_list(a)
    b_rows = _iter_to_list(b)

    # Validate dimensions
    if not a_rows or not b_rows:
        raise ValueError("Input matrices cannot be empty.")

    a_cols = len(a_rows[0])
    b_rows_count = len(b_rows)
    if a_cols != b_rows_count:
        raise ValueError(
            f"Incompatible matrix dimensions: {len(a_rows)}x{a_cols} cannot be "
            f"multiplied with {b_rows_count}x{len(b_rows[0])}."
        )

    b_cols = len(b_rows[0])

    if not all(len(row) == a_cols for row in a_rows):
        raise ValueError("All rows of matrix a must have the same length.")
    if not all(len(row) == b_cols for row in b_rows):
        raise ValueError("All rows of matrix b must have the same length.")

    # Skalarprodukt
    result: list[list[Number]] = []
    for i in range(len(a_rows)):
        result_row: list[Number] = []
        for j in range(b_cols):
            dot = sum(a_rows[i][k] * b_rows[k][j] for k in range(a_cols))
            result_row.append(dot)
        result.append(result_row)

    return result


# --- New functionality: rotating a 2D vector (e.g. moon around earth) ---

import math


def rotate_vector(vec: list[Number], angle_rad: Number) -> list[Number]:
    """Return a new 2‑D vector rotated by *angle_rad*.
    ``vec`` must be a list/tuple of two numbers ``[x, y]``.
    The rotation is counter‑clockwise and follows the standard
    rotation matrix:
    ``[x', y'] = [x cosθ - y sinθ, x sinθ + y cosθ]``.
    """
    if len(vec) != 2:
        raise ValueError("rotate_vector expects a 2‑D vector.")
    x, y = vec
    cos_a = math.cos(angle_rad)
    sin_a = math.sin(angle_rad)
    return [x * cos_a - y * sin_a, x * sin_a + y * cos_a]


def orbit_earth_moon(
    earth_pos: list[Number],
    moon_pos: list[Number],
    angular_speed_rad_per_time: Number,
    dt: Number,
) -> tuple[list[Number], list[Number]]:
    """Rotate *moon_pos* around *earth_pos*.

    Both *earth_pos* and *moon_pos* are ABSOLUTE [x, y] coordinates.
    The moon is moved by applying a rotation of ``omega * dt`` radians
    around the Earth. The Earth stays fixed, so repeatedly calling this
    function makes the moon orbit the Earth on a circle of constant radius.

    Returns the new positions ``(new_earth_pos, new_moon_pos)``.
    """
    # Translate moon to origin relative to earth
    relative = [moon_pos[0] - earth_pos[0], moon_pos[1] - earth_pos[1]]
    # Rotate the relative vector
    new_rel = rotate_vector(relative, angular_speed_rad_per_time * dt)
    # Translate back
    new_moon = [earth_pos[0] + new_rel[0], earth_pos[1] + new_rel[1]]
    return earth_pos, new_moon


def point_on_circle(
    center: list[Number], radius: Number, angle_rad: Number
) -> list[Number]:
    """Return the absolute [x, y] position on a circle.

    *center* is the absolute center, *radius* the distance from the center
    and *angle_rad* the angle in radians (0 points to the right,
    increasing angles rotate counter-clockwise).
    """
    return [
        center[0] + radius * math.cos(angle_rad),
        center[1] + radius * math.sin(angle_rad),
    ]


if __name__ == "__main__":
    # Test matmul
    m1 = [[1, 2], [3, 4]]
    m2 = [[5, 6], [7, 8]]
    res1 = matmul(m1, m2)
    assert res1 == [[19, 22], [43, 50]]

    m3 = [[1, 2, 3], [4, 5, 6]]
    m4 = [[7, 8], [9, 1], [2, 3]]
    res2 = matmul(m3, m4)
    assert res2 == [[31, 19], [85, 55]]

    # Test transposnieren
    m_square = [[1, 2], [3, 4]]
    assert transpose(m_square) == [[1, 3], [2, 4]]

    m_rect = [[1, 2, 3], [4, 5, 6]]
    assert transpose(m_rect) == [[1, 4], [2, 5], [3, 6]]

    # Test rotate mit Uhr
    assert rotate([[1, 2], [3, 4]]) == [[3, 1], [4, 2]]
    assert rotate([[1, 2, 3], [4, 5, 6]]) == [[4, 1], [5, 2], [6, 3]]

    # Testf rotate gegen Uhr
    assert rotate([[1, 2], [3, 4]], clockwise=False) == [[2, 4], [1, 3]]