# -*- coding: utf-8 -*-
"""
Created on Tue April  2 2024

@author: Markus Schake
"""
import sys, os
sys.path.append(os.getcwd())

import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np
from IPython import get_ipython
import json
from zernike_poly_xy_cartesian import*


#### Custom functions
    
class NumpyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)

# Pass this class name as a value to the 'cls' parameter
# of the json.dump() or the json.dumps() function


            
def generate_data(curvatureradius,freeaperture,Naperture,offsetvector,n_zi,m_zi,P_custom):
    print('Generates Dataset')
    print("Sphere with radius of curvature " + str(curvatureradius) + "/mm and free aperture " + str(freeaperture) +"/mm")
    
    
    x_aperture=np.linspace(-freeaperture/2,freeaperture/2,Naperture)
    y_aperture=np.linspace(-freeaperture/2,freeaperture/2,Naperture)
    
    xv, yv = np.meshgrid(x_aperture,y_aperture)
    
    aperture_radius=np.sqrt(xv**2+yv**2)
    nan_index_cut=np.where(aperture_radius>freeaperture/2)
    
    
    #Zernike Fringe Wyant indices (https://en.wikipedia.org/wiki/Zernike_polynomials#Wyant_indices)
    ZFR_nr=np.floor((1+(n_zi+np.abs(m_zi))/2)**2-2*np.abs(m_zi)+(1-np.sign(m_zi))/2)-1
    sort_ind=np.argsort(ZFR_nr)
    
    n_zi_ZFR=n_zi
    m_zi_ZFR=m_zi
        
            
    print('######################')
    print('n_zi_ZFR='+str(n_zi_ZFR))
    print('m_zi_ZFR='+str(m_zi_ZFR))
    print('######################')
    
    ################################################
    x_aperture_zern=np.linspace(-1,1,Naperture)
    y_aperture_zern=np.linspace(-1,1,Naperture)
    #Determine the Zernike raster for the whole aperture rectangle
    Z_nmi_ZFR=np.zeros((Naperture,Naperture,np.size(m_zi)))
    for j in range(np.size(m_zi)):
        Z_nmi_ZFR[:,:,j]=zernike_poly_xy_cartesian(n_zi_ZFR[j],m_zi_ZFR[j],x_aperture_zern,y_aperture_zern)
        
    X_zernike_input_ZFR=np.zeros((Naperture*Naperture,np.size(m_zi)));
    for j in range(np.size(m_zi)):
        temp=Z_nmi_ZFR[:,:,j]
        X_zernike_input_ZFR[:,j]=temp.flatten()

    ################################################
    #Calculate your custom Topography profile
    #Zerniek Terms in mm
    Pzfr=np.zeros((np.size(m_zi),))
    Pzfr[0:np.size(P_custom)]=P_custom
    Fitted_surface_ZFR_custom=X_zernike_input_ZFR @ Pzfr
    Fitted_surface_ZFR_custom=Fitted_surface_ZFR_custom.reshape((Naperture,Naperture))
    Fitted_surface_ZFR_custom[nan_index_cut]=float("nan")
    
    
    #Determine the Topography of the reference wavefront, that perfectly compensates the radius of curvature
    Ref_sphere=np.sqrt(np.abs((curvatureradius)**2-xv**2-yv**2))
    Ref_sphere=Ref_sphere-np.max(np.max(Ref_sphere))
    Ref_sphere[nan_index_cut] = float("nan")
    
    #Determine the profile for 5 different defocus positions
    Nexp=np.size(offsetvector)
    Height_data_matrix=np.zeros((Naperture,Naperture,Nexp))
    Height_data_matrix_diff=np.zeros((Naperture,Naperture,Nexp))
    for nind in range(Nexp):
        Heightprofile=np.sqrt(np.abs((curvatureradius-offsetvector[nind])**2-xv**2-yv**2))
        Heightprofile=Heightprofile-np.max(np.max(Heightprofile))
        #print(Heightprofile.dtype)
        Heightprofile[nan_index_cut] = float("nan")
        Height_data_matrix[:,:,nind]=Heightprofile+Fitted_surface_ZFR_custom
        Height_data_matrix_diff[:,:,nind]=Fitted_surface_ZFR_custom+Heightprofile-Ref_sphere
    #Calculate the surface points of a sphere with the specified information
    #The sphere is given by: curvatureradius^2=x^2+y^2+z^2
    #Determine the Topography information recorded in typical interferograms, by subtracting the reference sphere

    fig_sim, ax_sim = plt.subplots()
    im_sim = ax_sim.imshow((Height_data_matrix[:,:,0]), interpolation='bilinear', origin='lower', cmap=cm.jet, extent=[-freeaperture/2, freeaperture/2, -freeaperture/2, freeaperture/2])
    cb = plt.colorbar(im_sim)
    cb.set_label('Topography input/mm')
    ax_sim.set_title('Topography/mm')
    ax_sim.set_xlabel('x/mm')
    ax_sim.set_ylabel('y/mm')
    
    fig_sim2, ax_sim2 = plt.subplots()
    im_sim2 = ax_sim2.imshow((Fitted_surface_ZFR_custom*10**6), interpolation='bilinear', origin='lower', cmap=cm.jet, extent=[-freeaperture/2, freeaperture/2, -freeaperture/2, freeaperture/2])
    cb = plt.colorbar(im_sim2)
    cb.set_label('Sphericity/nm')
    ax_sim2.set_xlabel('x/mm')
    ax_sim2.set_ylabel('y/mm')
    
    
    #Save Data if not already present
    data_dict={ "Topography": Height_data_matrix,
               "Topography_diff":Height_data_matrix_diff,
               "x-axis":x_aperture,
               "y_axis":y_aperture,
               "defocus":offsetvector,
               "P_Zernike":P_custom,
               "curvatureradius":curvatureradius,
               "freeaperture":freeaperture}


    filepath_data2=os.getcwd()
    directory2 = os.listdir(filepath_data2)
    searchstring2 = "data_dict_sphere"
    data_dict_flag=False
    
    if(searchstring2!=""):
            for fname in directory2:
                if searchstring2 in fname:
                    data_dict_flag=True
                    print('Datafile already exists!')
    
    if(data_dict_flag==False):
        with open('data_dict_sphere.json', 'w') as fp:
            json.dump(data_dict, fp, cls=NumpyEncoder, indent="")
    
    return Height_data_matrix_diff*10**6 #nm
    




def determine_zernike(n_zi,m_zi,offsetvector,P_custom):

    offsetvector=offsetvector*10**6 #in nm
    
    N_cluster=np.size(offsetvector)   

    filepath_data2=os.getcwd()
    directory2 = os.listdir(filepath_data2)
    searchstring2 = "data_dict_sphere"
    data_dict_flag=False
    
    if(searchstring2!=""):
            for fname in directory2:
                if searchstring2 in fname:
                    data_dict_flag=True
                    pathname_load=os.path.join(os.getcwd(),'data_dict_sphere.json')
                    print('Datafile exists.')
    
    if(data_dict_flag==False):
        print('No data file!')
        return
    
    #Zernike Fringe Wyant indices (https://en.wikipedia.org/wiki/Zernike_polynomials#Wyant_indices)
    ZFR_nr=np.floor((1+(n_zi+np.abs(m_zi))/2)**2-2*np.abs(m_zi)+(1-np.sign(m_zi))/2)-1
    sort_ind=np.argsort(ZFR_nr)
    ZFR_count=np.size(n_zi)
    
    n_zi_ZFR=n_zi
    m_zi_ZFR=m_zi
    print('######################')
    print('n_zi_ZFR='+str(n_zi_ZFR))
    print('m_zi_ZFR='+str(m_zi_ZFR))
    print('######################')
    
    #clear old zernike data
    zernike_clustermatrix=np.zeros((ZFR_count,N_cluster))
    Imported_Height_array=[]
    
    #This flag is set if the zernike eval region was selected
    Region_selection_flag=False
    
    if pathname_load!=[]:
        
        #Load the *.json Datafile
        importfile = open(pathname_load) 
        data_dict = json.load(importfile)
        importfile.close()
        print("Loaded topography data from json file")
        
        Imported_Height_Matrix=np.array(data_dict['Topography_diff'])*10**6 #in nm
        for c_ind in range(N_cluster):
            Imported_Height=Imported_Height_Matrix[:,:,c_ind] # in nm
            print('Determine Zernike Coefficients for defocus' + str(offsetvector[c_ind]) + ' nm')
            clustername='Defocus='+ str(offsetvector[c_ind]) + ' nm '

            #The sample size on the camera does not change relevantly during the samll Defocus scan
            #Thus it is determined once for all clusters and set as the zernike evaluation region
            #This step is unnecessary if working with simulated data
            if(Region_selection_flag==False):
                print('Determining the region of valid data for Zernike fitting.')
            
                #scan the field in x and y direction and detect the circle of valid values
                #The matrix contains the notnan indices for axis 0 and 1 in a column each
                notnan_index_matrix=np.argwhere(np.isnan(Imported_Height)==False)
                cutmin_ax0=np.min(notnan_index_matrix[:,0])-1
                cutmax_ax0=np.max(notnan_index_matrix[:,0])+2
                cutmin_ax1=np.min(notnan_index_matrix[:,1])-1
                cutmax_ax1=cutmin_ax1+(cutmax_ax0-cutmin_ax0) #To get a square cut off
                
                #Cut the imported data to only have the non-NaN circle/rectangle left
                Cut_Height=Imported_Height[cutmin_ax0:cutmax_ax0,cutmin_ax1:cutmax_ax1]
                print('Shape of not NaN rectangle: '+ str(np.shape(Cut_Height)))
            
                #Prepare the zernike polynomials of different grade on a unit square
                Naperture=Imported_Height.shape[0]
                Naperture_cut=Cut_Height.shape[0]
                x_aperture=np.linspace(-1,1,Naperture)
                y_aperture=np.linspace(-1,1,Naperture)
                x_aperture_cut=np.linspace(-1,1,Naperture_cut)
                y_aperture_cut=np.linspace(-1,1,Naperture_cut)
                
                ################################################
                #Determine the Zernike raster for the cut aperture rectangle
                
                Z_nmi_cut_ZFR=np.zeros((Naperture_cut,Naperture_cut,np.size(m_zi)))
                for j in range(np.size(m_zi)):
                    Z_nmi_cut_ZFR[:,:,j]=zernike_poly_xy_cartesian(n_zi_ZFR[j],m_zi_ZFR[j],x_aperture_cut,y_aperture_cut)
                    
                X_zernike_input_cut_ZFR=np.zeros((Naperture_cut*Naperture_cut,np.size(m_zi)));
                for j in range(np.size(m_zi)):
                    temp_cut_ZFR=Z_nmi_cut_ZFR[:,:,j]
                    X_zernike_input_cut_ZFR[:,j]=temp_cut_ZFR.flatten()
                ################################################
                
                Region_selection_flag=True
            else:
                Cut_Height=Imported_Height[cutmin_ax0:cutmax_ax0,cutmin_ax1:cutmax_ax1]
                
                
            #Add the topo data to an array
            Imported_Height_array.append(Cut_Height)
            
            y_meas_input=Imported_Height.flatten()
            #Remove all nan elements from y_meas_input and those associated to it in X_zernike_input
            nan_index=np.argwhere(np.isnan(y_meas_input))
            notnan_index=np.argwhere(np.isnan(y_meas_input)==False)
            
            y_meas_input_cut=Cut_Height.flatten()
            #Remove all nan elements from y_meas_input and those associated to it in X_zernike_input
            nan_index_cut=np.argwhere(np.isnan(y_meas_input_cut))
            notnan_index_cut=np.argwhere(np.isnan(y_meas_input_cut)==False)
            
            #Fitting problem for coefficient fitting
            Id_i=np.eye(np.size(m_zi))
            Id_i[0,0]=0
            lambda_i=0
            
            X_zernike_input_nonan_cut_ZFR=X_zernike_input_cut_ZFR[notnan_index_cut[:,0],:]
            y_meas_input_nonan_cut_ZFR=y_meas_input_cut[notnan_index_cut[:,0]]
            
            Popt_input_cut_ZFR=np.linalg.inv((lambda_i*Id_i+np.transpose(X_zernike_input_nonan_cut_ZFR)@X_zernike_input_nonan_cut_ZFR)) @ np.transpose(X_zernike_input_nonan_cut_ZFR) @ y_meas_input_nonan_cut_ZFR

            print(clustername + ' ZFR Coefficients: ' + str(Popt_input_cut_ZFR))
            for ind_p in range(np.size(Popt_input_cut_ZFR)):
                print('ZFR'+str(ind_p)+': ' +'['+str(n_zi_ZFR[ind_p])+' ' +str(m_zi_ZFR[ind_p])+ ']' + str(Popt_input_cut_ZFR[ind_p]))
                zernike_clustermatrix[ind_p,c_ind]=Popt_input_cut_ZFR[ind_p].copy()

            
            Fitted_surface_cut_ZFR=X_zernike_input_cut_ZFR @ Popt_input_cut_ZFR
            Fitted_surface_cut_ZFR[nan_index_cut[:,0]] = float("nan")
            Fitted_surface_cut_ZFR=Fitted_surface_cut_ZFR.reshape((Naperture_cut,Naperture_cut)) #in mm
            
            Fitted_surface_ZFR=np.zeros((Naperture,Naperture))
            Fitted_surface_ZFR[cutmin_ax0:cutmax_ax0,cutmin_ax1:cutmax_ax1]=Fitted_surface_cut_ZFR  #in mm
            
            #fig, ax = plt.subplots()
            #im = ax.imshow((Cut_Height-Fitted_surface_cut_ZFR), interpolation='bilinear', origin='lower', cmap=cm.jet, extent=[-Naperture/2, Naperture/2, -Naperture/2, Naperture/2])
            #cb = plt.colorbar(im)
            #cb.set_label('Topography/nm')
            #ax.set_title(clustername + ' residual to ZFR fit')
            
            
        print('Determined Zernike coefficients as matrix:')

        if (ZFR_count>=11):
            #Determine the spherical aperture error at defocus ZFR3=0
            p_ZFR8=np.polyfit(zernike_clustermatrix[3,:],zernike_clustermatrix[8,:],1)
            polynom=np.poly1d(p_ZFR8)
            fit_points=polynom(np.sort(np.append(zernike_clustermatrix[3,:],0)))
            ZFR8_cor_val=polynom(0)
            print('The value ZFR8(0)='+str(ZFR8_cor_val)+' nm')
            
            #Determine the higher spherical aperture error at defocus ZFR3=0
            p_ZFR15=np.polyfit(zernike_clustermatrix[3,:],zernike_clustermatrix[15,:],1)
            polynom=np.poly1d(p_ZFR15)
            fit_points=polynom(np.sort(np.append(zernike_clustermatrix[3,:],0)))
            ZFR15_cor_val=polynom(0)
            print('The value ZFR15(0)='+str(ZFR15_cor_val)+' nm')
            
            plot_index=np.squeeze(np.argwhere(np.abs(P_custom)>0))
            
            fig_zern, ax_zern = plt.subplots()
            for ind_p in range(np.size(plot_index)):
                ax_zern.plot(zernike_clustermatrix[3,:],zernike_clustermatrix[plot_index[ind_p],:],label="ZFR"+str(plot_index[ind_p]),marker='x')
            ax_zern.set_xlabel('ZFR3/nm')
            ax_zern.set_ylabel(r'ZFR$i$/nm')
            plt.legend(loc='upper left')
            
            fig_zern2, ax_zern2 = plt.subplots()
            for ind_p in range(np.size(plot_index)):
                ax_zern2.plot(offsetvector,zernike_clustermatrix[plot_index[ind_p],:],label="ZFR"+str(plot_index[ind_p]),marker='x')
            ax_zern2.set_xlabel(r'$(z-z_{conf})$/nm')
            ax_zern2.set_ylabel(r'ZFR$i$/nm')
            plt.legend(loc='upper right')
            
            #Determine the average of all measurements from all clusters
            #This would be done with the data aquired during the focus scan to reduce image noise
            Height_avg_all=np.mean(Imported_Height_array,0)
            
            
            #Cut the imported data to only have the non-NaN circle/rectangle left
            Cut_Height_avg_all=Height_avg_all[cutmin_ax0:cutmax_ax0,cutmin_ax1:cutmax_ax1]

            
            #Calculate the Zernike polynomial ZFR8(0) on the raster Naperture x Naperture
            Pzfr8=np.zeros(np.shape(Popt_input_cut_ZFR))
            Pzfr8[8]=ZFR8_cor_val
            Fitted_surface_ZFR8_cor_cut=X_zernike_input_cut_ZFR @ Pzfr8
            Fitted_surface_ZFR8_cor_cut[nan_index_cut[:,0]] = float("nan")
            Fitted_surface_ZFR8_cor_cut=Fitted_surface_ZFR8_cor_cut.reshape((Naperture_cut,Naperture_cut))
            
            Fitted_surface_ZFR8_cor=np.ones((Naperture,Naperture))*float("nan")
            Fitted_surface_ZFR8_cor[cutmin_ax0:cutmax_ax0,cutmin_ax1:cutmax_ax1]=Fitted_surface_ZFR8_cor_cut
            
            #Calculate the Zernike polynomial ZFR15(0) on the raster Naperture x Naperture
            Pzfr15=np.zeros(np.shape(Popt_input_cut_ZFR))
            Pzfr15[15]=ZFR15_cor_val
            Fitted_surface_ZFR15_cor_cut=X_zernike_input_cut_ZFR @ Pzfr15
            Fitted_surface_ZFR15_cor_cut[nan_index_cut[:,0]] = float("nan")
            Fitted_surface_ZFR15_cor_cut=Fitted_surface_ZFR15_cor_cut.reshape((Naperture_cut,Naperture_cut))
            
            Fitted_surface_ZFR15_cor=np.ones((Naperture,Naperture))*float("nan")
            Fitted_surface_ZFR15_cor[cutmin_ax0:cutmax_ax0,cutmin_ax1:cutmax_ax1]=Fitted_surface_ZFR15_cor_cut
            
            #Determine the Surface with correction of Piston, Tip, Tilt and Defocus and the added ZFR8(0), ZFR15(0)
            y_meas_input_cut=Cut_Height_avg_all.flatten()
            #Remove all nan elements from y_meas_input and those associated to it in X_zernike_input
            nan_index_cut=np.argwhere(np.isnan(y_meas_input_cut))
            notnan_index_cut=np.argwhere(np.isnan(y_meas_input_cut)==False)
            
            X_zernike_input_nonan_cut_ZFR=X_zernike_input_cut_ZFR[notnan_index_cut[:,0],:]
            y_meas_input_nonan_cut_ZFR=y_meas_input_cut[notnan_index_cut[:,0]]
            Popt_input_all_ZFR_cut=np.linalg.inv((lambda_i*Id_i+np.transpose(X_zernike_input_nonan_cut_ZFR)@X_zernike_input_nonan_cut_ZFR)) @ np.transpose(X_zernike_input_nonan_cut_ZFR) @ y_meas_input_nonan_cut_ZFR
            Popt_input_pttd_ZFR_cut=np.zeros(np.shape(Popt_input_all_ZFR_cut))
            
            for ind_poly in [0,1,2,3,8,15]:
                Popt_input_pttd_ZFR_cut[ind_poly]=Popt_input_all_ZFR_cut[ind_poly]
                
            
            Fitted_surface_ZFR_pttd_cut=X_zernike_input_cut_ZFR @ Popt_input_pttd_ZFR_cut
            Fitted_surface_ZFR_pttd_cut[nan_index_cut[:,0]] = float("nan")
            Fitted_surface_ZFR_pttd_cut=Fitted_surface_ZFR_pttd_cut.reshape((Naperture_cut,Naperture_cut))
            
            Fitted_surface_ZFR_pttd=np.ones((Naperture,Naperture))*float("nan")
            Fitted_surface_ZFR_pttd[cutmin_ax0:cutmax_ax0,cutmin_ax1:cutmax_ax1]=Fitted_surface_ZFR_pttd_cut
            
            Surface_avg_all_DF0=Height_avg_all-Fitted_surface_ZFR_pttd+Fitted_surface_ZFR8_cor+Fitted_surface_ZFR15_cor
            
            
            fig_final, ax_final = plt.subplots()
            im_final = ax_final.imshow((Surface_avg_all_DF0), interpolation='bilinear', origin='lower', cmap=cm.jet, extent=[-freeaperture/2, freeaperture/2, -freeaperture/2, freeaperture/2])
            cb = plt.colorbar(im_final)
            cb.set_label('Topography/nm')
            ax_final.set_title('Defocus=0 (ZFR3=0) Topography/nm')
            ax_final.set_xlabel('x/mm')
            ax_final.set_ylabel('y/mm')
            
            
            #Save Data if not already present
            data_dict={ "Topography_processed": Surface_avg_all_DF0,
                       "x-axis":x_aperture,
                       "y_axis":y_aperture}


            filepath_data2=os.getcwd()
            directory2 = os.listdir(filepath_data2)
            searchstring2 = "data_dict_processed_sphere"
            data_dict_flag=False
            
            if(searchstring2!=""):
                    for fname in directory2:
                        if searchstring2 in fname:
                            data_dict_flag=True
                            print('Datafile already exists!')
            
            if(data_dict_flag==False):
                with open('data_dict_processed_sphere.json', 'w') as fp:
                    json.dump(data_dict, fp, cls=NumpyEncoder, indent="")
                    
            return Surface_avg_all_DF0
            
            
        else:
            print('You need at least Zernike of order 5 to process the data appropriately!')
            

    
 



##### Main Program ######################################################

if __name__ == "__main__":
    try:
        get_ipython().magic('clear')
        get_ipython().magic('reset -f')
        print('cleared variables and console')
    except:
        pass
    
    print('Executing Sphere_Simulator.py')
    
    ### Start of custom Parameters ###
    curvatureradius=100 #mm
    freeaperture=120 #mm
    Naperture=200 #Number of samples in x/y direction
    offsetvector=np.array([-500, -250, 0.0, 250, 500])*10**(-6) #Defocus in mm
    #The Fringe Zernike set with 36 coefficients is employed
    n_zi=[0,1,1,2,2,2,3,3,4,3,3,4,4,5,5,6,4,4,5,5,6,6,7,7,8,5,5,6,6,7,7,8,8,9,9,10]
    m_zi=[0,1,-1,0,2,-2,1,-1,0,3,-3,2,-2,1,-1,0,4,-4,3,-3,2,-2,1,-1,0,5,-5,4,-4,3,-3,2,-2,1,-1,0]
    
    #Zernike coefficients in mm
    P_custom=np.array([0,0,0,0,1.0,0,3.0,4.0,5.0,0,0,0,0,0,0,-5.0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0])*10**(-6) 
    
    plt.rcParams.update({'font.size': 12})
    ### End of custom Parameters ###
    
    
    Height_data_matrix_diff_input=generate_data(curvatureradius,freeaperture,Naperture,offsetvector,n_zi,m_zi,P_custom)
    
    Surface_avg_all_DF0=determine_zernike(n_zi,m_zi,offsetvector,P_custom)
    
    fig_control, ax_control = plt.subplots()
    im_control = ax_control.imshow((Surface_avg_all_DF0-Height_data_matrix_diff_input[:,:,2]), interpolation='bilinear', origin='lower', cmap=cm.jet, extent=[-freeaperture/2, freeaperture/2, -freeaperture/2, freeaperture/2])
    cb = plt.colorbar(im_control)
    cb.set_label('Topography/nm')
    ax_control.set_title('Input/Output Difference/nm')
    ax_control.set_xlabel('x/mm')
    ax_control.set_ylabel('y/mm')
    
    
