%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 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.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function MC_draw_figure(fname)
  % function MC_draw_figure(fname)
  % Makes the figure visualizing the results of the MC experiment stored in <fname>. Replicates Figure 1 in the paper.

  %%
  bPlotPEs = 1;
  bIncludePECIs = 1;
  bIncludeTruePEsBounds = 1;
  bIncludeCIofNaive = 1;
    
  bInternal = 0; % set to 0 to make figures for the paper
  
  iFigBoundsNew = 11;
  fig_BoundsNew = figure(iFigBoundsNew); clf;
  
  tic_ID = tic;
  fprintf('** file %s\n', fname);
  D = load(fname);
  R = D.arr_all_R{1}(1);
  rho_UV_arr = D.rho_UV_arr;
  bProbit = D.bProbit;
  
  fn_extract_field = @(s_field, idx1, idx2) fn_extract_arr_R_field(s_field, idx1, idx2, D.arr_all_R);
  
  dim_theta = length(R.theta); 
  nTheta = dim_theta;
  n_rho = length(rho_UV_arr);

%% ---------------------- PEM plots: for paper ------------------------------------
  if bPlotPEs
    %my_light_blue = [0.3    0.75   1];% dark blue = b = [0 0 1]
    my_light_blue = [0.2    0.6   1];% we want it closer to dark blue, to make more noticeable
    my_yellow     = [0.9290 0.6940 0.1250];
    my_green      = [0.4660 0.6740 0.1880];
    my_purple     = [0.75, 0, 0.75];
    arr_fn2plot = {@median}; %{@mean, @median}; 
    n_fn2plot = length(arr_fn2plot);
    %fn_extract_field = @(s_field, idx1, idx2) fn_extract_arr_R_field(s_field, idx1, idx2, arr_all_R);
      
    for i_fn = 1:n_fn2plot
      fn2plot = arr_fn2plot{i_fn};
      set(0, 'CurrentFigure', fig_BoundsNew);
      clf;
      i_theta = 1; % PE for X1 only
        nAPE_targets = qyn(bProbit, 1, 2); % Probit: {PEM-P}, Tobit: {PEM-E, PEM-P}
        for i_plot = 1:nAPE_targets
          bPEM_E = i_plot==1 && nAPE_targets==2;
          bPEM_Pr = ~bPEM_E;
          i_row = i_theta+nTheta*bPEM_Pr; %*(i_plot-1);
          s_PEM_label = qyn(bPEM_Pr, "Pr", "E");% superscript for PE. PE^{Tob} = PEM-E, PE^{Pr} = PEM-P
          PEM_true   = fn2plot(fn_extract_field(qyn(bPEM_E,'PEM_E_true','PEM_P_true'), i_theta, 1)); % use correct PEM_X_oracle field
          arr_PEM_UB = fn_extract_field('PEMs_Corr', i_row, 2);
          arr_PEM_LB = fn_extract_field('PEMs_Corr', i_row, 1);
          PEM_UB = fn2plot(arr_PEM_UB);
          PEM_LB = fn2plot(arr_PEM_LB);
          PEM_CI_lb = fn2plot(fn_extract_field('CI_PE_Corr'      , i_row, 1));
          PEM_CI_ub = fn2plot(fn_extract_field('CI_PE_Corr'      , i_row, 2));
          
          PEM_Naive   = fn2plot(fn_extract_field('PEM_Naive', i_row, 1));
          
          subplot(1,2,i_plot);
          hold on
          h(1) =plot(rho_UV_arr(:),PEM_UB(:)     ,'--b','LineWidth',4);
          h(2) =plot(rho_UV_arr(:),PEM_LB(:)     ,'--', 'LineWidth',4,'Color',my_light_blue);
          h(3) =plot(rho_UV_arr(:),PEM_true(:)   ,'s' , 'LineWidth',1,'Color',my_green, 'MarkerSize',10);
          h(4) =plot(rho_UV_arr(:),PEM_Naive(:)  ,'o',  'LineWidth',2,'Color',my_yellow,'MarkerSize',8);

          if bIncludeTruePEsBounds
            true_PEM_lb = fn2plot(fn_extract_field('mc_true_PEM_Corr', i_row, 1));
            true_PEM_ub = fn2plot(fn_extract_field('mc_true_PEM_Corr', i_row, 2));
            h(7)=plot(rho_UV_arr(:),true_PEM_lb(:) ,':k' ,'LineWidth', 2);%
            h(8)=plot(rho_UV_arr(:),true_PEM_ub(:) ,':k' ,'LineWidth', 2);%
          end

          if bIncludePECIs
            h(5)=plot(rho_UV_arr(:),PEM_CI_lb(:) ,'-.' ,'LineWidth',2,'Color',my_purple); 
            h(6)=plot(rho_UV_arr(:),PEM_CI_ub(:) ,'-.' ,'LineWidth',2,'Color',my_purple);
          end          
          
          if bIncludeCIofNaive
            PEM_Naive_arr    = fn_extract_field('PEM_Naive'   , i_row, 1);
            se_PEM_Naive_arr = fn_extract_field('se_PEM_Naive', i_row, 1);
            CIofNaive_LB     = fn2plot(PEM_Naive_arr - norminv(1-D.alpha/2)*se_PEM_Naive_arr);
            CIofNaive_UB     = fn2plot(PEM_Naive_arr + norminv(1-D.alpha/2)*se_PEM_Naive_arr);
            h(9) = plot(rho_UV_arr(:),CIofNaive_LB(:),'--' ,'LineWidth', 2,'MarkerSize',4,'Color',my_yellow);
            h(10)= plot(rho_UV_arr(:),CIofNaive_UB(:),'--' ,'LineWidth', 2,'MarkerSize',4,'Color',my_yellow);
          end
          
          s_PEM_legend_prefix = ""; %"avg PEM ";
          if i_plot==nAPE_targets
            legend(s_PEM_legend_prefix + ["UB" "LB" "true PE" "naive" ...% "oracle" ...
                                qyn(bIncludeTruePEsBounds, ["true bounds" "true bounds"], []) ...
                                qyn(bIncludePECIs, ["CI" "CI"], []) ...
                                qyn(bIncludeCIofNaive, ["CI naive" "CI naive"], [])] ...
                                ,'Location', 'best'...
                                ,'Orientation','horizontal');
            lgd=legend(h([1:5 7 9]));
            set(lgd,'Position',[0.15 qyn(bInternal, 0.00, 0.95) 0.7099 0.03]);% [left bottom width height]
          end

          s_x_label = "(" + char('a'+i_plot-1) + ")" + qyn(~bInternal, "", " PE-" + s_PEM_label);      
          xlabel(s_x_label, 'FontSize', 11);
        end
    end %for i_fn
  end %if bPlotPEs
  fprintf('Plotting took %1.1fs\n', toc(tic_ID));
end % of main function

%% Helper Functions

function arr_res = fn_extract_arr_R_field(s_field, idx1, idx2, arr_all_R)
  n_rho = length(arr_all_R);
  n_reps = length(arr_all_R{1});
  arr_res = zeros(n_reps, n_rho);
  for i_rho=1:n_rho
    for i_rep=1:n_reps
      arr_res(i_rep,i_rho) = arr_all_R{i_rho}(i_rep).(s_field)(idx1,idx2);
    end
  end
end

