from pyoptools.all import *
from numpy import pi

L1=SphericalLens(radius=25,curvature_s1=1./100.,curvature_s2=-1./100,thickness=10,material=material.schott["N-BK7"])
S1=CCD()
S=System(complist=[(L1,(0,0,200),(0,0,0)),
                   (S1,(0,0,400),(0,0,0))],n=1)
R=[Ray(pos=(0,0,0),dir=(0,.1,1),wavelength=.650),
   Ray(pos=(0,0,0),dir=(0,-.1,1),wavelength=.650),
   Ray(pos=(0,0,0),dir=(.1,0,1),wavelength=.650),
   Ray(pos=(0,0,0),dir=(-.1,0,1),wavelength=.650),
   Ray(pos=(0,0,0),dir=(0,0,1),wavelength=.650)]
S.ray_add(R)
S.propagate()
Plot3D(S,center=(0,0,200),size=(400,100),scale=2,rot=[(0,-pi/2,0),(pi/20,-pi/10,0)])

"""
Explanation: Following a ray through a system
Sometimes it is necessary to follow a ray that has been ray-traced through a system. In the following example 
To demonstrate this, a simple system will be simulated.
End of explanation
"""


r=R[0]
r.childs

"""
Explanation: Each ray has an attribute childs, this is a tuple that contains all the rays that continue after the intersection with a surface (this is filled after a propagation). When this tuple contains 0 elements, it means that this is the final ray of the trace or what is the same, this ray does not intersect any element of the system. If this tuple has one or more elements, it means that the ray intersected an optical surface. If the surface is partially reflective, 2 rays (transmitted and reflected) will be created. If the surface is a diffraction gratting, the tuple can contain as much elements as diffraction orders simulated.
The next example shows the content of childs for one of the rays simulated in the system above.
End of explanation
"""


r=R[0]
while True:
    print(r,"\n")
    if len(r.childs)==0: 
        break
    r=r.childs[0]
    

"""
Explanation: Next there is a simple example show all the rays that make ray-path of one of the rays in the simulation above:
End of explanation
"""


import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')


for r in R:
    X=[]
    Y=[]
    Z=[]
    while True:
        x,y,z=r.pos
        X.append(x)
        Y.append(y)
        Z.append(z)
    
        if len(r.childs)==0: 
            break
        r=r.childs[0]
    ax.plot(X,Y,Z,zdir="y")
ax.set(xlim=(-100, 100), ylim=(0, 400), zlim=(-100, 100))


"""
Explanation: In this example we are assuming all the rays have only one child. Care must be taken when this is not the case.
In the next example, the position information of each ray is obtained, and a plot is made:
End of explanation
"""


