clear all
close all
clc

%%  SET PARAMETERS
% Set resolution of output image (without zero-padding)
Mres = 80;
% folder with telescope data
data_folder = 'data_ch_191/';
% compute beam-correlation again or load results
compute = 0;
% station diameter (size of aperture image in each dimension) 
diameter = 40;

%%  LOAD DATA
% speed of light
c0 = 299792458;
% load parameters of recorded data
load([data_folder,'params']);
% frequency
f = channel*1e6*400/512;
% wavelength
lambda = c0/f;
% source coordinates (zenith angle,azimuth) (azimuth with respect to north)
s_tp = [90-source_elevation,source_azimuth]; 
% source coordniates (l,m) in direction cosines  
s_lm = [sind(s_tp(1)).*sind(s_tp(2)),sind(s_tp(1)).*cosd(s_tp(2))]; 
% relative EDA2 antenna positions (east,north,height)
enh = read_eda2_positions('EDA2_parameters/antenna_locations_eda2.txt');

% EDA2 flagged elements
[flagged_x,flagged_y] = read_eda2_flags('EDA2_parameters/instr_config_eda2.txt');
% Files with reference calibration
x_cal_file = ['chan_',num2str(channel),'_selfcal_pha_XX.txt'];
y_cal_file = ['chan_',num2str(channel),'_selfcal_pha_YY.txt'];
% load reference calibration solution
x_cal_ref = read_calibration_phase(['EDA2_parameters/',x_cal_file]);
y_cal_ref = read_calibration_phase(['EDA2_parameters/',y_cal_file]);


%%  COMPUTATION OF BEAM-CORRELATION
if compute == 1
    % load samples from hdf5 file
    data_filename = dir([data_folder,'/*.hdf5']);
    [x_pol,y_pol] = read_hd5_samples([data_folder,data_filename.name]);
    
    % covariance matrices
    Rxx = x_pol*x_pol';
    Rxx = Rxx/max(abs(Rxx(:)));
    Ryy = y_pol*y_pol';
    Ryy = Ryy/max(abs(Ryy(:)));
    
    % antenna coefficients for reference beam
    ref_coeff = beam_steering_lm(enh,lambda,s_lm);
    
    % resolution for lm-grid
    deltaL = lambda/diameter;
    deltaM = deltaL;
    
    % compute values of lm grid
    lc = -Mres/2*deltaL:deltaL:Mres/2*deltaL;
    mc = -Mres/2*deltaL:deltaM:Mres/2*deltaM;
    [Lc,Mc] = meshgrid(lc,mc);
    L = Lc + s_lm(1);
    M = Mc + s_lm(2);
    
    % reference signals
    x_ref = ref_coeff.'*x_pol;
    y_ref = ref_coeff.'*y_pol;
    
    % beam correlations x-polarisation
    Ax = zeros(size(L));
    % beam correlations y-polarisation
    Ay = zeros(size(L));

    % compute beam correlations for all beamforming coordinates on lm grid
    for idxL = 1:length(lc)
        for idxM = 1:length(mc)
            lm = [L(idxL,idxM),M(idxL,idxM)];
            coeff = beam_steering_lm(enh,lambda,lm);
            x_beam = coeff.'*x_pol;
            y_beam = coeff.'*y_pol;
            Ax(idxL,idxM) = x_beam*x_ref';
            Ay(idxL,idxM) = y_beam*y_ref';
        end
        fprintf("Progress:\t%3d%% \n",floor(100*idxL/length(lc)));
    end
    % normalize correlation values 
    Ax = Ax/max(abs(Ax(:)));
    Ay = Ay/max(abs(Ay(:)));
    
    % store computed correlation data
    save([data_folder,'beam_correlation_M_',num2str(Mres)], ...
        'Rxx','Ryy','ref_coeff','deltaL','deltaM','L','M','Ax','Ay');

elseif compute == 0
    load([data_folder,'beam_correlation_M_',num2str(Mres)]);
end

%%  PROCESSING
% size to which beam correlation is padded
Nfft = 512;

% apply 2-dimensional gaussian window and zero-phase zero-padding
window = gausswin(size(Ax,2),2)*gausswin(size(Ax,2),2)';
Ax_padded = zero_pad_2d(Ax.*window,Nfft,Nfft);
Ay_padded = zero_pad_2d(Ay.*window,Nfft,Nfft);

% compute 2 dimensional FFT of beam correlation and normalize
Ex = fftshift(fft2(fftshift(Ax_padded)));
Ey = fftshift(fft2(fftshift(Ay_padded)));

Ex = Ex/max(abs(Ex(:)));
Ey = Ey/max(abs(Ex(:)));

% coordinates in aperture domain
east = linspace(-lambda/(2*deltaL),lambda/(2*deltaL),Nfft);
north = linspace(-lambda/(2*deltaM),lambda/(2*deltaM),Nfft);
[EAST,NORTH] = meshgrid(east,north);

% pick out antenna excitations
x_val = get_antenna_values(EAST,NORTH,Ex,enh);
y_val = get_antenna_values(EAST,NORTH,Ey,enh);
x_power = 10*log10(abs(x_val)); 
y_power = 10*log10(abs(y_val));

% compute RMS to reference calibration
xflags = zeros(256,1); 
yflags = zeros(256,1);
xflags(x_power < -10) = 1; 
yflags(y_power < -10) = 1;
xflags(flagged_x) = 1; 
yflags(flagged_y) = 1;
x_phase_diff_all = 180/pi*angle(x_val.*x_cal_ref); 
y_phase_diff_all = 180/pi*angle(y_val.*y_cal_ref);
x_phase_diff = x_phase_diff_all(~xflags)-mean(x_phase_diff_all(~xflags));
y_phase_diff = y_phase_diff_all(~yflags)-mean(y_phase_diff_all(~yflags));
x_phase_diff_all = x_phase_diff_all - mean(x_phase_diff_all(~xflags));
y_phase_diff_all = y_phase_diff_all - mean(y_phase_diff_all(~yflags));
x_phase_rms = sqrt(mean(x_phase_diff.^2));
y_phase_rms = sqrt(mean(y_phase_diff.^2));
total_phase_rms = sqrt(mean([x_phase_diff.^2;y_phase_diff.^2]));

%%  PLOT RESULTS
beamCorr = plot_beam_corr(L-s_lm(1),M-s_lm(2),window.*Ax);
apertureDistr = plot_aperture_distr(EAST,NORTH,Ex,enh,flagged_x);
aperturePhase = plot_aperture_phase(EAST,NORTH,Ex,enh,flagged_y);
calibrationOffset = plot_calibration_offset(enh,x_phase_diff_all,xflags,0);
fprintf("RMS between phase calibrations: %f\n", total_phase_rms);