%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 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 [PEMs_Corr] = fn_ivtobit_PE_from_beta(h,beta)
  
  beta = beta(:); h = h(:);
  assert(length(beta)==length(h)+3);
  % beta = [theta', sig_V, sig_UV, sig_U]'
  
  theta     = beta(1:end-3);%';
  dim_theta = length(theta);

  sig_V     = beta(end-2); %sqrt(beta(end-2));
  sig_UV    = beta(end-1);
  sig_U     = beta(end); %sqrt(beta(end));
  
  sig2_Us_LB_1 = (theta(1)*sig_UV+sig_U^2)^2/(sig_V^2*theta(1)^2 + 2*theta(1)*sig_UV + sig_U^2);
  sig_Us_LB    = max(1e-7, sqrt(max(sig2_Us_LB_1, sig_U^2-theta(1)^2*sig_V^2)));

  PEMs_Corr   = zeros(2*dim_theta,2);
  theta_h     = theta'*h;
  abs_theta_h = abs(theta'*h);
  for i = 1:2*dim_theta
    arr_PE_sig_U = [sig_Us_LB; sig_U];% 2 or 3 pts matter: para after eq (16)
    if i>dim_theta %only for PEM-Pr
      if sig_Us_LB < abs_theta_h && abs_theta_h < sig_U
        arr_PE_sig_U(end+1) = abs_theta_h; %#ok<AGROW> 
      end
    end
    arr_PE_val     = arrayfun(@(s) fn_ivtobit_PEM(s, i, theta_h, theta), arr_PE_sig_U);
    [~, i_max]     = max(arr_PE_val);
    [~, i_min]     = min(arr_PE_val);
    PEMs_Corr(i,:) = [arr_PE_val(i_min) -arr_PE_val(i_max)];
  end
  PEMs_Corr(:,2) = -PEMs_Corr(:,2); % flip sign
end

  %% Helper Functions  
  function [PEM_E, PEM_P] = fn_ivtobit_PEM(sig_U_star, only_PEM_j,theta_h,theta)
    % When only_PEM_j is between 1 and 2*dim_theta computing only a single PEM (PEM_E_j or PEM_P_j) -- used for optimization
    dim_theta = length(theta);
    if only_PEM_j==0
      PEM_E = normcdf(theta_h/sig_U_star)*theta;
      PEM_P = normpdf(theta_h/sig_U_star)*theta/sig_U_star;
    elseif only_PEM_j<=dim_theta % compute APE_E(j)
      PEM_E_j = normcdf(theta_h/sig_U_star)*theta(only_PEM_j);
      PEM_E   = PEM_E_j;
    else
      assert(only_PEM_j<=2*dim_theta);
      % compute PEM_P(j) and return as scalar in PEM_E
      PEM_P_j = normpdf(theta_h/sig_U_star)*theta(only_PEM_j-dim_theta)/sig_U_star;
      PEM_E   = PEM_P_j;
    end
  end  

