# -*- coding: utf-8 -*-
"""
Created on Tue Jul  4 10:44:26 2023

@author: Markus Schake
"""

import sys, os
sys.path.append(os.getcwd())
import numpy as np
import math



def zernike_poly_xy_cartesian(n,m,x,y):
    R=np.sqrt(np.atleast_2d(x**2).T+(y**2));
    PHI=np.arctan2(y,np.atleast_2d(x).T);
    Z=np.zeros((len(y),len(x)));
    
    if m>=0:
        Z[:,:]=zernike_R_nm_cartesian(n,np.abs(m),R)*np.cos(m*PHI)
    else:
        Z[:,:]=zernike_R_nm_cartesian(n,np.abs(m),R)*np.sin(np.abs(m)*PHI)
        
    return Z
 
       
def zernike_R_nm_cartesian(n,m,R):
    Rnm=np.zeros((np.shape(R)));
    
    if iseven(n-m):
        for l in range(int(0.5*(n-m)+1)):
            Rnm=Rnm + zernike_R_nm_norm_cartesian(n,m,l)*R**(n-2*l);
    
    return Rnm


def zernike_R_nm_norm_cartesian(n,m,l):
    norm=((-1)**l)*math.factorial(int(n-l))/(math.factorial(int(l))*math.factorial(int(0.5*(n+m)-l))*math.factorial(int(0.5*(n-m)-l)))
    return norm


def isodd(numb):
    # check if a numbber is isodd(number)
    # returns 1 if the number is odd, the function returns 1 else 0
    X=np.where((numb/2)-np.floor(numb/2)>0,1,0)
    return X

def iseven(numb):
    # check if a numbber is iseven(number)
    # returns 1 if the number is even, the function returns 1 else 0
    X=np.where((numb/2)-np.floor(numb/2)==0,1,0)
    return X
    
    
### Example for plotting    
# =============================================================================
# if __name__ == "__main__":
#     try:
#         get_ipython().magic('clear')
#         get_ipython().magic('reset -f')
#         print('cleared variables and console')
#     except:
#         pass
#     
#     Naperture=100
#     x_aperture=np.linspace(-1,1,Naperture);
#     y_aperture=np.linspace(-1,1,Naperture);
#     
#     n=2
#     m=-2
#     Zern=zernike_poly_xy_cartesian(n,m,x_aperture,y_aperture)
#     
#     fig4, ax4 = plt.subplots(subplot_kw={"projection": "3d"})
#     X, Y = np.meshgrid(np.linspace(0,100,100), np.linspace(0,100,100))
#     ax4.set_xlabel('x-pixels/1', rotation=150)
#     ax4.set_ylabel('y-pixels/1')
#     ax4.set_zlabel('height/nm', rotation=60)
#     plt.title('Auxiliary Plot ' + 'n=' + str(n)+ 'm=' + str(m))
#     surf = ax4.plot_surface( X, Y, Zern, cmap=cm.coolwarm,linewidth=0, antialiased=False)
# =============================================================================
    
    
    
    
    