import numpy as np
import math
import timeit

# denominators for sine and cosine degree Taylor polynomials
# sin(x) ≈ x - x^3/6 + x^5/120, cos(x) ≈ 1 - x^2/2 + x^4/24
SIN_P, COS_P = np.array([1, -1/6]), np.array([1, -1/2])
 
def _approx_trig_gen_powers(x, degree):
    n = x.shape[0]
    out = np.empty((degree + 1, n), dtype=x.dtype)
    out[0] = 1
    if degree >= 1:
        out[1] = x
        for i in range(2, degree + 1):
            np.multiply(out[i - 1], x, out=out[i])
    return out

def method_approximate_trig(n_samples):
    raw = np.random.randint(0, 1 << 32, size=n_samples, dtype=np.uint32)
    angles = ((raw >> 2) / (1 << 30)) * (math.pi / 2)
    w1 = (2*(raw & 1)-1).astype(np.int8)
    w2 = (2*((raw >> 1) & 1) -1).astype(np.int8)
    pow = _approx_trig_gen_powers(angles, 3)
    pow_e, pow_o = pow[0::2], pow[1::2]
    return w1 * (COS_P @ pow_e), w2 * (SIN_P @ pow_o)
 
def _ngon_chord_gen_anchors(n):
    theta = np.linspace(0, 2 * np.pi, n, endpoint=False)
    return np.stack([np.cos(theta), np.sin(theta)], axis=1)

def method_select_ngon_chord(n_samples, anchors, logn):
    u = np.random.randint(0, 2**32, size=n_samples, dtype=np.uint32)
    ll = len(anchors)
    t_bits = 32 - logn
    t = (u & ((1 << t_bits) - 1)) / float(1 << t_bits)
    mask = (1 << logn) - 1
    a = (u >> (32 - logn)) & mask
    return (1 - t)[:, None] * anchors[a] + t[:, None] * anchors[(a+1)%ll]

def method_standard(n_samples):
    theta = np.random.uniform(0, 2 * np.pi, n_samples)
    return np.stack([np.cos(theta), np.sin(theta)], axis=1)

def method_normalize(n_samples):
    xy = np.random.normal(size=(n_samples, 2))
    return xy / np.linalg.norm(xy, axis=1, keepdims=True)

n_samples = 10_000_000
anchors, logn = _ngon_chord_gen_anchors(64), int(np.log2(64))
chord_time = timeit.timeit(lambda: method_select_ngon_chord(n_samples, anchors, logn), number=5)
std_time = timeit.timeit(lambda: method_standard(n_samples), number=5)
norm_time = timeit.timeit(lambda: method_normalize(n_samples), number=5)
approx_trig_time = timeit.timeit(lambda: method_approximate_trig(n_samples), number=5)