import numpy as np
from numpy.typing import NDArray
from time import perf_counter_ns


def _proj_original(b: NDArray, a: NDArray) -> NDArray:
    return (b.T @ (a / np.linalg.norm(a))) * (a / np.linalg.norm(a))

def _proj_new(b: NDArray, a: NDArray) -> NDArray:
    return (np.dot(b, a) / np.dot(a, a)) * a

def verify_projection():
    a = np.random.rand(3)
    b = np.random.rand(3)
    proj1 = _proj_original(b, a)
    proj2 = _proj_new(b, a)
    assert np.allclose(proj1, proj2)

def benchmark_projections(n_iter=100000):
    a = np.random.rand(3)
    b = np.random.rand(3)

    # Measure original implementation
    start_time = perf_counter_ns()
    for _ in range(n_iter):
        _proj_original(b, a)
    time_original = perf_counter_ns() - start_time

    # Measure new implementation
    start_time = perf_counter_ns()
    for _ in range(n_iter):
        _proj_new(b, a)
    time_new = perf_counter_ns() - start_time

    print(f"Original Implementation Time: {time_original / 1e6:.2f} ms")
    print(f"New Implementation Time: {time_new / 1e6:.2f} ms")

verify_projection()

print("Benchmarking both implementations...")
benchmark_projections()