from functools import wraps
from mpmath import mp

pi = mp.pi

def log2(x):
    return mp.log(x, 2)

def bits_per_vector(d, precision=16):
    return log2(d * precision)

def memoize(function):
    memo = {}

    @wraps(function)
    def wrapper(*args):
        try:
            return memo[args]
        except KeyError:
            rv = function(*args)
            memo[args] = rv
            return rv
    return wrapper

@memoize
def sphere(d):
    """ The volume of the d-1 sphere """
    with mp.workprec(53):
        return 2 ** (d / 2 * log2(mp.pi) + 1) / mp.gamma(d / 2)

def C(d, theta, prec=None):
    """
    The volume of a spherical cap of radius theta
    """
    prec = prec if prec else mp.prec
    with mp.workprec(prec):
        theta = mp.mpf(theta)
        d = mp.mpf(d)
        r = mp.betainc((d - 1) / 2, 1 / 2.0, x2=mp.sin(theta) ** 2, regularized=True) / 2
        return r

@memoize
def W(d, alpha, beta, theta, prec=None):
    """
    The volume of the intersection between spherical caps of radius alpha and beta
    that are at angle theta with eachother.
    """
    assert alpha <= mp.pi / 2
    assert beta <= mp.pi / 2
    assert 0 >= (mp.cos(beta) - mp.cos(alpha) * mp.cos(theta)) * (mp.cos(beta) * mp.cos(theta) - mp.cos(alpha))

    if theta >= alpha + beta:
        return mp.mpf(0.0)

    prec = prec if prec else mp.prec
    with mp.workprec(prec):
        alpha = mp.mpf(alpha)
        beta = mp.mpf(beta)
        theta = mp.mpf(theta)
        d = mp.mpf(d)
        c = mp.atan(mp.cos(alpha) / (mp.cos(beta) * mp.sin(theta)) - 1 / mp.tan(theta))

        def f_alpha(x):
            return mp.sin(x) ** (d - 2) * mp.betainc(
                (d - 2) / 2,
                1 / 2.0,
                x2=mp.sin(mp.re(mp.acos(mp.tan(theta - c) / mp.tan(x)))) ** 2,
                regularized=True,
            )

        def f_beta(x):
            return mp.sin(x) ** (d - 2) * mp.betainc(
                (d - 2) / 2, 1 / 2.0, x2=mp.sin(mp.re(mp.acos(mp.tan(c) / mp.tan(x)))) ** 2, regularized=True
            )

        S_alpha = mp.quad(f_alpha, (theta - c, alpha), error=True)[0] / 2
        S_beta = mp.quad(f_beta, (c, beta), error=True)[0] / 2

        return (S_alpha + S_beta) * sphere(d - 1) / sphere(d)

def N(d):
    return 1/C(d, pi/3)

def R(d, x, y, z, m, force_a=None):
    if x != y:
        print("x must equal y")
        return
    mem_req = log2(N(d)) + bits_per_vector(d)
    if x+y+z+m < mem_req:
        print(f"Memory error: {x+y+z+m} < {mem_req}")
        return
    if force_a is None:
        a = pi/2
        T = cost_table(d, x, y, z, m, a)
        while T["search_cost_per_iter"] > T["sort_cost_per_iter"]:
            a -= 0.001
            T = cost_table(d, x, y, z, m, a)
    else:
        a = force_a
    summarize(d, x, y, z, m, a)

def cost_table(d, x, y, z, m, a):
    out = {}
    out["machines"] = x + y + z
    # Number of vectors in the database
    out["vecs"] = log2(N(d))
    # Memory required for those vectors
    out["mem"] = out["vecs"] + bits_per_vector(d)
    # Number of buckets
    out["buckets"] = -log2(W(d, a, a, pi/3))
    # Number of vectors per bucket
    out["bucket_vecs"] = log2(N(d) * C(d, a))
    # Memory per bucket
    out["bucket_bits"] = out["bucket_vecs"] + bits_per_vector(d)
    # Number of buckets stored in one iteration
    out["buckets_per_iter"] = x + y + z + m - out["bucket_bits"]
    # Number of iterations needed to prevent bucket storage from
    # dominating memory.
    out["iters"] = out["buckets"] - out["buckets_per_iter"]
    # Cost of sorting full memory
    out["sort_cost_per_iter"] = out["machines"] + x + m
    # Cost to search one bucket
    out["search_cost_per_iter"] = out["buckets_per_iter"] + out["bucket_vecs"] + out["bucket_bits"]
    out["cost"] = out["iters"] + max(out["sort_cost_per_iter"], out["search_cost_per_iter"])
    return out

def summarize(d, x, y, z, m, a):
    print(f"Costs for sieving in dimension {d} on an R({x}, {y}, {z}, {m})")
    T = cost_table(d, x, y, z, m, a)
    print(f"a = {a}, cos(a) = {mp.cos(a)}")
    print(f"Vecs: {T['vecs']} ({T['vecs']/d} * {d})")
    print(f"Mem: {T['mem']}")
    print(f"Buckets: {T['buckets']} ({T['buckets']/d} * {d})")
    print(f"Iterations: {T['iters']} ({T['iters']/d} * {d})")
    print(f"Buckets / iter: {T['buckets_per_iter']}")
    print(f"Vecs / bucket: {T['bucket_vecs']} ({T['bucket_vecs']/d} * {d})")
    print(f"Mem / bucket: {T['bucket_bits']}")
    print(f"Sort cost / iter: {T['sort_cost_per_iter']}")
    print(f"Search cost / iter: {T['search_cost_per_iter']}")
    print(f"Cost: {T['cost']} ({T['cost']/d} * {d})")
    print()

if __name__ == "__main__":
    z = 10
    m = 50
    for d in [384, 576, 768, 960]:
        # smallest x such that N(d) vectors can be stored in x + x + z + m bits
        x = max(0,(log2(N(d)) + bits_per_vector(d) - z - m)/2);
        R(d, x, x, z, m)