###########################################################
# Description: This file can be used to export
# data from Lumerical STACK to be used in either
# Speos or Zemax simulations
#
# Copyright 2022 Ansys, Inc.
###########################################################

############ Part 1 Parameter settings  ############
##This part would define parameters we use in the simulation.
##Including :
##  the wavelength range
##  the designed reflected wavelengths
##  the target incident angle
##  the materials
##  the output file names for both Zemax and Speos
##   
####################################################

# choose a filename
fname_Speos = "STACK_coating.coated"; #.coated
fname_Zemax = "Zemax_coating.txt"; # can be a .txt file for Zemax

# Calculate R and T using STACK
lmin = 400e-9;
lmax = 700e-9;
lambda = linspace(lmin,lmax,201);
f = c/lambda;
theta = linspace(0,90,201);

# example of an N layer reflector to target a 3 wavelengths to reflect at one angle of incidence
N = 51; # number of layer pairs for each of 3 filters
N_apodization = 5; # smooth transition from n1 to n2 over this many layers
delta_n_max = 0.05; # index difference of n2
lambda_Rpeak = [450e-9,535e-9,632e-9]; # pick 3 wavelengths to reflect
theta_Rpeak = 30; # incident angle at which we want lambda_Rpeak to be correct

# use index of glass from Palik
n1 = getindex("SiO2 (Glass) - Palik",f);
n2 = n1 + delta_n_max; # film with an index higher than glass


############ Part 2 STACK  ############
## This part would first get the thickness of each layer by using constructive interference conditions.
## After the calculation, we us STACK to get the reflection and transmission of different polarizations.(S&P)
#######################################

mean_n1 = 0.5*real(interp(n1,lambda,lambda_Rpeak) + interp(n2,lambda,lambda_Rpeak));
theta_inside = asin( sin(theta_Rpeak*pi/180)/mean_n1 );
d = [0; 
     ones(2*N+1)*lambda_Rpeak(1)/4/mean_n1(1)/cos(theta_inside(1));  # filter lambda1
     ones(2*N+1)*lambda_Rpeak(2)/4/mean_n1(2)/cos(theta_inside(2));  # filter lambda2
     ones(2*N+1)*lambda_Rpeak(3)/4/mean_n1(3)/cos(theta_inside(3));  # filter lambda3
     0]; # define layers
n = meshgridy(d,n1); # define materials, start with all n1
n(1,:) = 1; # incidence in air
# define apodization function, raised cosine filter
pos = 1:N;
apodization = (pos < N_apodization)*(1-cos(pos/N_apodization*pi))/2 + 
              ((pos >= N_apodization) & (pos<=N+1-N_apodization)) +
              (pos > N+1-N_apodization)*(1-cos((N+1-pos)/N_apodization*pi))/2;

# fill the layers with n2 materials
for(i=1:N) {
    n(1+2*i,:) = n1 + (n2-n1)*apodization(i);
}
for(i=1:N) {
    n(2*N+2+2*i,:) = n1 + (n2-n1)*apodization(i);
}
for(i=1:N) {
    n(4*N+3+2*i,:) = n1 + (n2-n1)*apodization(i);
}

plot(1:length(d),real(n(:,find(lambda,550e-9))),"layer number","refractive index","refractive index at 550nm");
setplot("show legend",false);

# calculate RT of stack
RTstack = stackrt(n,d,f,theta);

# extract data
theta = RTstack.theta;
lambda = RTstack.lambda; #
Ro = matrix(length(lambda), 2, length(theta));
To = matrix(length(lambda), 2, length(theta));
Rp = RTstack.Rp;
Rs = RTstack.Rs;
Tp = RTstack.Tp;
Ts = RTstack.Ts;
ARpo = angle(-RTstack.rp)*180/pi; # negative sign due to P convention on reflection
ARso = angle(RTstack.rs)*180/pi; # negative sign due to P convention
ATpo = angle(RTstack.tp)*180/pi; # negative sign due to P convention
ATso = angle(RTstack.ts)*180/pi; # negative sign due to P convention

# resort by lambda to ensure increasing
map = sortmap(lambda);
lambda = lambda(map);
Ro(:, 1, :) = Rp(map, :);
Ro(:, 2, :) = Rs(map, :);
To(:, 1, :) = Tp(map, :);
To(:, 2, :) = Ts(map, :);
ARpo = ARpo(map,:);
ARso = ARso(map,:);
ATpo = ATpo(map,:);
ATso = ATso(map,:);

# add the 90 degree angle with reflection 100 and transmission 0
if(!almostequal(theta(end),90)) {
    theta = [ pinch(theta); 90 ];
    R = T = matrix(length(lambda),2,length(theta));
    R(:,:,1:end-1) = Ro;
    T(:,:,1:end-1) = To;
    ARp = ARs = ATp = ATs = matrix(length(lambda),length(theta));
    ARp(:,1:end-1) = ARpo;
    ARs(:,1:end-1) = ARso;
    ATp(:,1:end-1) = ATpo;
    ATs(:,1:end-1) = ATso;
} else {
    R = Ro;
    T = To;
    ARp = ARpo;
    ARs = ARso;
    ATp = ATpo;
    ATs = ATso;
}
R(:,:,end) = 1;
T(:,:,end) = 0;
ARp(:,end) = ARp(:,end-1);
ARs(:,end) = ARs(:,end-1);
ATp(:,end) = ATp(:,end-1);
ATs(:,end) = ATs(:,end-1);

# clamping function to eliminate force any points with small
# negative reflection to 0, or any points with transmission
# great than 100 to 100.
function clamp(x) {
    result = x;
    pos = find(result < 0);
    if(pos(1) > 0) {
        result(pos) = 0;
    }
    pos = find(result > 100);
    if(pos(1) > 0) {
        result(pos) = 100;
    }
    return result;
}

# clamp R and T
R = clamp(R);
T = clamp(T);
image(lambda,theta,pinch(R,2,1),"lambda (nm)","theta (degrees)","RP");
image(lambda,theta,pinch(R,2,2),"lambda (nm)","theta (degrees)","RS");
plot(lambda,pinch(R,3,find(theta,30)),"lambda (nm)","R (%)","R at 30 degrees incidence","linewidth=3");
legend("RP","RS");

##############################################################
############# Export a Speos file ############
##############################################################

function write_Speos_coated_file(fname,theta,lambda,R,T) {
    lambda = lambda * 1e9; #switch to nm
    R = R * 100; # switch to percent
    T = T * 100; # switch to percent
    
    # write the file
    delim = '	'; # actually a tab, paste from notepad
    
    # write the header
    format short;
    write(fname,"OPTIS-Coated surface file v1.0","overwrite");
    write(fname,"Coating surface");
    write(fname,num2str(length(theta))+" "+num2str(length(lambda))); # tab or space?
    
    # write the wavelength
    outstring = "";
    spacer = delim;
    for(i=1:length(lambda)) {
        outstring = outstring + spacer + num2str(lambda(i));
        spacer = delim+delim;
    }
    write(fname,outstring);
    
    # write the angle and R, T data
    for(j=1:length(theta)) {
        
        outstring = num2str(theta(j));
        for(i=1:length(lambda)) {
            outstring = outstring + delim + num2str(R(i,1,j)) + delim + num2str(T(i,1,j));
        }
        write(fname,outstring);
        
        outstring = "";
        for(i=1:length(lambda)) {
            outstring = outstring + delim + num2str(R(i,2,j)) + delim + num2str(T(i,2,j));
        }
        write(fname,outstring);
    }
}
write_Speos_coated_file(fname_Speos,theta,lambda,R,T);


##############################################################
############# Export a Zemax file ############
##############################################################

function write_Zemax_tabular_file(fname,theta,lambda,R,T,ARp,ARs,ATp,ATs) {
    # write the file
    delim = ' '; # actually a tab, paste from notepad
    
    lambda = lambda * 1e6; # switch to micron units
    
    
    # write the header
    format short;
    write(fname,"TABLE Lumerical_STACK_model","overwrite");
    # write the angle and R, T data
    for(j=1:length(theta)) {
        
        write(fname,"ANGL" + delim + num2str(theta(j)));
        outstring = "";
        for(i=1:length(lambda)) {
            outstring = outstring + "WAVE" + delim + num2str(lambda(i)) + delim +
            num2str(R(i,2,j)) + delim + num2str(R(i,1,j)) + delim +
            num2str(T(i,2,j)) + delim + num2str(T(i,1,j)) + delim +
            num2str(ARs(i,j)) + delim + num2str(ARp(i,j)) + delim + 
            num2str(ATs(i,j)) + delim + num2str(ATp(i,j));
            if(i < length(lambda)-0.5) {
                outstring = outstring + endl;
            }
        }
        write(fname,outstring);
        
    }
}
write_Zemax_tabular_file(fname_Zemax,theta,lambda,R,T,ARp,ARs,ATp,ATs);