% Run time flags
TopHat = false;
ZeroPad  = 1; % go to next 0,1 or 2 power of two?
DoFFTWk= false;
Log_bins = true;
Linear_bins = false;

% Choose the number of bins fr the 6D covariance
nkbins = 64;%64;%16;
nkthetabins = 48;%48;%8;
nkphibins   = 32;%32;


%-------------------------------------------------------------------
%1- Get the 2dfGRS galaxy density, construct selection function W following FKP,
%   place in zero-pad and Fourier transform 

% for 2dfGRS:
x_box  = 1007.0 ;%Mpc/h
y_box  =  907.4 ;%Mpc/h
z_box  =  411.2 ;%Mpc/h

box = [x_box y_box z_box];

% Load the survey selection function
load gal_density_256_256_128.mat

%Get the size of the loaded file
nbins = size(gal_density);
Nc = numel(gal_density);

% Get the total number of bins, including zero-padding
nbins_tot(1) = 2.^(nextpow2(nbins(1))+ZeroPad);
nbins_tot(2) = 2.^(nextpow2(nbins(2))+ZeroPad);
nbins_tot(3) = 2.^(nextpow2(nbins(3))+ZeroPad);

Nc_tot  = nbins_tot(1)*nbins_tot(2)*nbins_tot(3);
%sum(sum(sum(gal_density)))/Nc_tot ~ 4.3e-4 for 2dfGRS

% Get the origin shift in the zero padded map
shift(1:3)  = (nbins_tot(:)  - nbins(:))/2;

% Get the center in the zero padded map
center(1:3)  = nbins_tot(:)/2+1;

% Get the volume ratios of the zero padded map
ratio(1:3)  = nbins_tot(:)./nbins(:);

% Get the weigth function (defined in FKP)
meanP = 5000;
gal_weight = 1./(1. + gal_density*meanP);

% Fit the selection function into to zero padded map
W = zeros(nbins_tot(1), nbins_tot(2),nbins_tot(3));
if(TopHat)
   W(shift(1)+1:shift(1)+nbins(1), shift(2)+1:shift(2)+nbins(2), shift(3)+1:shift(3)+nbins(3)) = 1.0;
else
   W(shift(1)+1:shift(1)+nbins(1), shift(2)+1:shift(2)+nbins(2), shift(3)+1:shift(3)+nbins(3)) = gal_density.*gal_weight;
end

%Get the Norm (defined in FKP)
Norm = sqrt(sum(sum(sum(W.^2))))*sqrt(Nc)/3;
if(TopHat)
    Norm = Norm*3;
end

% Fourier transform the full map, normalize and save
% this corresponds to eq. 2.1.7 in FKP
W_name = ['./Wk_'  num2str(nbins(1)) '_' num2str(nbins(2)) '_' num2str(nbins(3)) '_zeroPad_' num2str(ZeroPad) '.mat' ];
if(DoFFTWk)
   disp('performing fft of W(x)')
   Wk = fftshift(fftn(W))/Norm;
   disp('Done!')
   save(W_name, 'Wk')
   disp('Saved Wk')
else
   load(W_name) 
   disp('Loaded Wk')
end
%clear W



%-----------------------------------------------
% 2- Get the bins for the survey

%-------
%radial

k_min = 2*pi/max([x_box*ratio(1) y_box*ratio(2) z_box*ratio(3)]); % h/Mpc
V = 2*pi/k_min; % Mpc/h
k_Ny =  max(2*pi./(box(:).*ratio(:)).*nbins(:));
k_max =2.0;%k_Ny;%/4;

%define survey bins and centers
%k_survey = logspace(log10(k_min),log10(k_max),nkbins);
k_survey = zeros(1,nkbins);
dkr = zeros(1,nkbins);
%k_edges(1)= k_survey(1) - (k_survey(2)*k_survey(1))^(0.5);
%k_edgess(nkbins+1)  = k_survey(nkbins) + (k_survey(nkbins)*k_survey(nkbins-1))^(0.5);


if(Log_bins)
   k_edges = logspace(log10(k_min),log10(k_max),nkbins + 1);
   for bin = 1:nkbins
     k_survey(bin) = (k_edges(bin)*k_edges(bin+1))^(0.5);
     dkr(bin) = k_edges(bin+1) -k_edges(bin);
   end
elseif(Linear_bins)
   k_edges = linspace(k_min,k_max,nkbins + 1);    
   for bin = 1:nkbins
     k_survey(bin) = (k_edges(bin)+k_edges(bin+1))*(0.5);
     dkr(bin) = k_edges(bin+1) -k_edges(bin);
   end
else
   error('Must choose a binning type (linear or log)') 
end
    
%-------
%angular

kthetamin = 0.0;
kthetamax = 2*pi;
ktheta = zeros(1,nkthetabins);
ktheta_edges = linspace(kthetamin,kthetamax,nkthetabins + 1);    

for bin = 1:nkthetabins
  ktheta(bin) = (ktheta_edges(bin)+ktheta_edges(bin+1))*(0.5);
end
dktheta = ktheta_edges(2) - ktheta_edges(1);

kphimin = 0;
kphimax = pi;
kphi = zeros(1,nkphibins);
kphi_edges = linspace(kphimin,kphimax,nkphibins + 1);    

for bin = 1:nkphibins
  kphi(bin) = (kphi_edges(bin)+kphi_edges(bin+1))*(0.5);
end
dkphi = kphi_edges(2) -kphi_edges(1);

disp('********************************************')
disp(['Convolution performed  using ' num2str(nkbins)  ' radial bins,'])
disp([num2str(nkthetabins) ' theta bins and ' num2str(nkphibins) ' phi bins.'])
disp('********************************************')

% Get min and max k to be read in the selection function table:
ktable_min(1:3) = -(center(:)-1)*k_min;
ktable_max(1:3) = (center(:)-1)*k_min;

disp('min and max modes of this survey, including zero-pading modes:')
disp([ktable_min ktable_max])


%-----------------------------------------------
% 3- Load P(k), interpolate onto k_survey, 
% then get fitted C_ell  matrice, 

% ps_file  =  load('Output_Z0.5-NGP.txt');

ps_file  = load('CAMB_Z0.5_high_k_WMAP5.dat');
p_survey = interp1(ps_file(:,1), ps_file(:,2), k_survey);


% Get Gaussian Covariance and C_ell
N_over_2 = (1./2.)*(4.0*pi*k_survey.^2.*dkr).*(V/2/pi)^3./(ratio(1)*ratio(2)*ratio(3));
Cov_gauss = p_survey.^2./N_over_2;
C_ell_gauss = 4*pi*Cov_gauss;

[C0 C2 C4 lambda U Ratio] = FittedCov(k_survey',p_survey,N_over_2);

PID =  fopen(['./Pk_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, p_survey, 'float32')
fclose(PID);


PID =  fopen(['./Cov_gauss_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, Cov_gauss, 'float32')
fclose(PID);

PID =  fopen(['./lambda_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, lambda, 'float32')
fclose(PID);

PID =  fopen(['./U_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, U, 'float32')
fclose(PID);

PID =  fopen(['./Ratio_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, Ratio, 'float32')
fclose(PID);

%-----------------
%--Poisson Error--
U_Poisson_400(1:nkbins)  = FitU1_new2([1.0193 0.0947 2.1021 2.5861], k_survey);
U_Poisson_1216(1:nkbins) = FitU1_new2([0.9987 0.2034 2.1553 2.3407], k_survey);
U_Poisson_8000(1:nkbins) = FitU1_new2([0.2966 3.3736 0.6099 0.6255], k_survey);

lambda_400 = 52.02;
lambda_1216 = 45.09;
lambda_8000 = 24.41;

rho_Poisson_400 = zeros(nkbins);
rho_Poisson_1216 = zeros(nkbins);
rho_Poisson_8000 = zeros(nkbins);
cov_Poisson_400 = zeros(nkbins);
cov_Poisson_1216 = zeros(nkbins);
cov_Poisson_8000 = zeros(nkbins);


Ratio_400(1:nkbins)  = exp(2.6936)*k_survey.^(2.1347);
Ratio_1216(1:nkbins) = exp(1.6533)*k_survey.^(2.1965);
Ratio_8000(1:nkbins) = exp(-0.4321)*k_survey.^(2.0347);

Gauss_400(1:nkbins) = (10^(4.31).^2)./N_over_2;
Gauss_1216(1:nkbins) = (10^(3.82).^2)./N_over_2;
Gauss_8000(1:nkbins) = (10^(3.0).^2)./N_over_2;

diag_Poisson_400(1:nkbins) = Gauss_400.*Ratio_400;
diag_Poisson_1216(1:nkbins) = Gauss_1216.*Ratio_1216;
diag_Poisson_8000(1:nkbins) = Gauss_8000.*Ratio_8000;

for i = 1:nkbins
for j = 1:nkbins
   rho_Poisson_400(i,j) = lambda_400*U_Poisson_400(i)*U_Poisson_400(j);
   rho_Poisson_1216(i,j) = lambda_1216*U_Poisson_1216(i)*U_Poisson_1216(j);
   rho_Poisson_8000(i,j) = lambda_8000*U_Poisson_8000(i)*U_Poisson_8000(j);
   
   cov_Poisson_400(i,j) = rho_Poisson_400(i,j)*sqrt(diag_Poisson_400(i)*diag_Poisson_400(j));
   cov_Poisson_1216(i,j) = rho_Poisson_1216(i,j)*sqrt(diag_Poisson_1216(i)*diag_Poisson_1216(j));
   cov_Poisson_8000(i,j) = rho_Poisson_8000(i,j)*sqrt(diag_Poisson_8000(i)*diag_Poisson_8000(j));
end
end

C0_Poisson_400 = 4*pi*cov_Poisson_400;
C0_Poisson_1216 = 4*pi*cov_Poisson_1216;
C0_Poisson_8000 = 4*pi*cov_Poisson_8000;

PID =  fopen(['./lambda_400_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, lambda_400, 'float32')
fclose(PID);
PID =  fopen(['./U_400_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, U_Poisson_400, 'float32')
fclose(PID);
PID =  fopen(['./Ratio_400_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, Ratio_400, 'float32')
fclose(PID);
PID =  fopen(['./Cov_Poisson_400_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, Gauss_400, 'float32')
fclose(PID);

PID =  fopen(['./lambda_1216_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, lambda_1216, 'float32')
fclose(PID);
PID =  fopen(['./U_1216_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, U_Poisson_1216, 'float32')
fclose(PID);
PID =  fopen(['./Ratio_1216_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, Ratio_1216, 'float32')
fclose(PID);
PID =  fopen(['./Cov_Poisson_1216_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, Gauss_1216, 'float32')
fclose(PID);

PID =  fopen(['./lambda_8000_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, lambda_8000, 'float32')
fclose(PID);
PID =  fopen(['./U_8000_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, U_Poisson_8000, 'float32')
fclose(PID);
PID =  fopen(['./Ratio_8000_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, Ratio_8000, 'float32')
fclose(PID);
PID =  fopen(['./Cov_Poisson_8000_nkbins_' num2str(nkbins) '_kmin_' num2str(k_min) '_kmax_' num2str(k_max) '.dat'], 'wb');
fwrite(PID, Gauss_8000, 'float32')
fclose(PID);

  return

