# Create a (length-1) RayVector from position and velocity 3-tuples (or 3-elements lists, arrays, ...)
x, y, z = 0.1, 0.2, 0.3  # meters
vx, vy, vz = 0.1, 0.2, np.sqrt(1 - 0.1**2 - 0.2**2)  # in units of c
r = (x, y, z)
v = (vx, vy, vz)
rv = batoid.RayVector(x, y, z, vx, vy, vz)

# Printing the ray reveals its attributes.
print(rv)

# You can access the position and velocity attributes either individually (x, y, z) or as
# 3-vectors (r)
print(rv.x)
print(rv.vz)

print(rv.r)
print(rv.v)

# Rays have a number of optional fields too
print(rv.t)  # The time in meters (see explanation above)
print(rv.wavelength)  # Vacuum wavelength in meters.
print(rv.flux)  # Flux in arbitrary units.
print(rv.vignetted)  # Whether or not the ray has been vignetted
print(rv.failed)  # Whether or not the ray is failed.  
                  # Usually this means an intersection between the Ray and a Surface could not be found.

# You can create rays with specific values for these optional fields too 
rv = batoid.RayVector(*r, *v, t=0.3, wavelength=400e-9, flux=2.3, vignetted=True)
print(rv)

"""
Explanation: RayVectors
The most fundamental object in batoid is the RayVector, which you can roughly think of as and array of photons.  RayVectors are defined primarily by the position vectors r and velocity vectors v of their constituent photons.  Positions in batoid are always specified in meters, and velocities are always in units of the speed of light in vacuum.  This makes the magnitude of the velocity vector equal to $1/n$, where $n$ is the refractive index of the medium in which each ray is currently propagating.  If you're just planning on geometrically tracing rays through reflective optics, then positions and velocities are the only required RayVector parameters.  The main RayVector constructor requires x, y, and z components of position and vx, vy, vz components of velocity.
The next most important RayVector attribute after position and velocity is wavelength, which is (almost always) required for tracing through refractive optics.  (The exception is if you define a refractive medium that doesn't depend on wavelength.  Real media do have wavelength-dependent refractive indices of course.)  Wavelengths are always specified in meters, and always in vacuum -- even when the rays being created aren't currently in vacuum.  If you don't specify a value for wavelength explicitly when constructing a RayVector, a (somewhat nonsensical) value of 0.0 meters will be used.
Next in the list of RayVector attributes is time t.  Batoid doesn't track time directly, but rather tracks the product of time and the speed of light in vacuum.  That product has dimensions of length, so time in batoid is measured in meters.  Note that this convention is complementary to the convention of tracking velocity in units of the speed of light in vacuum -- we still have the familiar relation $\Delta \vec{r} = \vec{v} \Delta t$.  The default value of time t for a RayVector is 0.0 meters.
The final fundamental attributes of a RayVector are a flux for each ray in arbitrary units and status indicators for whether rays in RayVector have become vignetted during tracing, or if batoid failed to successfully trace some rays in a RayVector for some reason.  These and a number of additional derived attributes will be discussed more below.
End of explanation
"""


print(rv.omega)  # angular frequency
print(rv.k)  # wave vector
print(rv.kx)  # individual components of the wave vector are also accessible directly

"""
Explanation: Derived Ray attributes
Two additional attributes derivable from the fundamental attributes detailed above are the wavevector $\vec{k}$ and the temporal frequency $\omega$.  The wavevector is defined by
$\vec{k} = \frac{2 \pi \vec{v}}{\lambda |\vec{v}|^2}$
(with $\lambda$ being the vacuum wavelength) and has units of radians per meter.
The normal physics definition of angular temporal frequency is $\omega = \frac{2 \pi c}{\lambda}$, but in batoid, we divide out the speed of light in vacuum to get
$\omega = \frac{2 \pi}{\lambda}$
The units are again radians per meter (but in contrast to $\vec{k}$ this time in a scalar variable).  This choice for the frequency means that $\omega t$ is a phase angle in batoid (just as $\vec{k} \cdot \vec{r}$ is).
End of explanation
"""


rv = batoid.RayVector(0,0,0, 0,0,1)
print(rv.r)
print(rv.positionAtTime(1.0))  # propagated from (0,0,0) to (0,0,1)
rv2 = rv.propagate(1.0)  # propagate ray forward to t=1
print(rv.r)  # propagate works in-place
print(rv2.r)  # the return value is a reference to the original
print(rv == rv2)  # so they're equal
print(rv is rv2)  # in fact, they're the same object
# If you want to preserver the original ray when performing a propagation, first make a copy.
rv3 = rv.copy().propagate(2.0)
print(rv.r)
print(rv3.r)

# Plot one period of the plane wave with wavelength 800e-9 m
wavelength = 800e-9
rv = batoid.RayVector(0,0,0, 0,0,1, t=0, wavelength=wavelength)
ts = np.linspace(0, wavelength)
plt.plot(ts, [np.cos(rv.phase([0,0,0], t)) for t in ts])
plt.axhline(0, c='k')
plt.xlabel("c t (meters)")
plt.ylabel("amplitude")
plt.show()

# We can also look at a fixed time over a small region of space
# Since the wave is traveling in the +z direction, its amplitude
# varies with z, but not with x or y
wavelength = 800e-9
rv = batoid.RayVector(0,0,0, 0,0,1, t=0, wavelength=wavelength)
xs = ys = zs = np.linspace(0, wavelength)
plt.plot(xs, [np.cos(rv.phase([x,0,0], 0)) for x in xs], label='x')
plt.plot(ys, [np.cos(rv.phase([0,y,0], 0))+0.1 for y in ys], label='y+0.1')
plt.plot(zs, [np.cos(rv.phase([0,0,z], 0)) for z in zs], label='z')
plt.xlabel("position (meters)")
plt.ylabel("amplitude")
plt.legend()
plt.show()

"""
Explanation: RayVector methods
The first set of methods available to RayVector is for propagation of photons through time.  The positionAtTime method accepts a time (still in meters) and returns the position of the RayVector photons at that past, present, or future time.  The propagate method similarly propagates the rays to a given time.  Note that the return value of propagate is a reference to the original RayVector (not a copy).
For the second set of RayVector methods, we change our conception of what a ray is: from a propagating infinitesimal point to a propagating plane wave.  I.e., we imagine that the velocity v attribute is that of a plane wave, and that the position r and time t together indicate a point in space-time where the wave amplitude is a maximum.  Note that these conditions will remain true as we propagate the wave in time.
The phase method accepts a position and time and returns the phase $\phi$ of these plane waves in radians.  The amplitude method returns the complex amplitudes $\exp(i \phi)$ .
End of explanation
"""


# An example of a square grid of rays:
rv1 = batoid.RayVector.asGrid(
    backDist=10,  # roughly how far back from the origin the rays are created
    nx=15,  # number of rays on a side
    lx=10.0,  # length of one side of the grid in meters
    theta_x=np.deg2rad(0.0),  # Field angle of rays in x-direction
    theta_y=np.deg2rad(0.0),  # Field angle of rays in y-direction
    wavelength=500e-9
)

# Polar RayVector with similar extent
rv2 = batoid.RayVector.asPolar(
    backDist=10,
    nrad=20,  # number of radii
    naz=50,  # number of azimuths on outermost ring.
    outer=5.0,  # outer radius.  Default inner radius is 0.0
    theta_x=np.deg2rad(0.0),
    theta_y=np.deg2rad(0.0),
    wavelength=500e-9
)

# Spokes RayVector with similar extent
rv3 = batoid.RayVector.asSpokes(
    backDist=10,
    rings=10,  # number of rings
    spokes=10,  # number of spokes
    outer=5.0,  # outer radius
    inner=2.5,  # inner radius
    theta_x=np.deg2rad(0.0),
    theta_y=np.deg2rad(0.0),
    wavelength=500e-9
)

plt.scatter(rv1.x, rv1.y, c='b', s=25)
plt.scatter(rv2.x, rv2.y, c='r', s=25)
plt.scatter(rv3.x, rv3.y, c='c', s=25)
plt.xlabel("x")
plt.ylabel("y")
plt.xlim(-6, 6)
plt.ylim(-6, 6)
plt.gca().set_aspect("equal")
plt.show()

plt.scatter(rv1.x, rv1.z, c='b', s=25)
plt.scatter(rv2.x, rv2.z, c='r', s=10)
plt.scatter(rv3.x, rv3.z, c='c', s=5)
plt.xlabel("x")
plt.ylabel("z")
plt.xlim(-6, 6)
plt.ylim(-1, 11)
plt.gca().set_aspect("equal")
plt.show()

# Repeat with a non-zero field angle.

# An example of a square grid of rays:
rv1 = batoid.RayVector.asGrid(
    backDist=10,  # roughly how far back from the origin the rays are created
    nx=15,  # number of rays on a side
    lx=10.0,  # length of one side of the grid in meters
    theta_x=np.deg2rad(5),  # Field angle of rays in x-direction
    theta_y=np.deg2rad(10),  # Field angle of rays in y-direction
    wavelength=500e-9
)

# Polar RayVector with similar extent
rv2 = batoid.RayVector.asPolar(
    backDist=10,
    nrad=20,  # number of radii
    naz=50,  # number of azimuths on outermost ring.
    outer=5.0,  # outer radius.  Default inner radius is 0.0
    theta_x=np.deg2rad(5),
    theta_y=np.deg2rad(10),
    wavelength=500e-9
)

# Spokes RayVector with similar extent
rv3 = batoid.RayVector.asSpokes(
    backDist=10,
    rings=10,  # number of rings
    spokes=10,  # number of spokes
    outer=5.0,  # outer radius
    inner=2.5,  # inner radius
    theta_x=np.deg2rad(5),
    theta_y=np.deg2rad(10),
    wavelength=500e-9
)

plt.scatter(rv1.x, rv1.y, c='b', s=25)
plt.scatter(rv2.x, rv2.y, c='r', s=25)
plt.scatter(rv3.x, rv3.y, c='c', s=25)
plt.xlabel("x")
plt.ylabel("y")
plt.xlim(-6, 6)
plt.ylim(-6, 6)
plt.gca().set_aspect("equal")
plt.show()

plt.scatter(rv1.x, rv1.z, c='b', s=25)
plt.scatter(rv2.x, rv2.z, c='r', s=10)
plt.scatter(rv3.x, rv3.z, c='c', s=5)
plt.xlabel("x")
plt.ylabel("z")
plt.xlim(-6, 6)
plt.ylim(-1, 11)
plt.gca().set_aspect("equal")
plt.show()

# If you have ipyvolume installed, you can visualize the ray grids in 3D.
has_ipv = True
try:
    import ipyvolume as ipv
except ImportError:
    has_ipv = False

if has_ipv:
    ipv.figure()
    ipv.quiver(rv1.x, rv1.y, rv1.z, rv1.vx, rv1.vy, rv1.vz, size=5, color='red')
    ipv.quiver(rv2.x, rv2.y, rv2.z, rv2.vx, rv2.vy, rv2.vz, size=5, color='blue')
    ipv.quiver(rv3.x, rv3.y, rv3.z, rv3.vx, rv3.vy, rv3.vz, size=5, color='cyan')
    ipv.xyzlim(10)
    ipv.view(0,90)
    ipv.show()

"""
Explanation: RayVector factory functions
It can be slow to construct a large grid of positions in python just to then assemble them into RayVectors.  To help speed up this construction, batoid includes a few RayVector factory functions that can be used to rapidly create commonly desired grids of rays.  These include RayVector.asGrid, RayVector.asPolar, and RayVector.asSpokes.  In each case, the goal is to sample either a plane wave or spherical wave incident upon some optical system by creating a grid of either parallel or co-radial co-phased Rays.  For RayVector.asGrid, the Rays are assembled onto a parallelogram (commonly just a square), for RayVector.asPolar the pattern is more circularly symmetric.  The RayVector.asSpokes function is a bit special, in that it does not uniformly sample the region of interest.  It does have a use, however, for integrating functions on a circle or annulus using Gaussian quadrature.  There are a number of different options for each of these factory functions to support different use cases.
End of explanation
"""


