import bpy
import bmesh
import numpy as np
verts = []
edges = []
faces = []

Nr     = 10 # number of points along the radius of the lens
Ntheta = 10 # number of points around the symmetry axis
R      = 2  # radius of the lens

def make_rot_mat_Z(theta):
    return np.array([[np.cos(theta), -np.sin(theta), 0.],[np.sin(theta), np.cos(theta), 0.], [0., 0., 1.]])

def create_vertices_spun_profile(verts, theta, r, x, y, z):
    for j in range(len(theta)):
        for i in range(len(r)):
            rotmat = make_rot_mat_Z(theta[j])
            p = np.array([[x[i],y[i],z[i]]]).transpose()
            prot = rotmat@p
            verts.append([prot[0], prot[1], prot[2]])
    return verts

def create_faces_from_vertices_spun_profile(faces, offset=0):
    Nr = len(r)
    Ntheta = len(theta)
    for j in range(Ntheta):
        for i in range(Nr-1):
            jp1 = (j+1) % Ntheta    # j+1 must loop back to 0
            faces.append([offset+j*Nr+i, offset+j*Nr+i+1, offset+jp1*Nr+i+1, offset+jp1*Nr+i])
    return faces

def create_cylinder_faces_from_vertices_spun_profile(faces, Ntheta, Nr, NvertsHalfLens):
    for j in range(1,Ntheta-1):
        jp1 = (j+1) % Ntheta    # j+1 must loop back to 0
        faces.append([j*Nr-1, jp1*Nr-1, NvertsHalfLens+jp1*Nr-1, NvertsHalfLens+j*Nr-1])
    # add two remaining faces
    faces.append([NvertsHalfLens+(Ntheta-1)*Nr-1, NvertsHalfLens+(Ntheta)*Nr-1, (Ntheta)*Nr-1, (Ntheta-1)*Nr-1])
    faces.append([NvertsHalfLens+(Ntheta)*Nr-1, (Ntheta+1)*Nr-1, (NvertsHalfLens+(Ntheta+1)*Nr-1) % (2*NvertsHalfLens), (Ntheta)*Nr-1])
    return faces

def computelens_radii_from_focal_length(f, n, R1=None, R2=None):
    ''' f is the focal length in meters, n is the refraction index of lens material.
        Lens is assumed to be in air with n~=1.
        If R1 is specified, R2 is computed so that the focal length of the lens is f.
        If R2 is specified, R1 is computed so that the focal length of the lens is f.
        If neither R1 nor N2 are specified, R1 = -R2 and the focal length of the lens is f.
        Returns (R1, R2).
    '''
    # 1/f == (n - 1) (1/R1 - 1/R2)
    if R1 != None and R2 != None:# why would you call the method with both values already known ???
        return (R1,R2)
    if R1 == None and R2 == None:
        R1 = 2*f(n-1)
        R2 = -R1
    elif R1 == None:
        R1 = (f*(-1 + n)*R2)/(f*(n-1) + R2)
    else:
        R2 = (f*(-1 + n)*R1)/(f*(n-1) - R1)
    return (R1,R2)

dtheta = 2*np.pi/Ntheta
r = np.linspace(0., R, Nr) # linear distribution of points
r = np.sqrt(R)*np.sqrt(r)   # equal-area distribution of points
theta = np.arange(0., 2*np.pi, dtheta)
if 0:# test function
    x = r
    y = np.zeros(r.shape)
    z = .2*(1.-r**2)# parabola
    verts = create_vertices_spun_profile(verts, theta, r, x, y, z)
    faces = create_faces_from_vertices_spun_profile(faces, offset=0)

else:# thin spherical lens
    R1 = 6      # positive = convex
    R2 = 10     # positive = concave
    tmin = 0.1   # minimum thickness for concave-concave lenses
    
    # compute lens profiles
    x = r
    y = np.zeros(r.shape)
    
    if R1 > 0:# convex R1
        z1 = np.sqrt(R1**2 - r**2) - np.sqrt(R1**2 - R**2) + tmin/2
    else:     # concave R1
        z1 = -R1 - np.sqrt(R1**2 - r**2) + tmin/2
    
    if R2 > 0:# convex R1
        z2 = -(np.sqrt(R2**2 - r**2) - np.sqrt(R2**2 - R**2)) - tmin/2
    else:     # concave R2
        z2 = -(-R2 - np.sqrt(R2**2 - r**2)) - tmin/2
    
    # create vertices
    verts = create_vertices_spun_profile(verts, theta, r, x, y, z1)
    NvertsHalfLens = len(verts)
    verts = create_vertices_spun_profile(verts, theta, r, x, y, z2)
    #print(NvertsHalfLens, len(verts))
    # create faces
    faces = create_faces_from_vertices_spun_profile(faces, offset=0)
    faces = create_faces_from_vertices_spun_profile(faces, offset=NvertsHalfLens)
    faces = create_cylinder_faces_from_vertices_spun_profile(faces, Ntheta, Nr, NvertsHalfLens)# BUGGY
    #print('faces =', faces)

#print(Nr, len(r))
#print(Ntheta, len(theta))
    
# create object from vertex data
name = "Spherical lens"
mesh = bpy.data.meshes.new(name)
obj  = bpy.data.objects.new(name, mesh)
col  = bpy.data.collections.get("Collection")
col.objects.link(obj)
bpy.context.view_layer.objects.active = obj
mesh.from_pydata(verts, edges, faces)

if 1:
    # remove doubles    
    bm = bmesh.new()     # create an empty BMesh
    bm.from_mesh(mesh)   # fill it in from a Mesh
    
    bmesh.ops.remove_doubles(bm, verts=bm.verts, dist=0.001)
    
    # Finish up, write the bmesh back to the mesh
    bm.to_mesh(mesh)
    bm.free()  # free and prevent further access
    
    mesh.validate()    
    mesh.update()

# cleanup mesh
if 0:
    bpy.ops.object.shade_smooth()
    bpy.ops.mesh.select_all(action='SELECT')
    bpy.ops.mesh.remove_doubles()
    bpy.ops.mesh.normals_make_consistent(inside=False)