text: # -*- coding: utf-8 -*-
  
  import numpy as _np
  from scipy.special import hermite, genlaguerre
  from scipy.interpolate import RectBivariateSpline
  
  
  USE_CV2 = False
  
  if USE_CV2:
      import cv2
  
  USE_SKIMAGE = False
  if USE_SKIMAGE:
      from skimage.restoration import unwrap_phase as _unwrap_phase
  else:
      #used in PhaseUnwrap
      # own implementation currently slower, but seems a little more stable
      # with jumpy phases and of course removes dependency on the extra package
      from .unwrap import unwrap_phase as _unwrap_phase
  
  from .units import deg
  from .field import Field
  from .subs import Inv_Squares
  
  
  def BeamMix(Fin1, Fin2):
      """
      Fout = BeamMix(F1, F2)
  
      :ref:`Addition of the fields F1 and F2. <BeamMix>`
  
      Args::
      
          F1, F2: input fields
          
      Returns::
        
          Fout: output field (N x N square array of complex numbers).
          
      Example:
      
      :ref:`Two holes interferometer <Young>`
      
      """
      if Fin1.field.shape != Fin2.field.shape:
          raise ValueError('Field sizes do not match')
      Fout = Field.copy(Fin1)
      Fout.field += Fin2.field
      return Fout
  
  def CircAperture(R, x_shift, y_shift, Fin):
      """
      Fout = CircAperture(R, x_shift, y_shift, Fin)
      
      :ref:`Propagates the field through a circular aperture. <CircAperture>`
  
      Args::
      
          R: radius of the aperture
          x_shift, y_shift: shift from the center
          Fin: input field
          
      Returns::
       
          Fout: output field (N x N square array of complex numbers).
              
      Example:
      
      :ref:`Diffraction from a circular aperture <circ_aperture>`
      
      """
      
      #from
      #https://stackoverflow.com/questions/44865023/
      # circular-masking-an-image-in-python-using-numpy-arrays
      Fout = Field.copy(Fin)
      
      Y, X = Fout.mgrid_cartesian
      Y = Y - y_shift
      X = X - x_shift
      
      dist_sq = X**2 + Y**2 #squared, no need for sqrt
      
      Fout.field[dist_sq > R**2] = 0.0
      return Fout
  
  def CircScreen(R, x_shift, y_shift, Fin):
      """
      Fout = CircScreen(R, x_shift, y_shift, Fin)
                  
      :ref:`Diffracts the field by a circular screen. <CircScreen>`
  
      Args::
      
          R: radius of the screen
          x_shift, y_shift: shift from the center
          Fin: input field
          
      Returns::
       
          Fout: output field (N x N square array of complex numbers).
              
      Example:
      
      :ref:`Spot of Poisson <Poisson>`
      
      """
      
      #from
      #https://stackoverflow.com/questions/44865023/
      # circular-masking-an-image-in-python-using-numpy-arrays
      Fout = Field.copy(Fin)
      
      Y, X = Fout.mgrid_cartesian
      Y = Y - y_shift
      X = X - x_shift
      dist_sq = X**2 + Y**2 #squared, no need for sqrt
      
      Fout.field[dist_sq <= R**2] = 0.0
      return Fout
  
  def GaussAperture(w, x_shift, y_shift, T, Fin):
      """
      Fout = GaussAperture(w, x_shift, y_shift, T, Fin)
      
      :ref:`Inserts an aperture with a Gaussian shape in the field. <GaussAperture>`
      
          :math:`F_{out}(x,y)= \\sqrt{T}e^{ -\\frac{ x^{2}+y^{2} }{2w^{2}} } F_{in}(x,y)`
  
      Args::
      
          w: 1/e intensity width
          x_shift, y_shift: shift from center
          T: center intensity transmission
          Fin: input field
      
      Returns::
      
          Fout: output field (N x N square array of complex numbers).
  
      """ 
  
      Fout = Field.copy(Fin)
      
      Y, X = Fout.mgrid_cartesian
      Y = Y - y_shift
      X = X - x_shift
  
      w2=w*w*2
      SqrtT=_np.sqrt(T)
      Fout.field*=SqrtT*_np.exp(-(X*X+Y*Y)/w2)
      return Fout
  
  def SuperGaussAperture(w, x_shift, y_shift, T, n, Fin):
      """
      Fout = SuperGaussAperture(w, x_shift, y_shift, T, n, Fin)
      
      :ref:`Inserts an aperture with a Gaussian shape in the field. <GaussAperture>`
      
          :math:`F_{out}(x,y)= \\sqrt{T}e^{ -\\left [ \\frac{ x^{2}+y^{2} }{2w^{2}} \\right ]^n } F_{in}(x,y)`
  
      Args::
      
          w: 1/e intensity width
          x_shift, y_shift: shift from center
          T: center intensity transmission
          n: power of the super Gaussian
          Fin: input field
      
      Returns::
      
          Fout: output field (N x N square array of complex numbers).
  
      """ 
  
      Fout = Field.copy(Fin)
      
      Y, X = Fout.mgrid_cartesian
      Y = Y - y_shift
      X = X - x_shift
  
      w2=w*w*2
      SqrtT=_np.sqrt(T)
      Fout.field*=SqrtT*_np.exp(-((X*X+Y*Y)/w2)**n)
      return Fout
  
  def GaussScreen(w, x_shift, y_shift, T, Fin):
      """
      Fout = GaussScreen(w, x_shift, y_shift, T, Fin)
      
      :ref:`Inserts a screen with a Gaussian shape in the field. <GaussScreen>`
  
          :math:`F_{out}(x,y)= \\sqrt{1-(1-T)e^{ -\\frac{ x^{2}+y^{2} }{w^{2}} }} F_{in}(x,y)`
  
      Args::
      
          w: 1/e intensity width
          x_shift, y_shift: shift from center
          T: center intensity transmission
          Fin: input field
      
      Returns::
      
          Fout: output field (N x N square array of complex numbers).
  
      """  
      Fout = Field.copy(Fin)
      
      Y, X = Fout.mgrid_cartesian
      Y = Y - y_shift
      X = X - x_shift
  
      w2=w*w
      Fout.field*=_np.sqrt(1-(1-T)*_np.exp(-(X*X+Y*Y)/w2))
      return Fout
      
  def GaussHermite( n, m, A, w0, Fin):
      """
      Fout = GaussHermite(m, n, A, w0, Fin)
      
      :ref:`Substitutes a Hermite-Gauss mode (beam waist) in the field. <GaussHermite>`
  
          :math:`F_{m,n}(x,y,z=0) = A H_m\\left(\\dfrac{\\sqrt{2}x}{w_0}\\right)H_n\\left(\\dfrac{\\sqrt{2}y}{w_0}\\right)e^{-\\frac{x^2+y^2}{w_0^2}}`
  
      Args::
          
          m, n: mode indices
          A: Amplitude
          w0: Guaussian spot size parameter in the beam waist (1/e amplitude point)
          Fin: input field
          
      Returns::
      
          Fout: output field (N x N square array of complex numbers).            
          
      Reference::
      
          A. Siegman, "Lasers", p. 642
  
      """
  
      Fout = Field.copy(Fin)
      
      Y, X = Fout.mgrid_cartesian
      #Y = Y - y_shift
      #X = X - x_shift
  
      sqrt2w0=_np.sqrt(2.0)/w0
      w02=w0*w0
  
      Fout.field  = A * hermite(m)(sqrt2w0*X)*hermite(n)(sqrt2w0*Y)*_np.exp(-(X*X+Y*Y)/w02)
      return Fout
      
  def GaussLaguerre(p, l, A, w0, Fin):
      """
      Fout = GaussLaguerre(p, l, A, w0, Fin)
  
      :ref:`Substitutes a Laguerre-Gauss mode (beam waist) in the field. <GaussLaguerre>`
  
          :math:`F_{p,l}(x,y,z=0) = A \\left(\\frac{\\rho}{2}\\right)^{\\frac{|l|}{2} }L^p_l\\left(\\rho\\right)e^{-\\frac{\\rho}{2}}\\cos(l\\theta)`,
          
          with :math:`\\rho=\\frac{2(x^2+y^2)}{w_0^2}`
  
      Args::
          
          p, l: mode indices
          A: Amplitude
          w0: Guaussian spot size parameter in the beam waist (1/e amplitude point)
          Fin: input field
          
      Returns::
      
          Fout: output field (N x N square array of complex numbers).            
          
      Reference::
      
          A. Siegman, "Lasers", p. 642
  
      """
  
      Fout = Field.copy(Fin)
      R, Phi = Fout.mgrid_polar
      w02=w0*w0
      la=abs(l)
      rho = 2*R*R/w02
      Fout.field = A * rho**(la/2) * genlaguerre(p,la)(rho) * _np.exp(-rho/2) * _np.cos(l*Phi)
      return Fout
  
  
  def IntAttenuator(att, Fin):
      """
      Fout = IntAttenuator(att, Fin)
      
      :ref:`Attenuates the intensity of the field. <IntAttenuator>`
          
          :math:`F_{out}(x,y)=\\sqrt{att}F_{in}(x,y)`
          
      Args::
      
          att: intensity attenuation factor
          Fin: input field
          
      Returns::
      
          Fout: output field (N x N square array of complex numbers).
     
      """
      Efactor = _np.sqrt(att) #att. given as intensity
      Fout = Field.copy(Fin)
      Fout.field *= Efactor
      return Fout
  
  def Intensity(flag, Fin):
      """
      I=Intensity(flag,Fin)
      
      :ref:`Calculates the intensity of the field. <Intensity>`
      
      :math:`I(x,y)=F_{in}(x,y).F_{in}(x,y)^*`
      
      Args::
      
          flag: 0= no normalization, 1=normalized to 1, 2=normalized to 255 (for bitmaps)
          Fin: input field
          
      Returns::
      
          I: intensity distribution (N x N square array of doubles)
  
      """
      I = _np.abs(Fin.field)**2
      if flag > 0:
          Imax = I.max()
          if Imax == 0.0:
              raise ValueError('Cannot normalize because of 0 beam power.')
          I = I/Imax
          if flag == 2:
              I = I*255
      return I
  
  def Interpol(new_size, new_N, x_shift, y_shift, angle, magnif, Fin):
      """
      Fout = Interpol(NewSize, NewN, x_shift, y_shift, angle, magnif, Fin)
      
      :ref:`Interpolates the field to a new grid size, grid dimension. <Interpol>`
      
      Args::
      
          NewSize: the new grid size
          NewN: the new grid dimension
          x_shift, y_shift: shift of the field
          angle: rotation of the field in degrees
          magnif: magnification of the field amplitude
          Fin: input field
          
      Returns::
          
          Fout: output field (Nnew x Nnew square array of complex numbers).
    
      """
  
      Fout = Field.begin(new_size, Fin.lam, new_N)
      Fout.field[:,:] = 0.0
      
      legacy = True
      if legacy:
          Pi = 3.141592654 #compare Cpp results numerically
      else:
          Pi = _np.pi #more accurate, but slightly different results
      angle *= Pi/180.
      cc=_np.cos(angle)
      ss=_np.sin(angle)
      
      if legacy:
          #dx defined differently
          size_old = Fin.siz
          old_number = Fin.N
          dx_old = size_old/(old_number-1)
          on21 = int(old_number/2)
          Xold = dx_old * _np.arange(-on21, old_number-on21)
          Yold = dx_old * _np.arange(-on21, old_number-on21)
      else:
          Xold = Fin.xvalues
          Yold = Fin.yvalues
      
      if legacy:
          dx_new = new_size/(new_N-1) #TODO legacy, once again without -1 seems correct
          nn21 = int(new_N/2)
          X0 = dx_new * _np.arange(-nn21, new_N-nn21)
          Y0 = dx_new * _np.arange(-nn21, new_N-nn21)
          X0, Y0 = _np.meshgrid(X0, Y0)
      else:
          dx_new = Fout.dx
          Y0, X0 = Fout.mgrid_cartesian #note swapped order!
      X0 -= x_shift
      Y0 -= y_shift
      Xnew = (X0*cc + Y0*ss)/magnif
      Ynew = (X0*(-ss) + Y0* cc)/magnif
      
      xmin, xmax = Xold[0], Xold[-1]
      ymin, ymax = Yold[0], Yold[-1]
      #filter strictly inside (not <=) since edge pixels seem wrong in interp
      filtmask = ((Xnew > xmin) & (Xnew < xmax) &
                  (Ynew > ymin) & (Ynew < ymax))
      # same goes for Cpp lightpipes, interpolating a 20x20 grid to a 20x20 grid
      # of same size will have 0s along the edges and only 18x18 useful pixels
      
      #instead of calling interp for all pixels, only call for those new pixels
      # who's coordinates (transformed to old) are inside old grid box
      Xmask = Xnew[filtmask] #flat list of X-values, not meshgrid anymore
      Ymask = Ynew[filtmask]
      
      use_scipy_interp = False
      if use_scipy_interp:
          ks = 1 #spline order: linear or higher
          interp_real = RectBivariateSpline(Xold, Yold, Fin.field.real,
                                            kx=ks, ky=ks)
          interp_imag = RectBivariateSpline(Xold, Yold, Fin.field.imag,
                                            kx=ks, ky=ks)
          
          out_real = interp_real(Xmask, Ymask, grid=False)
          out_imag = interp_imag(Xmask, Ymask, grid=False)
          out_comp = out_real + 1j* out_imag
          Fout.field[filtmask] = out_comp
      else:
          out_z = Inv_Squares(Xmask, Ymask, Fin.field, dx_old)
          Fout.field[filtmask] = out_z
      Fout.field /= magnif
      return Fout
  
  def MultIntensity(Intens, Fin):
      """
      Fout = MultIntensity(Intens, Fin)
  
      :ref:`Multiplies the field with a given intensity distribution. <MultIntensity>`
          
      Args::
          
          Intens: N x N square array of real numbers or scalar
          Fin: input field
          
      Returns::
          
          Fout: output field (N x N square array of complex numbers).
    
      """
      if not _np.isscalar(Intens):
          if Intens.shape != Fin.field.shape:
              raise ValueError('Intensity pattern shape does not match field size')
      Fout = Field.copy(Fin)
      Efield = _np.sqrt(Intens)
      Fout.field *= Efield
      return Fout
  
  
  def MultPhase(Phi, Fin):
      """
      Fout = MultPhase(Phase, Fin)
  
      :ref:`Multiplies the field with a given phase distribution. <MultPhase>`
          
      Args::
          
          Phase: N x N square array of real numbers or scalar
          Fin: input field
          
      Returns::
          
          Fout: output field (N x N square array of complex numbers).
    
      """
      if not _np.isscalar(Phi):
          if Phi.shape != Fin.field.shape:
              raise ValueError('Phase pattern shape does not match field size')
      Fout = Field.copy(Fin)
      Fout.field *= _np.exp(1j*Phi)
      return Fout
  
  
  def Normal(Fin):
      """
      Fout = Normal(Fin)
  
      :ref:`Normalizes the field using beam power. <Normal>`
      
          :math:`F_{out}(x,y)= \\frac{F_{in}(x,y)}{\\sqrt{P}}`
  
          with
          
          :math:`P=\\int \\int F_{in}(x,y)^2 dx dy`
      
      Args::
          
          Fin: input field
          
      Returns::
          
          Fout: output field (N x N square array of complex numbers).
    
      """
      Fabs = _np.abs(Fin.field)**2
      Fabs *= Fin.dx**2
      Ptot = Fabs.sum()
      if Ptot == 0.0:
          raise ValueError('Error in Normal(Fin): Zero beam power!')
      Fout = Field.copy(Fin)
      Fout.field *= _np.sqrt(1/Ptot)
      return Fout
  
  
  def Phase(Fin, unwrap = False, units='rad', blank_eps=0):
      """
      Phi=Phase(Fin)
      
      :ref:`Calculates the phase of the field. <Phase>`
      
      
      Args::
      
          Fin: input field
          unwrap: Call PhaseUnwarp on the extracted Phase, Default is False
          units: 'opd': returned in [meters] of optical path length
                  'lam': returned in multiples of lambda
                  'rad': returned in multiples of 2pi phase jumps (default)
          blank_eps: [fraction] of max. Intensity at which to blank the phase
              and replace the value with numpy.nan (e.g. 1e-3==0.1%)
              Set to 0 or None to disable
          
      Returns::
      
          Phi: phase distribution (N x N square array of doubles)
  
      """
      _2pi = 2*_np.pi
      Phi = _np.angle(Fin.field)
      if unwrap:
          Phi = PhaseUnwrap(Phi)
      
      if units=='opd':
          Phi = Phi/_2pi*Fin.lam #a PtV of 2pi will yield e.g. 1*lam=1e-6=1um
      elif units=='lam':
          Phi = Phi/_2pi #a PtV of 2pi=6.28 will yield 1 (as in 1 lambda)
      elif units=='rad':
          pass #a PtV of 2pi will yield 6.28 as requested
      else:
          raise ValueError('Unknown value for option units={}'.format(units))
          
      if blank_eps:
          I = Intensity(0,Fin)
          Phi[I<blank_eps*I.max()] = _np.nan
          
      return Phi
  
  def PhaseSpiral(Fin, m = 1):
      """ 
      Fout = PhaseSpiral(Fin, m = 1)
      Multiplies Fin with a spiral phase distribution.
      
      Args::
      
          required:
          Fin: input field
  
          optional:
          m = 1: order of the spiral.
      """
      Fout = Field.copy(Fin) 
      R, Phi = Fout.mgrid_polar  
      Fout.field *= _np.exp(1j * m * Phi)
      return Fout
  
  def PhaseUnwrap(Phi):
      """
      PhiOut=PhaseUnwrap(PhiIn)
      
      :ref:`Unwraps (removes jumps of pi radians) the phase. <PhaseUnwrap>`
      
      
      Args::
      
          PhiIn: input phase distribution
          
      Returns::
      
          PhiOut: unwrapped phase distribution (N x N square array of doubles)
  
      """
      PhiU = _unwrap_phase(Phi)
      return PhiU
  
  
  def Power(Fin):
      """
      P = Power(Fin)
  
      :ref:`Calculates the total power. <Power>`
          
      Args::
          
          Fin: input field
          
      Returns::
          
          P: output power (real number).
    
      """
      #TODO why does Normal() also sum dx**2 (==integral) while this does not??
      I = _np.abs(Fin.field)**2
      return I.sum()
  
  
  def RandomIntensity(seed, noise, Fin):
      """
      Fout = RandomIntensity(seed, noise, Fin)
  
      :ref:`Adds random intensity to the field <RandomIntensity>`
          
      Args::
          
          seed: seed number for the random noise generator
          noise: level of the noise
          Fin: input field
          
      Returns::
          
          Fout: output field (N x N square array of complex numbers).
    
      """
      #TODO implementation error in original LP: field error, not I error!
      # need to sqrt for that
      Fout = Field.copy(Fin)
      _np.random.seed(int(seed))
      N = Fout.N
      ranint = _np.random.rand(N, N)*noise
      Fout.field += ranint
      return Fout
  
  def RandomPhase(seed, maxPhase, Fin):
      """
      Fout = RandomPhase(seed, maxPhase, Fin)
  
      :ref:`Adds random phase to the field <RandomPhase>`
          
      Args::
          
          seed: seed number for the random noise generator
          maxPhase: maximum phase in radians
          Fin: input field
          
      Returns::
          
          Fout: output field (N x N square array of complex numbers).
    
      """
      #2020023 - ldo - tested similar result as Cpp version, although not 
      # 1:1 since seed is different in numpy
      Fout = Field.copy(Fin)
      _np.random.seed(int(seed))
      N = Fout.N
      ranphase = (_np.random.rand(N, N)-0.5)*maxPhase
      Fout.field *= _np.exp(1j * ranphase)
      return Fout
  
  
  def RectAperture(sx, sy, x_shift, y_shift, angle, Fin):
      """
      Fout = RectAperture(w, h, x_shift, y_shift, angle, Fin)
      
      :ref:`Propagates the field through a rectangular aperture. <RectAperture>`
  
      Args::
      
          w: width of the aperture
          h: height of the aperture
          x_shift, y_shift: shift from the center
          angle: rotation angle in degrees 
          Fin: input field
          
      Returns::
       
          Fout: output field (N x N square array of complex numbers).
  
      """
      Fout = Field.copy(Fin)
      yy, xx = Fout.mgrid_cartesian
      yy = yy - y_shift
      xx = xx - x_shift
      if angle!=0.0:
          ang_rad = -1*angle*deg #-1 copied from Cpp convention
          cc = _np.cos(ang_rad)
          ss = _np.sin(ang_rad)
          xxr = cc * xx + ss * yy
          yyr = -ss * xx + cc * yy
          yy, xx = yyr, xxr
      matchx = _np.abs(xx) > sx/2
      matchy = _np.abs(yy) > sy/2
      Fout.field[matchx | matchy] = 0.0
      return Fout
  
  
  def RectScreen(sx, sy, x_shift, y_shift, angle, Fin):
      """
      Fout = RectScreen(w, h, x_shift, y_shift, angle, Fin)
      
      :ref:`Diffracts the field by a rectangular screen. <RectScreen>`
  
      Args::
      
          w: width of the screen
          h: height of the screen
          x_shift, y_shift: shift from the center
          angle: rotation angle in degrees 
          Fin: input field
          
      Returns::
       
          Fout: output field (N x N square array of complex numbers).
  
      """
      Fout = Field.copy(Fin)
      yy, xx = Fout.mgrid_cartesian
      yy = yy - y_shift
      xx = xx - x_shift
      if angle!=0.0:
          ang_rad = -1*angle*deg #-1 copied from Cpp convention
          cc = _np.cos(ang_rad)
          ss = _np.sin(ang_rad)
          xxr = cc * xx + ss * yy
          yyr = -ss * xx + cc * yy
          yy, xx = yyr, xxr
      matchx = _np.abs(xx) <= sx/2
      matchy = _np.abs(yy) <= sy/2
      Fout.field[matchx & matchy] = 0.0
      return Fout
  
  
  def Strehl(Fin):
      """
      S = Strehl( Fin)
  
      :ref:`Calculates the Strehl value of the field <Strehl>`
          
      Args::
          
          Fin: input field
          
      Returns::
          
          S: Strehl value (real number).
    
      """
      normsq = _np.abs(Fin.field).sum()**2
      if normsq == 0.0:
          raise ValueError('Error in Strehl: Zero beam power')
      strehl = _np.real(Fin.field).sum()**2 + _np.imag(Fin.field).sum()**2
      strehl = strehl/normsq
      return strehl
  
  def SubIntensity(Intens, Fin):
      """
      Fout = SubIntensity(Intens, Fin)
  
      :ref:`Substitutes  a given intensity distribution in the field with. <SubIntensity>`
          
      Args::
          
          Intens: N x N square array of real numbers >= 0
          Fin: input field
          
      Returns::
          
          Fout: output field (N x N square array of complex numbers).
    
      """
      Fout = Field.copy(Fin)
      Intens = _np.asarray(Intens)
      if Intens.shape != Fout.field.shape:
          raise ValueError('Intensity map has wrong shape')
      phi = _np.angle(Fout.field)
      Efield = _np.sqrt(Intens)
      Fout.field = Efield * _np.exp(1j * phi)
      return Fout
  
  def SubPhase(Phi, Fin):
      """
      Fout = SubPhase(Phi, Fin)
  
      :ref:`Substitutes  a given phase distribution in the field with. <SubPhase>`
          
      Args::
          
          Phase: N x N square array of real numbers
          Fin: input field
          
      Returns::
          
          Fout: output field (N x N square array of complex numbers).
    
      """
      Fout = Field.copy(Fin)
      Phi = _np.asarray(Phi)
      if Phi.shape != Fin.field.shape:
          raise ValueError('Phase map has wrong shape')
      oldabs = _np.abs(Fout.field)
      Fout.field = oldabs * _np.exp(1j * Phi)
      return Fout
  
  
  