import numpy as np
import pylab as plt
import pyzdde.zdde as pyz

class Component(object):
  def __init__(self, zmx_file, zcontroller, lumultiplier=1e-3, wumultiplier=1e6):
    '''
      Initialise a component.

      [lumultiplier] is used to move from Zemax lens units to another
      physical dimension. By default, Zemax does its calculations in
      mm, so a [lumultiplier] of 1e-3 moves from mm to m.

      [wumultiplier] is used to move from SI units to Zemax wavelength
      units. By default, Zemax does its calculations in
      micron, so a [wumultiplier] of 1e6 moves from m to micron.
    '''
    self.file_pathname = zmx_file
    self.zcontroller = zcontroller
    self.lumultiplier = lumultiplier
    self.wumultiplier = wumultiplier

  def _doAnalysisWFE(self, fields, field_type, wavelength, sampling, 
                     verbose=True, debug=False):
    if not self.zcontroller.isFileAlreadyLoaded(self.file_pathname):
      self.zcontroller.loadZemaxFile(self.file_pathname)
    self.zcontroller.setWavelengthNumberOf(1)
    self.zcontroller.setWavelengthValue(float(wavelength)*self.wumultiplier, 1)
    
    data, headers = self.zcontroller.getAnalysisWFEForFields(fields, field_type,
                                                             sampling=sampling)
    
    return data, headers
  
  def _doRayTrace(self, fields, field_type, wavelength, verbose=True, 
                  debug=False):  
    if not self.zcontroller.isFileAlreadyLoaded(self.file_pathname):
      self.zcontroller.loadZemaxFile(self.file_pathname)
    self.zcontroller.setWavelengthNumberOf(1)
    self.zcontroller.setWavelengthValue(float(wavelength)*self.wumultiplier, 1)
    
    rays = self.zcontroller.doRayTraceForFields(fields, field_type=field_type, 
                                                px=0, py=0)
    tmp = []
    for ray in rays:
      tmp.append({
        'x': ray.x * self.lumultiplier,
        'y': ray.y * self.lumultiplier,
        'z': ray.z * self.lumultiplier,
        'dcos_l': ray.dcos_l,
        'dcos_m': ray.dcos_m,
        'dcos_n': ray.dcos_n
        })
    rays = tmp

    if debug:
      for idx, ray in enumerate(rays):
        if idx == 0:
          print
          print "field\tx\ty\tz"
        print idx, '\t', round(ray['x'], 2), '\t', round(ray['y'], 2), \
              '\t', round(ray['z'], 2)
        if idx == len(rays):
          print
          
      plt.plot()
      for field, ray in zip(fields, rays):
        plt.plot(rays['x'], ray['y'], 'o', label=str('[' + str(field[0]) + 
                                      ', ' + str(field[1]) + 
                                      ']'))
      plt.legend(loc='upper right', numpoints=1)
      plt.show()    

    return rays
    
  def getEFFL(self, wavelength, verbose=False):
    if verbose:
      print "Getting EFFL for component... "
      
    if not self.zcontroller.isFileAlreadyLoaded(self.file_pathname):
      self.zcontroller.loadZemaxFile(self.file_pathname)
    self.zcontroller.setWavelengthNumberOf(1)
    self.zcontroller.setWavelengthValue(float(wavelength)*self.wumultiplier, 1)
    return self.zcontroller.getLensData().EFL*self.lumultiplier
  
  def getENPD(self, wavelength, verbose=False):
    if verbose:
      print "Getting ENPD for component... "
      
    if not self.zcontroller.isFileAlreadyLoaded(self.file_pathname):
      self.zcontroller.loadZemaxFile(self.file_pathname)
    self.zcontroller.setWavelengthNumberOf(1)
    self.zcontroller.setWavelengthValue(float(wavelength)*self.wumultiplier, 1)
    return self.zcontroller.getPupilData().ENPD*self.lumultiplier

  def getEXPD(self, wavelength, verbose=False):
    if verbose:
      print "Getting EXPD for component... "
      
    if not self.zcontroller.isFileAlreadyLoaded(self.file_pathname):
      self.zcontroller.loadZemaxFile(self.file_pathname)
    self.zcontroller.setWavelengthNumberOf(1)
    self.zcontroller.setWavelengthValue(float(wavelength)*self.wumultiplier, 1)
    return self.zcontroller.getPupilData().EXPD*self.lumultiplier    

  def getWFNO(self, wavelength, verbose=False):
    if verbose:
      print "Getting WFNO for component... "
      
    if not self.zcontroller.isFileAlreadyLoaded(self.file_pathname):
      self.zcontroller.loadZemaxFile(self.file_pathname)
    self.zcontroller.setWavelengthNumberOf(1)
    self.zcontroller.setWavelengthValue(float(wavelength)*self.wumultiplier, 1)
    return self.zcontroller.getLensData().realWorkFNum

class Camera(Component):
  def __init__(self, camera_zmx_file, zcontroller, lumultiplier=1e-3, 
    wumultiplier=1e6):
    super(Camera, self).__init__(camera_zmx_file, zcontroller, lumultiplier,
      wumultiplier)
    pass

  def getImXY(self, fields, wavelength, verbose=True, debug=False):
    '''
      Trace the chief ray from each collimated field point through the camera 
      and work out the corresponding (x, y) positions at the image plane.
    '''
    
    if verbose:
      print "Tracing object angles through camera..."
      
    rays = self._doRayTrace(fields, 0, wavelength, verbose=verbose, 
                            debug=debug)
    ImXYs = []
    for ray in rays:
      ImXYs.append((ray['x'], ray['y']))

    return ImXYs
 
  def getWFE(self, fields, wavelength, sampling, verbose=True, debug=False):
    '''
      Get the pupil WFE after passing through the camera.
    '''
    return self._doAnalysisWFE(fields, 0, wavelength, sampling=sampling,
                               verbose=verbose, debug=debug)
    
class Collimator(Component):
  def __init__(self, collimator_zmx_file, zcontroller, lumultiplier=1e-3,
    wumultiplier=1e6):
    super(Collimator, self).__init__(collimator_zmx_file, zcontroller,
                                     lumultiplier, wumultiplier)
    pass
  
  def getOA(self, fields, wavelength, verbose=True, debug=False):
    '''
      Trace the chief ray for each field point in the slit through the 
      collimator and, using directional cosines and the relations 
      (p53 Zemax manual)
    
        tan(alpha) = direction_cosine(x)/direction_cosine(z)   .. (1)
        tan(beta) = direction_cosine(y)/direction_cosine(z)    .. (2)
      
      calculate the output field angles x and y (alpha and beta respectively).
    '''
    
    if verbose:
      print "Tracing object heights through collimator... "
      
    rays = self._doRayTrace(fields, 1, wavelength, verbose=verbose, 
                            debug=debug)
    
    OAs = []
    for ray in rays:
      #print ray['dcos_l'], ray['dcos_m'], ray['dcos_n']
      OAs.append((np.degrees(np.arctan(ray['dcos_l']/ray['dcos_n'])), 
                             np.degrees(np.arctan(ray['dcos_m']/ray['dcos_n']))))
      
    return OAs
  
  def getWFE(self, fields, wavelength, sampling, verbose=True, debug=False):
    '''
      Get the pupil WFE after passing through the collimator.
    '''
    return self._doAnalysisWFE(fields, 1, wavelength, sampling=sampling, 
                               verbose=verbose, debug=debug)

import numpy as np
import pylab as plt
import pyzdde.zdde as pyz

class Component(object):
  def __init__(self, zmx_file, zcontroller, lumultiplier=1e-3, wumultiplier=1e6):
    '''
      Initialise a component.

      [lumultiplier] is used to move from Zemax lens units to another
      physical dimension. By default, Zemax does its calculations in
      mm, so a [lumultiplier] of 1e-3 moves from mm to m.

      [wumultiplier] is used to move from SI units to Zemax wavelength
      units. By default, Zemax does its calculations in
      micron, so a [wumultiplier] of 1e6 moves from m to micron.
    '''
    self.file_pathname = zmx_file
    self.zcontroller = zcontroller
    self.lumultiplier = lumultiplier
    self.wumultiplier = wumultiplier

  def _doAnalysisWFE(self, fields, field_type, wavelength, sampling, 
                     verbose=True, debug=False):
    if not self.zcontroller.isFileAlreadyLoaded(self.file_pathname):
      self.zcontroller.loadZemaxFile(self.file_pathname)
    self.zcontroller.setWavelengthNumberOf(1)
    self.zcontroller.setWavelengthValue(float(wavelength)*self.wumultiplier, 1)
    
    data, headers = self.zcontroller.getAnalysisWFEForFields(fields, field_type,
                                                             sampling=sampling)
    
    return data, headers
  
  def _doRayTrace(self, fields, field_type, wavelength, verbose=True, 
                  debug=False):  
    if not self.zcontroller.isFileAlreadyLoaded(self.file_pathname):
      self.zcontroller.loadZemaxFile(self.file_pathname)
    self.zcontroller.setWavelengthNumberOf(1)
    self.zcontroller.setWavelengthValue(float(wavelength)*self.wumultiplier, 1)
    
    rays = self.zcontroller.doRayTraceForFields(fields, field_type=field_type, 
                                                px=0, py=0)
    tmp = []
    for ray in rays:
      tmp.append({
        'x': ray.x * self.lumultiplier,
        'y': ray.y * self.lumultiplier,
        'z': ray.z * self.lumultiplier,
        'dcos_l': ray.dcos_l,
        'dcos_m': ray.dcos_m,
        'dcos_n': ray.dcos_n
        })
    rays = tmp

    if debug:
      for idx, ray in enumerate(rays):
        if idx == 0:
          print
          print "field\tx\ty\tz"
        print idx, '\t', round(ray['x'], 2), '\t', round(ray['y'], 2), \
              '\t', round(ray['z'], 2)
        if idx == len(rays):
          print
          
      plt.plot()
      for field, ray in zip(fields, rays):
        plt.plot(rays['x'], ray['y'], 'o', label=str('[' + str(field[0]) + 
                                      ', ' + str(field[1]) + 
                                      ']'))
      plt.legend(loc='upper right', numpoints=1)
      plt.show()    

    return rays
    
  def getEFFL(self, wavelength, verbose=False):
    if verbose:
      print "Getting EFFL for component... "
      
    if not self.zcontroller.isFileAlreadyLoaded(self.file_pathname):
      self.zcontroller.loadZemaxFile(self.file_pathname)
    self.zcontroller.setWavelengthNumberOf(1)
    self.zcontroller.setWavelengthValue(float(wavelength)*self.wumultiplier, 1)
    return self.zcontroller.getLensData().EFL*self.lumultiplier
  
  def getENPD(self, wavelength, verbose=False):
    if verbose:
      print "Getting ENPD for component... "
      
    if not self.zcontroller.isFileAlreadyLoaded(self.file_pathname):
      self.zcontroller.loadZemaxFile(self.file_pathname)
    self.zcontroller.setWavelengthNumberOf(1)
    self.zcontroller.setWavelengthValue(float(wavelength)*self.wumultiplier, 1)
    return self.zcontroller.getPupilData().ENPD*self.lumultiplier

  def getEXPD(self, wavelength, verbose=False):
    if verbose:
      print "Getting EXPD for component... "
      
    if not self.zcontroller.isFileAlreadyLoaded(self.file_pathname):
      self.zcontroller.loadZemaxFile(self.file_pathname)
    self.zcontroller.setWavelengthNumberOf(1)
    self.zcontroller.setWavelengthValue(float(wavelength)*self.wumultiplier, 1)
    return self.zcontroller.getPupilData().EXPD*self.lumultiplier    

  def getWFNO(self, wavelength, verbose=False):
    if verbose:
      print "Getting WFNO for component... "
      
    if not self.zcontroller.isFileAlreadyLoaded(self.file_pathname):
      self.zcontroller.loadZemaxFile(self.file_pathname)
    self.zcontroller.setWavelengthNumberOf(1)
    self.zcontroller.setWavelengthValue(float(wavelength)*self.wumultiplier, 1)
    return self.zcontroller.getLensData().realWorkFNum

class Camera(Component):
  def __init__(self, camera_zmx_file, zcontroller, lumultiplier=1e-3, 
    wumultiplier=1e6):
    super(Camera, self).__init__(camera_zmx_file, zcontroller, lumultiplier,
      wumultiplier)
    pass

  def getImXY(self, fields, wavelength, verbose=True, debug=False):
    '''
      Trace the chief ray from each collimated field point through the camera 
      and work out the corresponding (x, y) positions at the image plane.
    '''
    
    if verbose:
      print "Tracing object angles through camera..."
      
    rays = self._doRayTrace(fields, 0, wavelength, verbose=verbose, 
                            debug=debug)
    ImXYs = []
    for ray in rays:
      ImXYs.append((ray['x'], ray['y']))

    return ImXYs
 
  def getWFE(self, fields, wavelength, sampling, verbose=True, debug=False):
    '''
      Get the pupil WFE after passing through the camera.
    '''
    return self._doAnalysisWFE(fields, 0, wavelength, sampling=sampling,
                               verbose=verbose, debug=debug)
    
class Collimator(Component):
  def __init__(self, collimator_zmx_file, zcontroller, lumultiplier=1e-3,
    wumultiplier=1e6):
    super(Collimator, self).__init__(collimator_zmx_file, zcontroller,
                                     lumultiplier, wumultiplier)
    pass
  
  def getOA(self, fields, wavelength, verbose=True, debug=False):
    '''
      Trace the chief ray for each field point in the slit through the 
      collimator and, using directional cosines and the relations 
      (p53 Zemax manual)
    
        tan(alpha) = direction_cosine(x)/direction_cosine(z)   .. (1)
        tan(beta) = direction_cosine(y)/direction_cosine(z)    .. (2)
      
      calculate the output field angles x and y (alpha and beta respectively).
    '''
    
    if verbose:
      print "Tracing object heights through collimator... "
      
    rays = self._doRayTrace(fields, 1, wavelength, verbose=verbose, 
                            debug=debug)
    
    OAs = []
    for ray in rays:
      #print ray['dcos_l'], ray['dcos_m'], ray['dcos_n']
      OAs.append((np.degrees(np.arctan(ray['dcos_l']/ray['dcos_n'])), 
                             np.degrees(np.arctan(ray['dcos_m']/ray['dcos_n']))))
      
    return OAs
  
  def getWFE(self, fields, wavelength, sampling, verbose=True, debug=False):
    '''
      Get the pupil WFE after passing through the collimator.
    '''
    return self._doAnalysisWFE(fields, 1, wavelength, sampling=sampling, 
                               verbose=verbose, debug=debug)

Initialise a component.        [lumultiplier] is used to move from Zemax lens units to another       physical dimension. By default, Zemax does its calculations in       mm, so a [lumultiplier] of 1e-3 moves from mm to m.        [wumultiplier] is used to move from SI units to Zemax wavelength       units. By default, Zemax does its calculations in       micron, so a [wumultiplier] of 1e6 moves from m to micron. Trace the chief ray from each collimated field point through the camera        and work out the corresponding (x, y) positions at the image plane. Get the pupil WFE after passing through the camera. Trace the chief ray for each field point in the slit through the        collimator and, using directional cosines and the relations        (p53 Zemax manual)              tan(alpha) = direction_cosine(x)/direction_cosine(z)   .. (1)         tan(beta) = direction_cosine(y)/direction_cosine(z)    .. (2)              calculate the output field angles x and y (alpha and beta respectively). #print ray['dcos_l'], ray['dcos_m'], ray['dcos_n'] Get the pupil WFE after passing through the collimator.