#!/usr/bin/env python
# coding: utf-8

# RayOpt Tutorial
# ========
# 
# 

# In[1]:


get_ipython().run_line_magic('pylab', 'inline')
import warnings
import numpy as np
import matplotlib.pyplot as plt

import rayopt as ro

# ignore matplotlib and numpy warning each other
warnings.simplefilter("ignore", FutureWarning)
# ignore floating point exceptions
np.seterr(divide="ignore", invalid="ignore")
# by default only print 4 significant digits
np.set_printoptions(precision=4)


# In[2]:


s = ro.system_from_yaml("""
object:
  pupil:
    radius: 1
elements:
- {}
- {distance: 1, material: 1.5, roc: 5, radius: 1}
- {distance: .2, material: 1.0}
- {}
""")
s.update()
print(s)


# In[3]:


ro.Analysis(s)


# In[4]:


lib = ro.Library.one()
for g in lib.session.query(
    ro.library.Material).filter(
    ro.library.Material.name.contains("N-BK7")):
    print(g.name, g.catalog.name, g.catalog.source)


# In[5]:


s[1].material = lib.get("material", "SCHOTT-BK|N-BK7", source="rii")
s.object.angle = np.deg2rad(5)
s.fields = 0, .7, 1.
s.update()
s.paraxial.resize()
s.resize_convex()
s.paraxial.refocus()
ro.Analysis(s)


# In[6]:


s.object.pupil.radius = .1
q = ro.GaussianTrace(s)
fig, ax = plt.subplots()
s.plot(ax)
q.plot(ax, color="red", scale=10)
s.object.pupil.radius = 1.


# In[7]:


print(q)


# In[8]:


p = ro.PolyTrace(s, kmax=3)
print(p)


# In[9]:


s.paraxial.height[1], s.paraxial.lagrange, p.stvwof[-1, 0, 1]


# In[10]:


q = 1
#q *= 2
#q *= s.paraxial.height[1]
#q *= s.paraxial.lagrange
_ = p.seidel(*p.st())
print(np.array([v*q for k, v in _[:5]]))
print("\n".join(s.paraxial.print_transverse3()))


# In[11]:


print("\n".join(p.print_trace()))


# In[12]:


y = np.atleast_2d([1., 0, 0])
-s[1].surface_sag(y), p.stvwof[1, -1].view(p.Simplex)(*y.T[0]**2, [0.], [0.])


# In[ ]: