%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Part of the replication package for the paper
%   "Marginal Effects for Probit and Tobit with Endogeneity"
%   by Kirill S. Evdokimov, Ilze Kalnina, and Andrei Zeleneev.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% This program generates stored_grid_CV_for_max_1-sided.mat
% These are quantiles of max(Z_1,Z_2), where (Z_1,Z2) are standard Normal with correlation rho
clear

bLarge = 1; % 0 for testing

tic;
fprintf('\n');
arr_qqq = [0.995 0.9975];

n_sim_0 = qyn(bLarge, 4e7, 1e6); 
%actual 'n_sim' is determined below, after sampling
rng(0, 'mrg32k3a');

qq = norminv(min(arr_qqq))-0.005; % lower bound on the  quantile values we'll compute
n_tot = 0;
n_dropped = 0;
ZZZ_1 = [];
iter_ZZZ = 1;
n_sim_1 = min(qyn(bLarge, 1e8, 1e7), n_sim_0*10); %round(n_sim/2);
while nrows(ZZZ_1)<n_sim_0
  ZZZ = randn(n_sim_1,2);
  ix = ZZZ(:,1).^2 + ZZZ(:,2).^2 >= qq^2; %only these rows of ZZZ will be affecting arr_Q
  n_tot = n_tot + n_sim_1;
  n_dropped = n_dropped + sum(~ix);
  ZZZ(~ix,:) = [];
  ZZZ_1 = [ZZZ_1; ZZZ];
  fprintf('ZZZ iter %3d, n_have = %8.3f mln (%5.1f%%),  t = %5.2fs\n', iter_ZZZ, nrows(ZZZ_1)/1e6, nrows(ZZZ_1)/n_sim_0*100, toc);
  iter_ZZZ = iter_ZZZ+1;
end

p_dropped = n_dropped/n_tot;
arr_qqq_adj = (arr_qqq-p_dropped)/(1-p_dropped);
ZZZ = ZZZ_1;
clear ZZZ_1 ix;
toc
n_actual = nrows(ZZZ);
n_sim = n_tot;
fprintf('Effective n_tot = %7.2e, reduced to n_used = %7.2e  (ratio = %5.2f)\n', n_tot, n_actual, n_tot/n_actual);

fn_q = @(rS) quantile([max(ZZZ*rS,[],2); max(ZZZ*[rS(2,:); rS(1,:)],[],2)],arr_qqq_adj);

gr = 0.024*cumsum(0.97.^(0:130));
bbb = 0.995;
gr = 0.2 + (bbb-0.2)*(gr/gr(end));
d1 = gr(end)-gr(end-1);
gr2 = d1*cumsum(0.9.^(1:100));
gr2 = bbb + (1-1e-8-bbb)*(gr2/gr2(end));
arr_rho = [-1:0.05:0 0.025:0.025:0.2 gr gr2 1];
assert(all(diff(arr_rho)>0));

n_rho = length(arr_rho);
arr_Q = zeros(n_rho,length(arr_qqq));
for i_rho = 1:n_rho
  rho = arr_rho(i_rho);
  rS = cholcov([1 rho; rho 1]);
  rS = qyn(size(rS,1)==1, [rS; 0 0], rS); %make rS square
  arr_Q(i_rho,:) = fn_q(rS);
  if mod(i_rho, 10)==0
    fprintf('%3d/%d, t = %4.1fs\n', i_rho, n_rho, toc);
  end
end

arr_Q_0 = arr_Q;
took_total = toc;
fprintf('In total, took %4.3fs\n', toc);

%%
% make a small bias adjustment, since we know the exact quantiles for rho = {-1,0,1}
exact_Q_m1 = -norminv((1-arr_qqq)/2); % rho = -1
exact_Q_0  = norminv(arr_qqq.^0.5);
exact_Q_1  = norminv(arr_qqq);
[~,idx_0] = min(abs(arr_rho-0)); assert(abs(arr_rho(idx_0))<1e-7); % find zero in arr_idx

arr_Q = fn_monotonizing_iron(arr_Q);
d_qm1 = arr_Q(1,:)-exact_Q_m1;
d_q0 = arr_Q(idx_0,:)-exact_Q_0;
d_q1 = arr_Q(end,:)-exact_Q_1;      
bias_corr = (d_q1/2 - d_q0 + d_qm1/2).*arr_rho(:).^2 + (d_q1/2 - d_qm1/2).*arr_rho(:) + d_q0;
arr_Q = arr_Q - bias_corr;
arr_Q = fn_monotonizing_iron(arr_Q);

%%
if n_sim>=1e8 && bLarge %save results
  arr_rho = arr_rho(:);
  arr_alpha = 1-arr_qqq;
  n_sim_reduced = nrows(ZZZ);
  save('stored_grid_CV_for_max_1-sided.mat', 'arr_alpha', 'arr_rho', 'arr_Q', 'n_sim', 'n_sim_reduced');
end

%% Helper function
function arr = fn_monotonizing_iron(arr)
  [n, m] = size(arr);
  for i=2:n
    for j=1:m
      if arr(i,j)>arr(i-1,j)
        arr(i,j)=arr(i-1,j);
      end
    end
  end
end
