from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
from sys import argv

R = 5
p = 2
fmax = (R + p)/(4*(np.pi**2)*R)

def f(u, v):
    return (R + p*np.cos(v))/(4*(np.pi**2)*R)

def points(f, fmax, n=500):
    points = []
    R = np.random.random
    while True:
        u = R()*2*np.pi
        v = R()*2*np.pi

        if fmax*R() < f(u, v): points.append((u, v))
        if len(points) == n: break

    return points


def torus(points):
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    plt.hold(True)

    u = np.linspace(0, 2*np.pi, 100)
    v = np.linspace(0, 2*np.pi, 100)

    x = np.outer((R + p*np.cos(v)), np.cos(u))
    y = np.outer((R + p*np.cos(v)), np.sin(u))
    z = np.outer(p*np.sin(v), np.ones(np.size(u)))

    pu = np.array([p[0] for p in points])
    pv = np.array([p[1] for p in points])

    px = (R + p*np.cos(pv))*np.cos(pu)
    py = (R + p*np.cos(pv))*np.sin(pu)
    pz = p*np.sin(pv)

    ax.plot_surface(x, y, z,  rstride=10, cstride=10, color='w', alpha=.5)
    ax.scatter(px, py, pz, color = 'r', s=10, alpha=1)

    plt.show()

if __name__ == "__main__":
    if len(argv)>1:
        n = int(argv[1])
        points = points(f, fmax, n)
    else:
        points = points(f, fmax)
    torus(points)