from math import pi
import numpy as np
from .transform import gl_array

import OpenGL.GL as GL


def disc(r, n1, dir=1):
    """Return ``(vertices, normals, indices, mode)`` of a disc in the X-Y
    plane with center at Z=X=Y=0."""
    t1 = np.linspace(0, 2*pi, n1, dtype=np.float32)
    vertices = r * np.stack([
        np.cos(t1),
        np.sin(t1),
        0 * t1,
    ]).T
    vertices = np.vstack((gl_array([0, 0, 0]), vertices))
    normals = gl_array([0, 0, 1]) * np.ones((n1, 1), dtype=np.float32) * dir
    indices = np.arange(n1 + 1, dtype=np.uint32)
    return (vertices, normals, indices, GL.GL_TRIANGLE_FAN)


def cylinder(l, r, n1, n0=2):
    """Return ``(vertices, normals, indices, mode)`` of a cylinder with radius
    ``r`` around the Z axis, starting at ``Z=0``, ending at ``Z=l``."""
    z0 = np.zeros(n0, dtype=np.float32)[:, None]
    t0 = np.linspace(0, l,    n0, dtype=np.float32)[:, None]
    t1 = np.linspace(0, 2*pi, n1, dtype=np.float32)[None, :]
    centers = np.stack([z0, z0, t0])
    normals = np.stack([
        np.cos(t1),
        np.sin(t1),
        0 * t1,
    ]) + z0
    vertices = centers + r * normals
    return _tube_strip(vertices, normals)


def torus_arc(r0, r1, n0, n1, ang0=2*pi, ang1=2*pi):
    """
    Return (vertices, normals, draw indices, mode) of rounded pipe a.k.a. torus
    with x-y plane symmetry centered at ``(0, 0, 0)``.

    The returned vertices are generated by revolving a circular segment C1
    about the coplanar z axis, with C1's center coordinates being described by
    the circular segment C0. Both circular segments are desribed by triples:

        (r, n, ang) = (radius, num_points, angle).
    """
    # e_phi = (cos(s0), sin(s0), 0)
    # e_z   = (0, 0, 1)
    # center = r * e_phi
    # phi, z = cos(s1), sin(s1)
    # result = center + phi * e_phi + z * e_z
    z0 = np.zeros(n0, dtype=np.float32)[:, None]
    t0 = np.linspace(0, ang0, n0, dtype=np.float32)[:, None]
    t1 = np.linspace(0, ang1, n1, dtype=np.float32)[None, :]
    c0, s0 = np.cos(t0), np.sin(t0)
    c1, s1 = np.cos(t1), np.sin(t1)
    centers = np.stack([c0, s0, z0]) * r0
    normals = np.stack([
        c1 * c0,
        c1 * s0,
        s1 + z0,
    ])
    vertices = centers + r1 * normals
    return _tube_strip(vertices, normals)


def _tube_strip(vertices, normals):
    n0, n1 = vertices.shape[1:]
    indices = np.arange(n0 * n1, dtype=np.uint32).reshape((n0, n1))
    strips = np.stack((
        indices[:-1, :],
        indices[+1:, :],
    ))
    vertices = vertices.transpose((1, 2, 0)).reshape((-1, 3))
    normals = normals.transpose((1, 2, 0)).reshape((-1, 3))
    strips = strips.transpose((1, 2, 0)).reshape((-1,))
    return vertices, normals, strips, GL.GL_TRIANGLE_STRIP
