import pykat.exceptions as pkex
import numpy as np
import math
import copy
import warnings
import cmath
from scipy.special import hermite
from pykat.SIfloat import SIfloat

def fit_beam_size(z_data, w_data, w0guess, z0guess, lam=1064e-9, show_plot=True,
                  plotpts=100, title='Beam scan fit', w2zero=True, filename=None):
    """
    Function to fit a beam size (w0 and z) to a 2D intensity array. Plotting can
    be switched on or off.

    Returns best fit of (w0, z)
    
    Contributed by Paul Fulda on 15/01/16
    """
    
    import scipy.optimize
    
    def zw0z02w(z, w0, z0, lam):
        z_R = pl.pi*w0**2/lam
        w = w0*pl.sqrt(1+((z-z0)/z_R)**2)
        return w
    
    popt,pcov = scipy.optimize.curve_fit(lambda z_data, w0, z0: zw0z02w(z_data, w0, z0, lam), z_data, w_data, p0=[w0guess,z0guess])
    w0out=popt[0]
    z0out=popt[1]
    
    z_fit = pl.linspace(min(z_data),max(z_data),plotpts)
    w_fit = zw0z02w(z_fit, w0out, z0out, lam)


    if show_plot or filename is not None:
        import pylab as pl
        
        um=1e6
        
        pl.figure()
        pl.plot(z_data,w_data*um,'bo', label = 'Data')
        pl.plot(z_fit,w_fit*um,'b', label = 'Fit')
        pl.tight_layout
        pl.grid()
        pl.xlabel('Position [m]')
        pl.ylabel('Beam size [$\mu$m]')
        pl.xlim((min(z_data),max(z_data)))
        
        if w2zero:
            pl.ylim((0,max(w_data)*um))
            
        pl.legend(loc=0)
        pl.title(title)
        
        if filename is not None:
            pl.savefig(filename)
            
    return w0out, z0out      
    
    

class gauss_param(object):
    """
    Use beam_param instead, will be the future name of this object.
    
    Gaussian beam complex parameter
    
    beam_param is effectively a complex number with extra
    functionality to determine beam parameters.
    
    Defaults to 1064e-9m for wavelength and refractive index 1
    usage:
        q = gauss_param(w0=w0, z=z)
        q = gauss_param(z=z, zr=zr)
        q = gauss_param(w=w, rc=rc)
        q = gauss_param(q=a) # where a is a complex number
        
        or change default wavelength and refractive index with:
        
        q = gauss_param(wavelength, nr, w0=w0, zr=zr)
    """
    
    def __init__(self, wavelength=1064e-9, nr=1, *args, **kwargs):
        if self.__class__ != beam_param:
            warnings.warn("Name changed. Use beam_param instead of gauss_param.")
            
        self.__q = None
        self.__lambda = SIfloat(wavelength)
        self.__nr = SIfloat(nr)
        
        if len(args) == 1:
            self.__q = complex(args[0])
        
        elif len(kwargs) == 1:
            if "q" in kwargs:
                self.__q = complex(kwargs["q"])        
            else:
                raise pkex.BasePyKatException("Must specify: z and w0 or z and zr or rc and w or q, to define the beam parameter")
                
        elif len(kwargs) == 2:        
            
            if "w0" in kwargs and "z" in kwargs:
                q = SIfloat(kwargs["z"]) + 1j * math.pi*SIfloat(kwargs["w0"])**2/(self.__lambda/self.__nr)
            elif "z" in kwargs and "zr" in kwargs:
                q = SIfloat(kwargs["z"]) + 1j * SIfloat(kwargs["zr"]) 
            elif "rc" in kwargs and "w" in kwargs:
                one_q = 1 / SIfloat(kwargs["rc"]) - 1j * SIfloat(wavelength) / (math.pi * SIfloat(nr) * SIfloat(kwargs["w"])**2)
                q = 1/one_q
            else:
                raise pkex.BasePyKatException("Must specify: z and w0 or z and zr or rc and w or q, to define the beam parameter")
                
            self.__q = q
        else:
            raise pkex.BasePyKatException("Incorrect usage for gauss_param constructor")
    
    @property
    def wavelength(self): return self.__lambda
    @wavelength.setter
    def wavelength(self,value): self.__lambda = SIfloat(value)
    
    @property
    def nr(self): return self.__nr
    
    @property
    def q(self): return self.__q
    
    @property
    def z(self): return self.__q.real
    
    @property
    def zr(self): return self.__q.imag
    
    @property
    def w(self):
        return np.abs(self.__q)* np.sqrt(self.__lambda / (self.__nr * math.pi * self.__q.imag))
    
    def beamsize(self, z=None, wavelength=None, nr=None, w0=None):

        if z is None:
            z = self.z
        else:
            z = np.array(z)
                
        if wavelength is None:
            wavelength = self.wavelength
        else:
            wavelength = np.array(wavelength)
            
        if nr is None:
            nr = self.nr
        else:
            nr = np.array(nr)
            
        if w0 is None:
            w0 = self.w0
        else:
            w0 = np.array(w0)
        
        q = z + 1j * math.pi * w0 **2 / wavelength
        
        return np.abs(q)*np.sqrt(wavelength / (nr * math.pi * q.imag))
    
    def gouy(self, z=None, wavelength=None, nr=None, w0=None):
        if z is None:
            z = self.z
        else:
            z = np.array(z)
                
        if wavelength is None:
            wavelength = self.wavelength
        else:
            wavelength = np.array(wavelength)
            
        if nr is None:
            nr = self.nr
        else:
            nr = np.array(nr)
            
        if w0 is None:
            w0 = self.w0
        else:
            w0 = np.array(w0)
        
        q = z + 1j * math.pi * w0 **2 / wavelength
        
        return np.arctan2(q.real, q.imag)
        
    @property
    def w0(self):
        return np.sqrt(self.__q.imag * self.__lambda / (self.__nr * math.pi))    

    @property
    def Rc(self):
        def __rc(z, zr):
            if z != 0:
                return z * (1 + (zr/z)**2)
            else:
                return float("inf")
                
        v = np.vectorize(__rc)
        
        return v(self.z, self.zr)
    
    def curvature(self, z=None, wavelength=None, nr=None, w0=None):
        if z is None:
            z = self.z
        else:
            z = np.array(z)
                
        if wavelength is None:
            wavelength = self.wavelength
        else:
            wavelength = np.array(wavelength)
            
        if nr is None:
            nr = self.nr
        else:
            nr = np.array(nr)
            
        if w0 is None:
            w0 = self.w0
        else:
            w0 = np.array(w0)
        
        q = z + 1j * math.pi * w0 **2 / wavelength
        
        return q.real * (1+ (q.imag/q.real)**2)
        
    @staticmethod
    def overlap(q1, q2):
        """
        Computes the projection from one beam parameter to another to give a measure of the
        overlap between the two beam parameters.
        
        This function was provided by Paul Fulda and Antonio Perreca, which came originally
        from Chris Mueller.
        
        Added on 20/4/2015
        """
        return abs(4*q1.imag * q2.imag)/abs(q1.conjugate()-q2)**2
        
    def conjugate(self):
        return beam_param(self.__lambda, self.__nr, self.__q.conjugate())
    
    def __abs__(self):
        return abs(complex(self.__q))
        
    def __complex__(self):
        return self.__q
    
    def __str__(self):
        return str(self.__q)
    
    def __mul__(self, a):
        return beam_param(self.__lambda, self.__nr, self.__q * complex(a))
    
    def __imul__(self, a):
        self.__q *= complex(a)
        return self
        
    __rmul__ = __mul__
    
    def __add__(self, a):
        return beam_param(self.__lambda, self.__nr, self.__q + complex(a))
    
    def __iadd__(self, a):
        self.__q += complex(a)
        return self
        
    __radd__ = __add__
    
    def __sub__(self, a):
        return beam_param(self.__lambda, self.__nr, self.__q - complex(a))
    
    def __isub__(self, a):
        self.__q -= complex(a)
        return self
        
    def __rsub__(self, a):
        return beam_param(self.__lambda, self.__nr, complex(a) - self.__q)
    
    def __div__(self, a):
        return beam_param(self.__lambda, self.__nr, self.__q / complex(a))
    
    def __idiv__(self, a):
        self.__q /= complex(a)
        return self
    
    def __pow__(self, q):
        return beam_param(self.__lambda, self.__nr, self.__q**q)

    def __neg__(self, q):
        return beam_param(self.__lambda, self.__nr, -self.__q)
        
    def __eq__(self, q):
        if q is None:
            return False
            
        return complex(q) == self.__q
        
    @property
    def real(self): return self.__q.real
    @real.setter
    def real(self, value): self.__q.real = SIfloat(value)
    
    @property
    def imag(self): return self.__q.imag
    @imag.setter
    def imag(self, value): self.__q.imag = SIfloat(value)

    # reverse beam direction 
    def reverse(self):
        self.__q = -1.0 * self.__q.real + 1j * self.__q.imag


class beam_param(gauss_param):
    pass
    
class HG_beam(object):
    
    def __init__(self, qx, qy=None, n=0, m=0):
        self._qx = copy.deepcopy(qx)
        self._2pi_qrt = math.pow(2.0/math.pi, 0.25)
        
        if qy is None:
            self._qy = copy.deepcopy(qx)
        else:
            self._qy = copy.deepcopy(qy)
    
        self._n = int(n)
        self._m = int(m)
        self._hn = hermite(self._n)
        self._hm = hermite(self._m)
        self._calc_constants()
        
    @property
    def n(self): return self._n
    @n.setter
    def n(self,value): 
        self._n = int(value)
        self._calc_constants()
        self._hn = hermite(self._n)

    @property
    def m(self): return self._m
    @m.setter
    def m(self,value): 
        self._m = int(value)
        self._calc_constants()
        self._hm = hermite(self._m)
            
    @property
    def q(self):
        if self._qx.q == self._qy.q:
            return self._qx.q
        else:
            return (self._qx.q, self._qy.q)
    @q.setter
    def q(self, value):
        if value.__class__ == beam_param:
            self._qx = copy.deepcopy(value)
            self._qy = copy.deepcopy(value)
        else:
            self._qx = beam_param(q=complex(value))
            self._qy = beam_param(q=complex(value))
    
    @property
    def qx(self):
        return self._qx.q
        
    @qx.setter
    def qx(self, value):
        if value.__class__ == beam_param:
            self._qx = copy.deepcopy(value)
        else:
            self._qx = beam_param(q=complex(value))
    
    @property
    def qy(self):
        return self._qy.q
        
    @qy.setter
    def qy(self, value):
        if value.__class__ == beam_param:
            self._qy = copy.deepcopy(value)
        else:
            self._qy = beam_param(q=complex(value))
    
    @property
    def constant_x(self):
        return self.__xpre_const
        
    @property
    def constant_y(self):
        return self.__ypre_const
        
    def _calc_constants(self):
        self.__xpre_const = math.pow(2.0/math.pi, 0.25)
        self.__xpre_const *= np.sqrt(1.0/(self._qx.w0 * 2**(self._n) * np.math.factorial(self._n)))
        self.__xpre_const *= np.sqrt(1j*self._qx.imag / self._qx.q)
        self.__xpre_const *= ((1j*self._qx.imag * self._qx.q.conjugate())/(-1j*self._qx.imag * self._qx.q)) ** ( self._n/2.0)
        
        self.__ypre_const = math.pow(2.0/math.pi, 0.25)
        self.__ypre_const *= np.sqrt(1.0/(self._qy.w0 * 2**(self._m) * np.math.factorial(self._m)))
        self.__ypre_const *= np.sqrt(1j*self._qy.imag / self._qy.q)
        self.__ypre_const *= ((1j*self._qy.imag * self._qy.q.conjugate())/(-1j*self._qy.imag * self._qy.q)) **(self._m/2.0)
    
        self.__sqrt2_wxz = math.sqrt(2) / self._qx.w
        self.__sqrt2_wyz = math.sqrt(2) / self._qy.w
        
        self.__kx =  2*math.pi / self._qx.wavelength
        self.__ky =  2*math.pi / self._qy.wavelength
        
        self.__invqx = 1/ self._qx.q
        self.__invqy = 1/ self._qy.q
    
    def Un(self, x):
        return self.__xpre_const * self._hn(self.__sqrt2_wxz * x) * np.exp(-0.5j * self.__kx * x*x * self.__invqx)
    
    def Um(self, y):
        return self.__ypre_const * self._hm(self.__sqrt2_wyz * y) * np.exp(-0.5j * self.__ky * y*y * self.__invqy)
        
    def Unm(self, x, y):
        _un = self.Un(x)  
        _um = self.Um(y)
        return np.outer(_un, _um)
        
    def plot(self, ndx=100, ndy=100, xscale=4, yscale=4):
        import pylab
        
        xrange = xscale * np.linspace(-self._qx.w, self._qx.w, ndx)
        yrange = yscale * np.linspace(-self._qy.w, self._qy.w, ndy)

        dx = xrange[1]-xrange[0]
        dy = yrange[1]-yrange[0]

        data = self.Unm(xrange,yrange)

        fig = pylab.figure()
        axes = pylab.imshow(np.abs(data), aspect=dx/dy, extent=[min(xrange),max(xrange),min(yrange),max(yrange)])
        pylab.xlabel('x [m]')
        pylab.ylabel('y [m]')
        cbar = fig.colorbar(axes)
        pylab.show()